import pytest
from gotranx import exceptions
from gotranx import ode

def test_ODE_with_incomplete_component_raises_ComponentNotCompleteError(parser, trans):
    expr = "states(x=1, y=2)\n dy_dt=0"
    tree = parser.parse(expr)
    result = trans.transform(tree)
    with pytest.raises(exceptions.ComponentNotCompleteError) as e:
        ode.ODE(result.components)
    assert str(e.value) == "Component None is not complete. Missing state derivatives for ['x']"

@pytest.mark.parametrize(
    "input, output",
    [
        ((), set()),
        (("x", "y", "z"), set()),
        (("x", "y", "y"), {"y"}),
        (("x"), set()),
        (("y", "y", "y"), {"y"}),
        (("x", "y", "y", "x"), {"x", "y"}),
    ],
)
def test_find_duplicates(input, output):
    assert ode.find_duplicates(input) == output

def test_ODE_with_duplicates_raises_DuplicateSymbolError(parser, trans):
    expr = "states(y=2)\n dy_dt=0 \n y=42"
    tree = parser.parse(expr)
    with pytest.raises(exceptions.DuplicateSymbolError) as e:
        ode.ODE(trans.transform(tree).components)

    assert "Found multiple definitions for {'y'}" in str(e.value)

def test_make_ode_with_duplicates_raises_DuplicateSymbolError(parser, trans):
    expr = "states(y=2)\n dy_dt=0 \n y=42"
    tree = parser.parse(expr)
    with pytest.raises(exceptions.DuplicateSymbolError) as e:
        ode.make_ode(trans.transform(tree).components)

    assert "Found multiple definitions for {'y'}" in str(e.value)

def test_make_ode(parser, trans):
    expr = """
    # This is an ODE.
    # Here is another line.
    states("First component", "X-gate", x = 1, xr=3.14)
    states("First component", "Y-gate", y = 1)
    states("Second component", z=1)
    parameters("First component", a=1, b=2)
    parameters("Second component", c=3)

    expressions("First component")
    dx_dt=a+1
    d = a + b * 2 - 3 / c
    dy_dt = 2 * d - 1
    dxr_dt = (x / b) - d * xr

    expressions("Second component")
    dz_dt = 1 + x - y
    """
    tree = parser.parse(expr)

    result = ode.make_ode(*trans.transform(tree), name="TestODE")

    assert result.text == "This is an ODE. Here is another line."
    assert repr(result) == "ODE(TestODE, num_states=4, num_parameters=3)"

    assert result.num_components == 2

    states = result.states
    assert len(states) == result.num_states == 4
    state_names = set()
    state_symbols = set()
    for state in states:
        # assert isinstance(state, atoms.TimeDependentState)
        state_names.add(state.name)
        state_symbols.add(str(state.symbol))

    assert state_names == {"x", "xr", "y", "z"}
    # assert state_symbols == {"x(t)", "xr(t)", "y(t)", "z(t)"}

    parameters = result.parameters
    assert len(parameters) == result.num_parameters == 3
    parameter_names = set()
    for parameter in parameters:
        parameter_names.add(parameter.name)

    assert parameter_names == {"a", "b", "c"}

    state_derivatives = result.state_derivatives
    assert len(state_derivatives) == 4
    state_derivative_names = set()
    state_derivative_symbols = set()
    for state_derivative in state_derivatives:
        # assert isinstance(state_derivative.state, atoms.TimeDependentState)
        state_derivative_names.add(state_derivative.name)
        state_derivative_symbols.add(str(state_derivative.symbol))

    assert state_derivative_names == {"dx_dt", "dxr_dt", "dy_dt", "dz_dt"}
    # assert state_derivative_symbols == {
    #     "Derivative(x(t), t)",
    #     "Derivative(xr(t), t)",
    #     "Derivative(y(t), t)",
    #     "Derivative(z(t), t)",
    # }

    intermediates = result.intermediates
    assert len(intermediates) == 1

def test_ode_sorted_assignments(parser, trans):
    expr = """
    states(V=0)
    dV_dt = x * z
    x = y + z + V
    z = a - 2
    a = y + 3
    y = 1
    """
    tree = parser.parse(expr)
    result = ode.make_ode(*trans.transform(tree), name="TestODE")
    assert result.sorted_assignments == ("y", "a", "z", "x", "dV_dt")

@pytest.mark.parametrize(
    "assignments_only, expected",
    [
        (True, ("y", "a", "z", "x", "dV_dt")),
        (False, ("V", "y", "a", "z", "x", "dV_dt")),
    ],
)
def test_sort_assignment(assignments_only, expected, parser, trans):
    expr = """
    states(V=0)
    dV_dt = x * z
    x = y + z + V
    z = a - 2
    a = y + 3
    y = 2 * V
    """
    tree = parser.parse(expr)
    result = ode.make_ode(*trans.transform(tree), name="TestODE")
    sorted_assignments = ode.sort_assignments(
        result.intermediates + result.state_derivatives,
        assignments_only=assignments_only,
    )
    assert sorted_assignments == expected


from dataclasses import dataclass
from pollination_dsl.function import Function, command, Inputs, Outputs

@dataclass
class TimeStepData(Function):
    """Write a time step data JSON file to be consumed by the TimeStepImages."""

    time_step_file = Inputs.file(
        description='A file generated by an annual recipe that has all the time steps.'
        ' an example of such a file is sun-up-hours.txt generated by the'
        ' direct-sun-hours recipe.',
        path='time_step_file.txt'
    )

    periods_file = Inputs.file(
        description='A file that contains the periods and color associated with each'
        ' period.',
        path='periods.json'
    )

    @command
    def time_step_data(self):
        return 'honeybee-vtk config time-step-data'\
            ' --time-step-file time_step_file.txt --periods-file periods.json'\
            ' --file-name _time_step_data'

    output_file = Outputs.file(
        description='Created file.',
        path='_time_step_data.json'
    )


from django.shortcuts import render, redirect
from django.db.models import Q
from django.http import HttpResponseRedirect, JsonResponse
from django.views.generic import TemplateView
from django.views.generic.base import View
from blogapi.models import *
import datetime

# Create your views here.

class AdminBaseView(TemplateView):
    def __init__(self):
        TemplateView.__init__(self)
        self.isAuthenticated = False
    def authenticate(self, request):
        if not request.user.is_authenticated():
            return HttpResponseRedirect('/admin/login/?next=/admin/blogs/')
        self.isAuthenticated = True
        return

class BlogsView(AdminBaseView):
    template_name = "admin/blogs.html"
    def get(self, request, *args, **kwargs):

        # make sure we're authenticated
        redirect = self.authenticate(request)
        if not self.isAuthenticated:
            return redirect

        blogs = Blog.objects.filter(page_type=Blog.BLOG)
        context = { "blogs": blogs }
        return render(request, self.template_name, context)

class PagesView(AdminBaseView):
    template_name = "admin/pages.html"
    def get(self, request):

        # make sure we're authenticated
        redirect = self.authenticate(request)
        if not self.isAuthenticated:
            return redirect

        blogs = Blog.objects.filter(~Q(page_type = Blog.BLOG))
        context = { "blogs": blogs }
        return render(request, self.template_name, context)

class DeleteBlogView(View):
    def get(self, request, *args, **kwargs):
        id = kwargs['id']
        blog = Blog.objects.get(pk=id)
        if blog != None: blog.delete()
        return HttpResponseRedirect('/admin/blogs/')

class DeleteContentView(View):
    def delete(self, request, *args, **kwargs):
        id = kwargs['id']
        content = Content.objects.get(pk=id)
        if content != None: content.delete()
        return JsonResponse({'id': content.id})

class BlogBaseView(AdminBaseView):
    def getItemFromList(self, list, index):
        try:
            return list[index]
        except IndexError:
            return None

    def getIdFromList(self, list, index):
        try:
            id = list[index]
            return id if int(id) != -1 else None
        except (IndexError, ValueError):
            return None
    def addBlogContents(self, blog, request, ids):
        orderList = request.POST.getlist('position')
        bodyList = request.POST.getlist('body')
        typeList = request.POST.getlist("type")
        captionList = request.POST.getlist('caption')
        imageList = request.FILES.getlist('image')
        imageFlagList = request.POST.getlist("hasImage")

        textType = ContentType.objects.filter(type_name = 'TEXT').first()
        imageType = ContentType.objects.filter(type_name = 'IMAGE').first()
        itemCounter = 0
        imgCounter = 0
        imgTypeCounter = 0
        bodyCounter = 0

        # update blog title/page type if we need to
        blog.blog_title = request.POST['title']
        blog.page_type = request.POST['page-type']
        blog.save()

        for typeId in typeList:

            try:
                typeId = int(typeId)
            except ValueError:
                typeId = 0

            if typeId == imageType.id:
                hasImage = self.getItemFromList(imageFlagList, imgTypeCounter)
                hasImage = bool(int(hasImage)) # convert "0"/"1" to false/true
                image = self.getItemFromList(imageList, imgCounter) if hasImage else None
                if (image != None):
                    imgCounter += 1
                    data = image.file.read()
                    ext = image.content_type
                    defaults = {
                        'content_type': imageType,
                        'content_data': data,
                        'file_extension': ext,
                        'created_date': datetime.datetime.now(),
                        'content_order': self.getItemFromList(orderList, itemCounter),
                        'content_caption': self.getItemFromList(captionList, imgTypeCounter)
                    }
                else:
                    defaults = {
                        'content_type': imageType,
                        'content_order': self.getItemFromList(orderList, itemCounter),
                        'content_caption': self.getItemFromList(captionList, imgTypeCounter)
                    }
                blog.contents.update_or_create(
                    id = self.getIdFromList(ids, itemCounter),
                    defaults = defaults
                )
                imgTypeCounter += 1
            if typeId == textType.id:
                text = self.getItemFromList(bodyList, bodyCounter)
                bodyCounter += 1
                blog.contents.update_or_create(
                    id = self.getIdFromList(ids, itemCounter),
                    defaults = {
                        'content_type': textType,
                        'content_text': text,
                        'created_date': datetime.datetime.now(),
                        'content_order': self.getItemFromList(orderList, itemCounter),
                    }
                )
            itemCounter += 1

class EditBlogView(BlogBaseView):
    template_name = "admin/blog-form.html"
    def get(self, request, **kwargs):

        # make sure we're authenticated
        redirect = self.authenticate(request)
        if not self.isAuthenticated:
            return redirect

        id = kwargs['id']
        blog = Blog.objects.get(pk=id)
        if blog != None:
            context = {
                "blog": blog,
                "types": ContentType.objects.all(),
                "pageTypes": Blog.PAGE_TYPES
            }
            return render(request, self.template_name, context)
        else:
            return HttpResponseRedirect('/admin/blogs/')

    def post(self, request, **kwargs):

        # make sure we're authenticated
        redirect = self.authenticate(request)
        if not self.isAuthenticated:
            return redirect

        blog = Blog.objects.get(pk=request.POST['blogid'])
        contentIds = request.POST.getlist('contentid')
        self.addBlogContents(blog, request, contentIds)

        link = '/admin/blogs/' if blog.page_type == Blog.BLOG else '/admin/pages/'
        return HttpResponseRedirect(link)

class AddNewBlogView(BlogBaseView):
    template_name = "admin/blog-form.html"
    def get(self, request, *args, **kwargs):

        # make sure we're authenticated
        redirect = self.authenticate(request)
        if not self.isAuthenticated:
            return redirect

        context = {
            "types": ContentType.objects.all(),
            "pageTypes": Blog.PAGE_TYPES
        }

        return render(request, self.template_name, context)

    def post(self, request):

        # make sure we're authenticated
        redirect = self.authenticate(request)
        if not self.isAuthenticated:
            return redirect

        contentIds = request.POST.getlist('contentid')
        blog = Blog(
            blog_title = request.POST['title'],
            pub_date = datetime.datetime.now(),
            user = request.user
        )
        blog.save()
        self.addBlogContents(blog, request, contentIds)
        return HttpResponseRedirect('/admin/blogs/')


#! /usr/bin/env python
from __future__ import print_function, division

from astropy.io import fits
from numba import njit, float64, prange
import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d

from astrobits.sourcecounts import dNdS

@njit([float64[:](float64[:], float64[:, :], float64, float64)], parallel=True)
def R(xs, beam, pixelsize, minS):
    """
    pixelsize: area of cell in steradians
    """
    Rxs = np.empty_like(xs)
    beam = beam.flatten()

    for i in prange(0, len(xs)):
        x = xs[i]
        S = np.absolute(x / beam)
        _dNdS = dNdS(S) * pixelsize  # differential source count per pixel
        _dNdS[S < minS] = 0
        _dNdS[S > 75] = 0
        Rxs[i] = np.sum(_dNdS / beam)

    return Rxs

def p(Rxs):
    rws = np.fft.fft(Rxs)
    return np.exp(rws - rws[0])

def P(pws):
    return np.fft.ifft(pws)

# Set up Gaussian beam
xs, ys = np.meshgrid(range(-200, 201), range(-200, 201))
scale = 1  #arcseconds
xs, ys = xs * scale, ys * scale
FWHM = 70  # arcseconds
sigma = FWHM / (2 * np.sqrt(2 * np.log(2)))
beam = np.exp(-(xs**2 + ys**2) / (2 * sigma**2))

pixelsize = scale / 3600  # degrees
pixelarea = pixelsize**2 * (np.pi**2 / 180**2)  # steradians

for minS in [0.15e-3, 0.1e-3, 0.05e-3, 0.025e-3, 0.01e-3, 0.005e-3, 0.001e-3]: # [1e-3, 0.75e-3, 0.5e-3, 0.25e-3, 0.1e-3]:
    # Calculate R(x) for a logarithmic spacing of xs
    print("Calculating Rlogxs")
    logxs = np.logspace(-7, 1.7, num=10000)
    dlogxs = logxs[1:] - logxs[:-1]
    logxs = (logxs[1:] + logxs[:-1]) / 2
    Rlogxs = R(logxs, beam, pixelarea, minS=minS)

    # Interpolate Rxs onto regular xs grid
    print("Interpolating onto linear array")
    xs = np.linspace(logxs[0], logxs[-1], 1e8)
    dx = (xs[-1] - xs[0]) / (len(xs) - 1)
    Rxs = interp1d(logxs, Rlogxs, 'linear')(xs) * dx
    print("Done")

    pws = p(Rxs)
    PD = P(pws)

    # Normalise PD = 1 across distribution
    PD = PD / np.sum(PD * dx)

    # Calculate the 'width' of the distribution
    PD_integral = np.cumsum(PD * dx)
    low_idx = np.argsort(abs(PD_integral - 0.25))[0]
    high_idx = np.argsort(abs(PD_integral - 0.75))[0]
    width = (xs[high_idx] - xs[low_idx]) / 1.349
    print("Width:", width)

    print("Done. Plotting...")
    idx = np.all([xs < 5e-3], axis=0)
    plt.plot(xs[idx], PD[idx], label="S > %g Jy; $\sigma$ %g Jy" % (minS, width))

plt.legend()
plt.xlabel("D [Jy/beam]")
plt.ylabel("P(D) [(Jy/beam)$^{-1}$]")
plt.savefig("confusion-%d-arcseconds.pdf" % FWHM, transparent=True)
plt.show()
exit()