"""Actions to compute noise projections and budgets."""

import logging
from collections import defaultdict

import numpy as np

from ...detectors.compute.quantum import QShot0Workspace, QShotNWorkspace
from ...env import warn
from ...solutions import BaseSolution
from .base import Action
from .lti import FrequencyResponse

LOGGER = logging.getLogger(__name__)


class NoiseProjectionSolution(BaseSolution):
    def plot(self, output_node=None, lower=0.1, upper=3, *, ax=None, **kwargs):
        import matplotlib.pyplot as plt

        if output_node is None:
            output_node = self.output_nodes[0]

        if ax is None:
            fig = plt.gcf()
            if len(fig.axes) == 0:
                fig.subplots(1, 1)
            ax = fig.axes[0]

        total = np.sqrt((self.out[output_node] ** 2).sum(1))
        rng = lower * total.min(), upper * total.max()
        noises_to_plot = np.any(self.out[output_node] > rng[0], 0)
        if any(noises_to_plot):
            ax.loglog(self.f, np.abs(self.out[output_node][:, noises_to_plot]))
            ax.loglog(self.f, np.abs(total), c="k", ls="-.", lw=2)
            ax.legend((*np.array(self.noises)[noises_to_plot], "Total"))
            ax.set_ylim(*rng)
        else:
            warn("No noise data to plot in this solution")

        ax.set_ylabel(
            f"ASD [{output_node if not self.scaling else self.scaling}/$\\sqrt{{\\mathrm{{Hz}}}}$]"
        )
        ax.set_xlabel("Frequency [Hz]")


# IMPORTANT: renaming this class impacts the katscript spec and should be avoided!
class NoiseProjection(Action):
    def __init__(self, f, *output_nodes, scaling=None, name="loop"):
        if len(output_nodes) == 0:
            raise ValueError(
                "At least one output node must be specified to compute noise projection to"
            )
        super().__init__(name)
        process = lambda x: x.full_name if type(x) is not str else x

        self.f = f
        self.scaling = process(scaling) if scaling is not None else None
        self.output_nodes = tuple(process(o) for o in output_nodes)

        if len(self.output_nodes) > len(set(self.output_nodes)):
            raise ValueError(
                f"The same output node has been requested multiple times {self.output_nodes}"
            )

    def _do(self, state):
        sol = NoiseProjectionSolution(self.name)
        sol.f = self.f
        sol.output_nodes = self.output_nodes
        sol.scaling = self.scaling
        # create a list of callables func(fsig) to get the ASD noises
        noise_ASDs = {
            name: el.ASD.lambdify(state.model.fsig.f)
            for name, el in state.model.noises.items()
        }
        # labels for noises
        sol.noises = list(el.name for _, el in state.model.noises.items())
        # Keep track of which nodes have what noise injected into them
        noise_node_map = defaultdict(list)
        for el in state.model.noises:
            noise_node_map[el.node.full_name].append(el.name)
        # Collect any extra outputs that should be calculated during the fsig sweep. This
        # is to make efficient use of the filling and solving that this is already doing
        # to extract quantum noise, or others, as needed. Some of these outputs will be
        # signal frequency independant, such as standard shot-noise calculations, so just
        # compute them once.
        # TODO eventually handle qnoised detectors, which are frequency dependant
        fsig_indep_output = []
        for dws in state.sim.readout_workspaces:
            if isinstance(dws, (QShot0Workspace, QShotNWorkspace)):
                added = 0  # don't calculate anything if we aren't modelling the nodes
                # The quantum shot noise detectors will be on the optiical
                # input node, which we can't inject a signal into for computing
                # the noise propagation. Here we need to get the electrical outputs
                # of the readout and just put the noise in there
                for n in dws.owner.signal_nodes:
                    if n.full_name in state.sim.signal.nodes:
                        noise_node_map[n.full_name].append(dws.oinfo.name)
                        added += 1
                if added:
                    fsig_indep_output.append(dws)

        # None actually added
        if len(fsig_indep_output) == 0:
            fsig_indep_output = None

        # Compute all the required transfer functions for noise propagation
        # NOTE: We use _do directly here because we just want to call the action
        # on this state, rather than `run` which will try and create a new state.
        # This is fine, as long as we have requested all the options it needs in
        # _requests. We can't make this frequency response in the init as we do
        # not have the model to grab all the various noise and shot-noise nodes
        self.input_nodes = tuple(noise_node_map.keys())
        sol.freqresp = FrequencyResponse(
            self.f, self.input_nodes, self.output_nodes
        )._do(state, fsig_indep_output)

        # Get any shot noise outputs from the solution
        if fsig_indep_output is not None:
            for dws in fsig_indep_output:
                # Make a simple callable to work wiht the noise ASD functions
                noise_ASDs[dws.oinfo.name] = lambda f: sol.freqresp.extra_outputs[
                    dws.oinfo.name  # noqa: B023
                ]
                sol.noises.append(dws.oinfo.name)

        # get a map from nodes->noise, for noise->node index in output
        inv_noise_node_map = {}
        for k, v in noise_node_map.items():
            for n in v:
                inv_noise_node_map[n] = self.input_nodes.index(k)
        # Convert all the ASDs into PSDs
        sol.PSDs = np.array(
            tuple(np.ones_like(self.f) * fn(self.f) ** 2 for fn in noise_ASDs.values()),
            dtype=float,
        ).T
        # Use this to broadcast from the frequency response output to get the right
        # transfer function for each noise source
        inp_indices = tuple(inv_noise_node_map[name] for name in noise_ASDs.keys())

        # Here we can compute some projection for calculating equivalent noise budgets
        if self.scaling:
            sol.scaling_solution = FrequencyResponse(
                self.f, self.scaling, self.output_nodes, open_loop=True
            )._do(state)
        sol.out = {}
        # compute abs(H)**2 for noise projection of PSDS
        HH = np.zeros(
            (len(self.f), len(self.output_nodes), len(inp_indices)), dtype=float
        )
        np.abs(sol.freqresp.out[:, :, inp_indices], out=HH)
        np.multiply(HH, HH, out=HH)
        # The final index of HH is the output node index, so we can quickly iterate over
        # them here to project the noises
        for i, output_node in enumerate(self.output_nodes):
            # sqrt(output**2/node**2 * node**2/Hz) => output/rtHz
            sol.out[output_node] = np.sqrt(HH[:, i, :] * sol.PSDs)
            if self.scaling:
                # output / scaling
                # sqrt(output**2/node**2 * node**2/Hz) => scaling/rtHz
                sol.out[output_node] /= np.abs(sol.scaling_solution.out[:, 0, :])
        return sol

    def _requests(self, model, memo, first=True):
        memo["changing_parameters"].append("fsig.f")
        memo["output_nodes"].extend(self.output_nodes)
        if self.scaling:
            memo["input_nodes"].append(self.scaling)
        memo["input_nodes"].extend(
            (el.node.full_name for n, el in model.noises.items())
        )
