import numpy as np

from foxes.core import PartialWakesModel
from foxes.models.wake_models.dist_sliced import DistSlicedWakeModel
from foxes.models.rotor_models.grid import GridRotor
from foxes.utils import wd2uv, uv2wd
import foxes.variables as FV
import foxes.constants as FC


class PartialDistSlicedWake(PartialWakesModel):
    """
    Partial wakes for distance sliced wake models,
    making use of their structure.

    The evaluations are optinally done on a grid rotor
    that can differ from the algorithm's rotor model.

    Parameters
    ----------
    n : int, optional
        The `GridRotor`'s `n` parameter
    wake_models : list of foxes.core.WakeModel, optional
        The wake models, default are the ones from the algorithm
    wake_frame : foxes.core.WakeFrame, optional
        The wake frame, default is the one from the algorithm
    rotor_model : foxes.core.RotorModel, optional
        The rotor model, default is the one from the algorithm
    **kwargs : dict, optional
        Additional parameters for the `GridRotor`

    Attributes
    ----------
    rotor_model : foxes.core.RotorModel
        The rotor model, default is the one from the algorithm
    grotor : foxes.models.rotor_models.GridRotor
        The grid rotor model

    """

    def __init__(
        self, n=None, wake_models=None, wake_frame=None, rotor_model=None, **kwargs
    ):
        super().__init__(wake_models, wake_frame)

        self.rotor_model = rotor_model
        self.grotor = None if n is None else GridRotor(n=n, calc_vars=[], **kwargs)

    def __repr__(self):
        if self.grotor is not None:
            return super().__repr__() + f"(n={self.grotor.n})"
        elif self.rotor_model is not None and isinstance(self.rotor_model, GridRotor):
            return super().__repr__() + f"(n={self.rotor_model.n})"
        else:
            return super().__repr__()

    def initialize(self, algo, verbosity=0):
        """
        Initializes the model.

        This includes loading all required data from files. The model
        should return all array type data as part of the idata return
        dictionary (and not store it under self, for memory reasons). This
        data will then be chunked and provided as part of the mdata object
        during calculations.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        verbosity : int
            The verbosity level, 0 = silent

        Returns
        -------
        idata : dict
            The dict has exactly two entries: `data_vars`,
            a dict with entries `name_str -> (dim_tuple, data_ndarray)`;
            and `coords`, a dict with entries `dim_name_str -> dim_array`

        """
        idata = super().initialize(algo, verbosity)

        if self.rotor_model is None:
            self.rotor_model = algo.rotor_model
        if self.grotor is None:
            self.grotor = self.rotor_model

        for w in self.wake_models:
            if not isinstance(w, DistSlicedWakeModel):
                raise TypeError(
                    f"Partial wakes '{self.name}': Cannot be applied to wake model '{w.name}', since not an DistSlicedWakeModel"
                )

        self.YZ = self.var("YZ")
        self.W = self.var(FV.WEIGHT)

        algo.update_idata([self.rotor_model, self.grotor], 
            idata=idata, verbosity=verbosity)
        
        return idata

    def new_wake_deltas(self, algo, mdata, fdata):
        """
        Creates new initial wake deltas, filled
        with zeros.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        mdata : foxes.core.Data
            The model data
        fdata : foxes.core.Data
            The farm data

        Returns
        -------
        wake_deltas : dict
            Keys: Variable name str, values: any

        """
        n_rpoints = self.grotor.n_rotor_points()
        n_points = fdata.n_turbines * n_rpoints
        wake_deltas = {}
        for w in self.wake_models:
            w.init_wake_deltas(algo, mdata, fdata, n_points, wake_deltas)

        return wake_deltas

    def contribute_to_wake_deltas(
        self, algo, mdata, fdata, states_source_turbine, wake_deltas
    ):
        """
        Modifies wake deltas by contributions from the
        specified wake source turbines.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        mdata : foxes.core.Data
            The model data
        fdata : foxes.core.Data
            The farm data
        states_source_turbine : numpy.ndarray of int
            For each state, one turbine index corresponding
            to the wake causing turbine. Shape: (n_states,)
        wake_deltas : Any
            The wake deltas object created by the
            `new_wake_deltas` function

        """
        # calc coordinates to rotor centres:
        wcoos = self.wake_frame.get_wake_coos(
            algo, mdata, fdata, states_source_turbine, fdata[FV.TXYH]
        )

        # get x coordinates:
        x = wcoos[:, :, 0]
        del wcoos

        # evaluate grid rotor:
        n_states = fdata.n_states
        n_turbines = fdata.n_turbines
        n_rpoints = self.grotor.n_rotor_points()
        n_points = n_turbines * n_rpoints
        points = self.grotor.get_rotor_points(algo, mdata, fdata).reshape(
            n_states, n_points, 3
        )
        wcoos = self.wake_frame.get_wake_coos(
            algo, mdata, fdata, states_source_turbine, points
        )
        yz = wcoos.reshape(n_states, n_turbines, n_rpoints, 3)[:, :, :, 1:3]
        del points, wcoos

        # evaluate wake models:
        for w in self.wake_models:

            wdeltas, sp_sel = w.calc_wakes_spsel_x_yz(
                algo, mdata, fdata, states_source_turbine, x, yz
            )

            wsps = np.zeros((n_states, n_turbines, n_rpoints), dtype=bool)
            wsps[:] = sp_sel[:, :, None]
            wsps = wsps.reshape(n_states, n_points)

            for v, wdel in wdeltas.items():

                d = np.zeros((n_states, n_turbines, n_rpoints), dtype=FC.DTYPE)
                d[sp_sel] = wdel
                d = d.reshape(n_states, n_points)[wsps]

                try:
                    superp = w.superp[v]
                except KeyError:
                    raise KeyError(
                        f"Model '{self.name}': Missing wake superposition entry for variable '{v}' in wake model '{w.name}', found {sorted(list(w.superp.keys()))}"
                    )

                wake_deltas[v] = superp.calc_wakes_plus_wake(
                    algo,
                    mdata,
                    fdata,
                    states_source_turbine,
                    wsps,
                    v,
                    wake_deltas[v],
                    d,
                )

    def evaluate_results(
        self, algo, mdata, fdata, wake_deltas, states_turbine, update_amb_res=False
    ):
        """
        Updates the farm data according to the wake
        deltas.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        mdata : foxes.core.Data
            The model data
        fdata : foxes.core.Data
            The farm data
            Modified in-place by this function
        wake_deltas : Any
            The wake deltas object, created by the
            `new_wake_deltas` function and filled
            by `contribute_to_wake_deltas`
        states_turbine : numpy.ndarray of int
            For each state, the index of one turbine
            for which to evaluate the wake deltas.
            Shape: (n_states,)
        update_amb_res : bool
            Flag for updating ambient results

        """
        amb_res = self.get_data(FV.AMB_RPOINT_RESULTS, mdata)
        rpoints = self.get_data(FV.RPOINTS, mdata)
        rweights = self.get_data(FV.RWEIGHTS, mdata)
        wweights = self.grotor.rotor_point_weights()
        n_wpoints = self.grotor.n_rotor_points()
        n_states, n_turbines, n_rpoints, __ = rpoints.shape
        st_sel = (np.arange(n_states), states_turbine)

        uv = None
        if (FV.WS in amb_res and FV.WD not in amb_res) or (
            FV.WS not in amb_res and FV.WD in amb_res
        ):
            raise KeyError(
                f"Model '{self.name}': Missing one of the variables '{FV.WS}', '{FV.WD}' in ambient rotor results: {list(amb_res.keys())}"
            )

        elif FV.WD in amb_res and np.any(
            np.min(amb_res[FV.WD], axis=2) != np.max(amb_res[FV.WD], axis=2)
        ):

            wd = amb_res[FV.WD].reshape(n_states, n_turbines, n_rpoints)[st_sel]
            ws = amb_res[FV.WS].reshape(n_states, n_turbines, n_rpoints)[st_sel]
            uv = wd2uv(wd, ws, axis=-1)
            uv = np.einsum("spd,p->sd", uv, rweights)
            del ws, wd

        wres = {}
        for v, ares in amb_res.items():
            if v == FV.WS and uv is not None:
                wres[v] = np.linalg.norm(uv, axis=-1)
            elif v == FV.WD and uv is not None:
                wres[v] = uv2wd(uv, axis=-1)
            else:
                wres[v] = ares.reshape(n_states, n_turbines, n_rpoints)[st_sel]
                wres[v] = np.einsum("sp,p->s", wres[v], rweights)
            wres[v] = wres[v][:, None]
        del amb_res, uv

        wdel = {}
        for v, d in wake_deltas.items():
            wdel[v] = d.reshape(n_states, n_turbines, n_wpoints)[st_sel]
        for w in self.wake_models:
            w.finalize_wake_deltas(algo, mdata, fdata, wres, wdel)
        for v in wdel.keys():
            wdel[v] = np.einsum("sp,p->s", wdel[v], wweights)[:, None]

        for v in wres.keys():
            if v in wake_deltas:
                wres[v] += wdel[v]
                if update_amb_res:
                    mdata[FV.AMB_RPOINT_RESULTS][v][st_sel] = wres[v]
            wres[v] = wres[v][:, None]

        self.rotor_model.eval_rpoint_results(
            algo, mdata, fdata, wres, np.array([1.0]), states_turbine=states_turbine
        )

    def finalize(self, algo, verbosity=0):
        """
        Finalizes the model.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        verbosity : int
            The verbosity level

        """
        if self.rotor_model.initialized:
            self.rotor_model.finalize(algo, verbosity)
        if self.grotor is not None and self.grotor.initialized:
            self.grotor.finalize(algo, verbosity)
        super().finalize(algo, verbosity)
