from abc import abstractmethod

from foxes.core.data_calc_model import DataCalcModel
import foxes.variables as FV

class FarmDataModel(DataCalcModel):
    """
    Abstract base class for models that modify
    farm data.

    Parameters
    ----------
    pre_rotor : bool
        Flag for running this model before
        running the rotor model.
    
    Attributes
    ----------
    pre_rotor : bool
        Flag for running this model before
        running the rotor model.

    """

    def __init__(self, pre_rotor=False):
        super().__init__()
        self.pre_rotor = pre_rotor

    @abstractmethod
    def output_farm_vars(self, algo):
        """
        The variables which are being modified by the model.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        
        Returns
        -------
        output_vars : list of str
            The output variable names

        """
        return []

    @abstractmethod
    def calculate(self, algo, mdata, fdata):
        """"
        The main model calculation.

        This function is executed on a single chunk of data,
        all computations should be based on numpy arrays.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        mdata : foxes.core.Data
            The model data
        fdata : foxes.core.Data
            The farm data
        
        Returns
        -------
        results : dict
            The resulting data, keys: output variable str.
            Values: numpy.ndarray with shape (n_states, n_turbines)

        """
        pass

    def run_calculation(
            self, 
            algo, 
            *data, 
            out_vars,
            **calc_pars
        ):
        """
        Starts the model calculation in parallel, via
        xarray's `apply_ufunc`.

        Typically this function is called by algorithms.

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        *data : tuple of xarray.Dataset
            The input data
        out_vars: list of str
            The calculation output variables
        **calc_pars : dict, optional
            Additional arguments for the `calculate` function
        
        Returns
        -------
        results : xarray.Dataset
            The calculation results

        """
        return super().run_calculation(algo, *data, out_vars=out_vars,
                                        loop_dims=[FV.STATE], 
                                        out_core_vars=[FV.TURBINE, FV.VARS],
                                        **calc_pars)

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

        Parameters
        ----------
        algo : foxes.core.Algorithm
            The calculation algorithm
        results : xarray.Dataset
            The calculation results
        clear_mem : bool
            Flag for deleting model data and
            resetting initialization flag
        verbosity : int
            The verbosity level

        """
        super().finalize(algo, clear_mem=clear_mem, verbosity=verbosity)
