import numpy as np
from collections import defaultdict
import functools
import operator
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.text import Text
from mpl_toolkits.mplot3d import Axes3D

# create docs
# sphinx-apidoc -F -o docs rf_linkbudget/

try:
    from .rfmath import RFMath
except ImportError:
    from rfmath import RFMath


class simResult():
    """
    | A class which gets generated by Circuit.simulate(...)
    | It hold convenience functions for plotting and accessing the right datsets

    Parameters
    ----------
    network : str
        circuit network
    data : dictionary
        circuit data, from a simulation
    **kwargs : kwargs
        all other kwargs gets saved in a variable with the same name
    """

    def __init__(self, network, data, **kwargs):
        """
        Parameters
        ----------
        network : str
            circuit network
        data : dictionary
            circuit data, from a simulation
        **kwargs : kwargs
            all other kwargs gets saved in a variable with the same name
        """
        self.network = network
        self.nodes = set()  # all different nodes in network
        self.data = data
        self.freq = [key for key in data]
        self.power = [key for key in data[self.freq[0]]]

        self.graphes = set(functools.reduce(operator.iconcat, [list(network[x][y].nodes) for x in network for y in network[x]], []))
        self.nodes = set([g.parent.name for g in self.graphes])

        for key, value in kwargs.items():
            setattr(self, key, value)

        dat = defaultdict(dict)  # convert dictionarys to pandas DataFrame
        for freq in data:
            for pwr in data[freq]:
                dat[pwr][freq] = pd.DataFrame(data[freq][pwr])

        self.pandas = pd.DataFrame(dat)

    @property
    def param(self):
        """
        | access the parameters as a list

        Returns
        ----------
        parameter : list
            list of parameters
        """
        tmp = self.pandas.applymap(lambda x: list(x.index))
        return set(functools.reduce(operator.concat, functools.reduce(operator.concat, tmp.values.tolist())))

    def __repr__(self):
        out = 'simulation result:\n'
        out += '@ frequencies : {}\n'.format(self.freq)
        out += '@ input power : {}\n'.format(self.power)
        return out

    def setNoiseBandwidth(self, B):
        """
        | set the noisebandwidth
        | changes the following parameters

            - n : noise power
            - SNR : signal to noise ratio
            - DYN : dynamic
        """

        df = self.pandas

        def setNoise(x):
            for i in range(len(x.loc['n'])):
                tn = x.loc['Tn'].iat[i]
                p = x.loc['p'].iat[i]
                sfdr = x.loc['SFDR'].iat[i]

                n = RFMath.convert_T_n(tn, B)
                snr = RFMath.calc_SNR(p, n)
                dyn = RFMath.calc_Dynamic(snr, sfdr)

                x.loc['n'].iat[i] = n
                x.loc['SNR'].iat[i] = snr
                x.loc['DYN'].iat[i] = dyn

            return x

        df = df.applymap(lambda x: setNoise(x))
        self.pandas = df

    def extractValues(self, param, freq, power, keys=None):
        """
        | convencience function to access the data in a structured way

        Parameters
        ----------
        param : str or list of str
            a single parameter or a list of parameters which shall be extracted
        freq : float or list of floats
            a single frequency or a list of frequencies which shall be extracted
        power : float or list of floats
            a single power value or a list of power values which shall be extracted
        keys : str or list of str, optional
            a single device port or a list of device ports.
            only these values are extracted.

        Returns
        ----------
        params : (ports, data)
            a tuple with all ports and the corresponding data
        """
        df = self.pandas.loc[freq][power]
        val = df.loc[param].values
        ports = list(df.loc[param].index)

        if keys is not None:
            ports = keys
            val = df.loc[param].reindex(keys, fill_value=None)

        return (ports, val)

    def extractLastValues(self, key, freq=None, power=None):
        """
        | convencience function to access the data at the last device port.
        | an easy way to export all "end" values to compare systems with each other

        Parameters
        ----------
        key : str
            a single device port or a list of device ports.
            only these values are extracted.
        freq : float, optional
            a single frequency or a list of frequencies which shall be extracted
        power : float, optional
            a single power value or a list of power values which shall be extracted

        Returns
        ----------
        params : pandas.DataFrame
            data
        """
        df = self.pandas
        df = df.applymap(lambda x: x.iloc[:, -1][key])

        if freq is not None and power is None:
            df = df.loc[freq][:]
        elif freq is None and power is not None:
            df = df.loc[:][power]
        elif freq is not None and power is not None:
            df = df.loc[freq][power]
        else:
            pass

        return df

    def plot_chain(self, param=None, freq=None, power=None, keys=None):
        """
        | plots a window with responsive fields which can be switched between (click or mousewheel)
        | plot_chain plots all data from the start to the end, by all ports between
        | we can select which ports to show, be defining the keys
        | when param, freq or power is not defined, all parameters will be shown

        Parameters
        ----------
        param : str or list of str, optional
            a single parameter or a list of parameters which shall be plotted
        freq : float or list of floats, optional
            a single frequency or a list of frequencies which shall be plotted
        power : float or list of floats, optional
            a single power value or a list of power values which shall be plotted
        keys : str or list of str, optional
            a single device port or a list of device ports.
            only these values are extracted.

        Notes
        -----
        return handler must remain, otherwise the plot will not be responsive anymore

        Returns
        ----------
        handler : callback handlers
        """
        freq = self.freq if freq is None else freq
        power = self.power if power is None else power
        param = self.param if param is None else param
        freq = freq if type(freq) is list else [freq]
        power = power if type(power) is list else [power]
        param = param if type(param) is set() else list(param)
        param = param if type(param) is list else [param]

        fig, ax = plt.subplots()
        fig.subplots_adjust(bottom=0.2)
        l, = plt.plot([], [])
        plt.xticks(rotation=45)
        ax.grid(True)

        bbox_props = dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9, )
        tb = [ax.text(1.015, 0.95, param[0], picker=True, bbox=bbox_props, horizontalalignment='left', rotation=0, transform=ax.transAxes),
              ax.text(1.015, 0.85, freq[0], picker=True, bbox=bbox_props, horizontalalignment='left', rotation=0, transform=ax.transAxes),
              ax.text(1.015, 0.75, power[0], picker=True, bbox=bbox_props, horizontalalignment='left', rotation=0, transform=ax.transAxes)]

        class Index(object):
            v = {'param': param, 'freq': freq, 'power': power}
            i = dict(zip(['param', 'freq', 'power'], [0, 0, 0]))
            t = dict(zip(['param', 'freq', 'power'], tb))
            ref = {'fig': fig, 'plot': l, 'axis': ax}
            parent = self

            def _cb_rotate(self, event, direction, idx):
                self.i[idx] = np.mod(self.i[idx] + direction, len(self.v[idx]))
                self.redraw()

            def redraw(self):
                cparam, cfreq, cpower = [self.v['param'][self.i['param']], self.v['freq'][self.i['freq']], self.v['power'][self.i['power']]]
                xt, val = self.parent.extractValues(cparam, cfreq, cpower, keys=keys)
                self.t['param'].set_text(cparam)
                self.t['freq'].set_text('{:.1f}MHz'.format(cfreq/1e6))
                self.t['power'].set_text(cpower)
                self.ref['plot'].set_data(range(len(xt)), val)
                self.ref['axis'].set_xticks(range(len(xt)))
                self.ref['axis'].set_xticklabels(xt)
                self.ref['axis'].set_ylabel(cparam)
                self.ref['axis'].set_title("Performance: freq={:.1f}MHz, Pin={:.1f}dBm".format(cfreq/1e6, cpower), picker=True)
                self.ref['axis'].relim()
                self.ref['axis'].autoscale_view()
                self.ref['fig'].canvas.draw_idle()

        callback = Index()
        callback.redraw()

        def onpick1(event):
            if isinstance(event.artist, Text):
                text = event.artist
                direction = 1 if event.mouseevent.button == 1 or event.mouseevent.button == 'up' else -1
                if callback.t['param'] == text:
                    callback._cb_rotate(event=event, direction=direction, idx='param')
                elif callback.t['freq'] == text:
                    callback._cb_rotate(event=event, direction=direction, idx='freq')
                elif callback.t['power'] == text:
                    callback._cb_rotate(event=event, direction=direction, idx='power')

        cid = fig.canvas.mpl_connect('pick_event', onpick1)
        handler = (cid, callback)
        return handler

    def plot_total(self, param=None, freq=None, power=None):
        """
        | plots a window with responsive fields which can be switched between (click or mousewheel)
        | plot_total plots all "and" data in respect to frequency or power
        | we can select which ports to show, be defining the keys
        | when param, freq or power are not defined, all parameters will be shown

        Parameters
        ----------
        param : str or list of str, optional
            a single parameter or a list of parameters which shall be plotted
        freq : float or list of floats, optional
            a single frequency or a list of frequencies which shall be plotted
        power : float or list of floats, optional
            a single power value or a list of power values which shall be plotted

        Notes
        -----
        return handler must remain, otherwise the plot will not be responsive anymore

        Returns
        ----------
        handler : callback handlers
        """
        freq = self.freq if freq is None else freq
        power = self.power if power is None else power
        param = self.param if param is None else param
        freq = freq if type(freq) is list else [freq]
        power = power if type(power) is list else [power]
        param = param if type(param) is set() else list(param)
        param = param if type(param) is list else [param]

        if len(freq) == 1:
            varDependend = 'freq'
        elif len(power) == 1:
            varDependend = 'power'
        else:
            varDependend = 'freq'

        fig, ax = plt.subplots()
        fig.subplots_adjust(bottom=0.2)
        l, = plt.plot([], [])
        ax.grid(True)

        bbox_props = dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9, )
        tb = [ax.text(1.015, 0.95, param[0], picker=True, bbox=bbox_props, horizontalalignment='left', rotation=0, transform=ax.transAxes),
              ax.text(1.015, 0.85, freq[0], picker=True, bbox=bbox_props, horizontalalignment='left', rotation=0, transform=ax.transAxes),
              ax.text(1.015, 0.75, power[0], picker=True, bbox=bbox_props, horizontalalignment='left', rotation=0, transform=ax.transAxes)]

        class Index(object):
            v = {'param': param, 'freq': freq, 'power': power}
            i = dict(zip(['param', 'freq', 'power'], [0, 0, 0]))
            t = dict(zip(['param', 'freq', 'power'], tb))
            ref = {'fig': fig, 'plot': l, 'axis': ax}
            varDep = varDependend
            parent = self

            def _cb_rotate(self, event, direction, idx):
                self.i[idx] = np.mod(self.i[idx] + direction, len(self.v[idx]))
                self.redraw()

            def redraw(self):
                cparam, cfreq, cpower = [self.v['param'][self.i['param']], self.v['freq'][self.i['freq']], self.v['power']] if self.varDep == 'freq' else [self.v['param'][self.i['param']], self.v['freq'], self.v['power'][self.i['power']]]
                df = self.parent.extractLastValues(cparam, cfreq, cpower)
                self.t['param'].set_text(cparam)
                self.t['freq'].set_text('{:.1f}MHz'.format(cfreq/1e6)) if self.varDep == 'freq' else self.t['freq'].set_text('freq')
                self.t['power'].set_text('power') if self.varDep == 'freq' else self.t['power'].set_text(cpower)
                self.ref['plot'].set_data(list(df.index), list(df.values))
                self.ref['axis'].set_ylabel(cparam)
                self.ref['axis'].set_title("Concat {} plot".format(cparam))
                self.ref['axis'].relim()
                self.ref['axis'].autoscale_view()
                self.ref['fig'].canvas.draw_idle()

            def swapDependendVar(self):
                if self.varDep == 'freq':
                    self.varDep = 'power'
                else:
                    self.varDep = 'freq'
                self.redraw()

        callback = Index()
        callback.redraw()

        def onpick1(event):
            if isinstance(event.artist, Text):
                text = event.artist
                direction = 1 if event.mouseevent.button == 1 or event.mouseevent.button == 'up' else -1

                if event.mouseevent.button == 1:
                    # clicked on item
                    if callback.t['freq'] == text or callback.t['power'] == text:
                        callback.swapDependendVar()
                else:
                    if callback.t['param'] == text:
                        callback._cb_rotate(event=event, direction=direction, idx='param')
                    elif callback.t['freq'] == text:
                        callback._cb_rotate(event=event, direction=direction, idx='freq')
                    elif callback.t['power'] == text:
                        callback._cb_rotate(event=event, direction=direction, idx='power')
                    else:
                        pass

        cid = fig.canvas.mpl_connect('pick_event', onpick1)
        handler = (cid, callback)
        return handler

    def plot_total_simple(self, param, freq=None, power=None, axis=None):
        """
        | plots a window with the defined parameter
        | plot_total_simple is a simpler variant of the other plots
        | it is not a responvice view, but it is easy to show different parameters at once

        Parameters
        ----------
        param : str or list of str
            a single parameter or a list of parameters which shall be plotted
        freq : float or list of floats, optional
            a single frequency or a list of frequencies which shall be plotted
        power : float or list of floats, optional
            a single power value or a list of power values which shall be plotted
        axis : axis object, optional
            currently not in use

        Notes
        -----
        return handler must remain, otherwise the plot will not be responsive anymore

        Examples
        ---------
        >>> t = sim1.plot_total_simple(['SFDR', 'SNR', 'DYN'], freq=430e6)
        >>> plt.show()

        Returns
        ----------
        handler : callback handlers
        """
        # Todo, use axis parameter to show different plots in one...
        fig, axis = plt.subplots()

        for p in param:
            df = self.extractLastValues(p, freq, power)

            if freq is not None and power is None:
                plt.plot(df, label=p)
                axis.set_xlabel('input frequency [MHz]')
            elif freq is None and power is not None:
                plt.plot(df.index/1e6, df.values, label=p)
                axis.set_xlabel('input power [dBm]')
            elif freq is not None and power is not None:
                raise ValueError('is a scalar')
            else:
                raise ValueError('two dimensional plot')
        # axis.set_ylabel('Values')
        plt.xticks(rotation=45)
        plt.xticks(df.index, df.index)
        plt.legend()
        fig.set_tight_layout({"pad": .0})
        plt.grid(True)

    def plot_surface(self, param=None, freq=None, power=None, keys=None):
        """
        | plots a window with responsive fields which can be switched between (click or mousewheel)
        | plot_surface plots a surface plot with all data in respect to al ports and power or frequency
        | we can select which ports to show, be defining the keys
        | when param, freq or power are not defined, all parameters will be shown

        Parameters
        ----------
        param : str or list of str, optional
            a single parameter or a list of parameters which shall be plotted
        freq : float or list of floats, optional
            a single frequency or a list of frequencies which shall be plotted
        power : float or list of floats, optional
            a single power value or a list of power values which shall be plotted
        keys : str or list of str, optional
            a single device port or a list of device ports.
            only these values are extracted.

        Notes
        -----
        return handler must remain, otherwise the plot will not be responsive anymore

        Returns
        ----------
        handler : callback handlers
        """
        freq = self.freq if freq is None else freq
        power = self.power if power is None else power
        param = self.param if param is None else param
        freq = freq if type(freq) is list else [freq]
        power = power if type(power) is list else [power]
        param = param if type(param) is set() else list(param)
        param = param if type(param) is list else [param]

        # create figure
        fig = plt.figure()
        ax = fig.gca(projection='3d')

        surf = ax.plot([], [])
        fig.tight_layout()

        class Index(object):
            mesh = np.meshgrid(power, list(range(len(keys))))
            v = {'param': param, 'freq': freq, 'power': power}
            i = dict(zip(['param', 'freq'], [0, 0]))
            t = dict(zip(['param', 'freq'], [0, 0]))
            bbox_props = dict(boxstyle="round", fc="w", ec="0.5", alpha=0.9, )
            ref = {'fig': fig, 'plot': surf, 'axis': ax}
            ports = keys
            parent = self

            def _cb_rotate(self, event, direction, idx):
                self.i[idx] = np.mod(self.i[idx] + direction, len(self.v[idx]))
                self.redraw()

            def redraw(self):
                cparam, cfreq, cpower = [self.v['param'][self.i['param']], self.v['freq'][self.i['freq']], self.v['power']]

                (X, Y) = self.mesh
                Z = np.zeros(shape=(X.shape[0], X.shape[1]))
                for col, c in enumerate(Z):
                    for row, r in enumerate(c):
                        xt, val = self.parent.extractValues(cparam, freq=cfreq, power=cpower[row], keys=[self.ports[col]])
                        Z[col][row] = val.values[0]

                self.ref['axis'].cla()
                self.t['param'] = self.ref['axis'].text2D(0.05, 0.95, cparam, transform=ax.transAxes, bbox=self.bbox_props, picker=True)
                self.t['freq'] = self.ref['axis'].text2D(0.05, 0.90, cfreq, transform=ax.transAxes, bbox=self.bbox_props, picker=True)
                self.ref['axis'].plot_surface(X, Y, Z, cmap=cm.coolwarm,
                                              linewidth=1, antialiased=True, alpha=0.75, vmin=-174, vmax=-156, edgecolor='none')

                self.ref['axis'].set_title("surface plot : {}".format(cparam))
                self.ref['axis'].set_yticks(range(len(self.ports)))
                self.ref['axis'].set_yticklabels([str(x) for x in self.ports])
                self.ref['axis'].relim()
                self.ref['axis'].autoscale_view()
                self.ref['fig'].canvas.draw_idle()
                pass

        callback = Index()
        callback.redraw()

        def onpick1(event):
            if isinstance(event.artist, Text):
                text = event.artist
                direction = 1 if event.mouseevent.button == 1 or event.mouseevent.button == 'up' else -1

                if callback.t['param'] == text:
                    callback._cb_rotate(event=event, direction=direction, idx='param')
                elif callback.t['freq'] == text:
                    callback._cb_rotate(event=event, direction=direction, idx='freq')
                else:
                    pass

        cid = fig.canvas.mpl_connect('pick_event', onpick1)
        handler = (cid, callback)
        return handler

# ============================================================================


if __name__ == "__main__":
    pass
