"""Frequency analysis tools."""

import copy
import weakref
import numpy as np

from finesse.element import ModelElement
from finesse.parameter import float_parameter
from finesse.symbols import Symbol

from libc.stdlib cimport free, malloc

cdef class FrequencyContainer:
    """Contains an array of frequency objects and their associated
    faster access C struct information. 
    """
    def __cinit__(self, *args, **kwargs):
        self.size = 0
        self.frequency_info = NULL
        self.carrier_frequency_info = NULL

    def __init__(self, frequencies, FrequencyContainer carrier_cnt=None):
        self.size = len(frequencies)
        self.frequencies = tuple(frequencies)
        self.frequency_info = <frequency_info_t*> malloc(self.size * sizeof(frequency_info_t))
        if not self.frequency_info:
            raise MemoryError()
        if carrier_cnt:
            if not carrier_cnt.frequency_info:
                raise MemoryError()
            self.carrier_frequency_info = carrier_cnt.frequency_info
        else:
            self.carrier_frequency_info = NULL

    def __dealloc__(self):
        if self.size:
            free(self.frequency_info)

    def get_info(self, Py_ssize_t index):
        if not (0 <= index <= self.size):
            raise IndexError()
        rtn = {
            "f" : self.frequency_info[index].f,
            "index" : self.frequency_info[index].index,
            "audio_lower_index" : self.frequency_info[index].audio_lower_index,
            "audio_upper_index" : self.frequency_info[index].audio_upper_index,
            "audio_order" : self.frequency_info[index].audio_order,
            "audio_carrier_index" : self.frequency_info[index].audio_carrier_index,
        }
        return rtn

    cdef initialise_frequency_info(self):
        cdef Py_ssize_t i, j, cidx
        cdef FrequencyContainer fcnt

        for i in range(self.size):
            self.frequency_info[i].f = <double>self.frequencies[i].f
            self.frequency_info[i].index = <Py_ssize_t>self.frequencies[i].index
            self.frequency_info[i].audio_order = <int>self.frequencies[i].audio_order

            if self.frequency_info[i].audio_order:
                assert(self.carrier_frequency_info)
                cidx = <Py_ssize_t>self.frequencies[i].audio_carrier_index
                self.frequency_info[i].audio_carrier_index = cidx
                self.frequency_info[i].f_car = &self.carrier_frequency_info[cidx].f
                # Update carrier info so it knows about this sideband
                if self.frequency_info[i].audio_order == 1:
                    self.carrier_frequency_info[cidx].audio_upper_index = <Py_ssize_t>self.frequencies[i].index
                elif self.frequency_info[i].audio_order == -1:
                    self.carrier_frequency_info[cidx].audio_lower_index = <Py_ssize_t>self.frequencies[i].index
                else:
                    raise Exception("Unexpected")
        
        # if self.is_audio:
        #     for i in range(len(self.unique_fcnt)):
        #         fcnt = self.unique_fcnt[i]
        #         for j in range(fcnt.size):
        #             fcnt.frequency_info[j].f = <double>fcnt.frequencies[j].f
        #             fcnt.frequency_info[j].index = <Py_ssize_t>fcnt.frequencies[j].index

    cdef update_frequency_info(self):
        cdef Py_ssize_t i, j, cidx
        cdef FrequencyContainer fcnt

        for i in range(self.size):
            self.frequency_info[i].f = <double>self.frequencies[i].f

def generate_frequency_list(model):
    """
    For a given model a symbolic list of frequencies is generated.
    The result can be used to generate a set of frequencies bins to
    be modelled in a simulation.

    This method relies on using :class:`.Symbol`. Using
    symbolic statements this method attempts to isolate uniqe frequency
    bins whilst leaving those changing during a simulation present.

    Returns
    -------
    List of :class:`.Symbol`
    """

    def unique_indices(arr):
        """
        Simple unique element finder which doesn't require any
        greater or less than operations. This isn't tuned for
        efficiency at all.
        """
        lst = list(arr)
        lst2 = list(arr)
        for i in range(len(lst)):
            if lst.count(lst2[i]) > 1:
                lst[i] = None

        return [_ for _, __ in enumerate(lst) if __ is not None]

    fn_eval = np.vectorize(lambda x: x.eval())
    fn_eval2 = np.vectorize(lambda x: x.eval(keep_changing_symbols=True))
    # fn_subs = np.vectorize(
    #     lambda x, **kwargs: x.eval(keep_changing_symbols=True, subs=kwargs)
    # )
    fn_is_changing = np.vectorize(lambda x: x.is_changing)

    source_frequencies = []
    source_components = []
    modulation_frequencies = []

    for comp in model._frequency_generators:
        s = comp._source_frequencies()
        source_frequencies.extend(s)
        source_components.extend((comp,) * len(s))

        m = comp._modulation_frequencies()
        modulation_frequencies.extend(m)

    # Now to prune the frequency list
    Nm = len(modulation_frequencies)
    Ns = len(source_frequencies)

    if Ns == 0:
        raise Exception("There are no source frequencies present in the model")

    if Nm == 0:
        Fsym = np.array(source_frequencies)
    else:
        # First we make a list with all possible combinations of frequencies
        Fsym = (
            np.vstack((np.atleast_2d(modulation_frequencies),) * Ns)
            + np.hstack((np.atleast_2d(source_frequencies).T,) * Nm)
        ).flatten()
        Fsym = np.hstack((source_frequencies, Fsym))

    # Take all the frequency values which definitely won't be changing
    not_changing = Fsym[np.bitwise_not(fn_is_changing(Fsym))]
    if len(not_changing) == 0:
        not_changing = []
    else:
        # ... and select all those which are unique
        values, idx, rev, counts = np.unique(fn_eval(not_changing), True, True, True)
        not_changing = not_changing[idx]

    # First select only the changing frequency bins
    changing = Fsym[fn_is_changing(Fsym)]
    if len(changing) == 0:
        changing = []
    else:
        # We need to use a different unique finding function that doesn't rely on
        # using > or < than comparisons like np.unique. I'm not sure how to implement
        # > and < for Symbols sensibly
        idx = unique_indices(fn_eval2(changing))
        changing = changing[idx]

    # Finally sort the indicies so that upper and lower sidebands are grouped together. This
    # help with numerical errors when iterating outputs later so that similarly sized elements are
    # grouped this isn't always true though, just generally in regards to upper and lower sidebands.
    final = np.hstack((changing, not_changing))
    srt_idx = np.argsort(abs(fn_eval(final)))
    return final[srt_idx]


@float_parameter("f", "Frequency", units="Hz", validate="_validate_fsig", is_default=True)
class Fsig(ModelElement):
    def __init__(self, name, value):
        super().__init__(name)
        self.f = value

    def _validate_fsig(self, value):
        if value is None or isinstance(value, Symbol):
            return value
        elif value <= 0:
            raise Exception("fsig value must be > 0 Hz")
        else:
            return value


cdef class Frequency:
    """Represents a frequency "bin" with a specific index.

    The value of the frequency is calculated from the name of the frequency.

    Parameters
    ----------
    name : str
        Name of the frequency.

    order : int, optional
        The order of the frequency, defaults to zero.
    """
    def __init__(
        self,
        name,
        sim,
        symbol,
        index=None,
        audio=False,
        audio_carrier_index=None,
        audio_order=0,
    ):
        self.__name = name
        self.__symbol = symbol
        self.__index = index
        self.__is_audio = audio
        #self.__sim = sim
        self.__symbol_changing = symbol.is_changing
        self.__start_value = self.__symbol.eval()
        self.__lambdified = eval(f"lambda: {symbol}", sim.model.elements)

        if audio:
            if audio_carrier_index is None:
                raise Exception("Audio frequency carrier must be specified")

            if audio_order not in (-1, 1):
                raise Exception("Audio frequency order must be -1 or +1")

            self.__order = audio_order
            self.__carrier = audio_carrier_index
        else:
            self.__order = 0
            self.__carrier = 0

    def __repr__(self):
        return (
            f"<{self.__class__.__name__} {self.name} (f={self.symbol}={self.f}) at "
            f"{ hex(id(self)) }>"
        )

    def __str__(self):
        return f"<{self.name} is {self.f}Hz (Frequency)>"

    def __deepcopy__(self, memo):
        raise Exception(
            "Frequency objects cannot be deepcopied as they are associated with Simulations"
        )

    #@property
    #def simulation(self):
    #    return self.__sim

    @property
    def symbol(self):
        return self.__symbol

    @property
    def f(self):
        if self.__symbol_changing:
            return self.__lambdified()
        else:
            return self.__start_value

    @property
    def is_audio(self):
        """Is this an audio sideband frequency?"""
        return self.__is_audio

    @property
    def audio_carrier_index(self):
        """The carrier frequency.

        :getter: Returns the carrier frequency index (read-only).
        """
        return self.__carrier

    @property
    def name(self):
        """Name of the frequency object.

        :getter: Returns the name of the frequency (read-only).
        """
        return self.__name

    @property
    def audio_order(self):
        """Audio modulation order of this frequency.

        :getter: Returns the order of the frequency (read-only).
        """
        return self.__order

    @property
    def index(self):
        """Index of the frequency object.

        :getter: Returns the index of the frequency (read-only).
        """
        return self.__index
