"""
==============================================================
                      gcm_toolkit Main Class
==============================================================
 This class is the user interace of the gcm_toolkit
 functionalities. The goal is to have a clean and easy to use
 environment for new users of GCMs while allowing direct
 access to the data for more experienced users.
==============================================================
"""
import xarray

from .core import writer as wrt
from .utils import gcm_plotting as gcmplt
from .utils import manipulations as mani
from .utils import read_and_write as raw
from .gcm_dataset_collection import GCMDatasetCollection
from .core.units import ALLOWED_PUNITS, ALLOWED_TIMEUNITS
from .utils.passport import is_the_data_basic
from .utils.interface import PrtInterface


class GCMT:
    """
    The main gcm_toolkit class with which the user can interact.

    Attributes
    ----------
    models : dict
        Dictionary containing all of the 3D GCM models that are stored in the
        memory, with their respective tags.

    Methods
    -------
    read_raw(self, gcm, data_path, tag=None):
        Read in the raw data output from a GCM.
    read_reduced(self, data_path, tag=None):
        Read in the previously reduced GCM
    """

    def __init__(self, p_unit="bar", time_unit="day", write="on"):
        """
        Constructor for the gcm_toolkit class.

        Parameters
        ----------
        p_unit: str, optional
            Set the unit that is used internally for pressure related things
        time_unit: str, optional
            Set the unit that is used internally for time related things
        write: str, optional
            Set the level of logging information. Check core.writer for more infos
        """

        # Initialize empty dictionary to store all GCM models
        self._models = GCMDatasetCollection()

        # check units
        if p_unit not in ALLOWED_PUNITS:
            raise ValueError(
                f"Please use a pressure unit from {ALLOWED_PUNITS}"
            )
        self.p_unit = p_unit

        if time_unit not in ALLOWED_TIMEUNITS:
            raise ValueError(
                f"Please use a time unit from {ALLOWED_TIMEUNITS}"
            )
        self.time_unit = time_unit

        # initialize writing function (to file or to console)
        wrt.writer_setup(write)

        # print welcome message and information
        wrt.write_hline()
        hello = "Welcome to gcm_toolkit"
        wrt.write_message(
            hello,
            color="WARN",
            spacing=(wrt.Writer.line_length - len(hello)) // 2,
        )
        wrt.write_hline()
        wrt.write_status("STAT", "Set up gcm_toolkit")
        wrt.write_status("INFO", "pressure units: " + self.p_unit)
        wrt.write_status("INFO", "time units: " + self.time_unit)

    # ==============================================================================================
    #   Data handling
    # ==============================================================================================

    @property
    def models(self):
        """
        Shortcut to return all GCMs in memory.

        Returns
        -------
        selected_models : GCMDatasetCollection or xarray Dataset
            All models in self._models
        """
        return self.get_models()

    def __iter__(self):
        """
        Returns an iterator of GCMT
        """
        return iter(self.get_models(always_dict=True).items())

    def __getitem__(self, tag):
        """
        Access models of this dataset

        Parameters:
        -----------
        tag: str
            Tag of the model that should be returned

        Returns
        -------
        selected_model : xarray Dataset
        """
        if tag not in self._models.keys():
            raise ValueError(
                "The provided tag does not exist in the collection"
            )
        if isinstance(tag, str):
            return self.get_one_model(tag=tag, raise_error=True)

    def __setitem__(self, tag, dsi):
        """
        Add or replaces a dataset

        Parameters
        ----------
        tag: str
           Tag at which the model should be stored
        dsi: xarray Dataset
           Dataset to add
        """
        self._replace_model(tag, dsi)

    def __len__(self) -> int:
        return len(self._models)

    def __bool__(self) -> bool:
        return bool(self._models)

    def get(self, tag, default=None):
        """
        Pythonic get, will return default if tag is not available.

        Parameters
        ----------
        tag: str
            Name of the model that should be returned
        default: Any
            Default value in case that the tag is not available in the dataset

        Returns
        -------
        dsi or default, depending on whether dsi is available
        """
        if dsi := self.get_one_model(tag, raise_error=False):
            return dsi
        return default

    def get_one_model(self, tag=None, raise_error=True):
        """
        Like get_models, but raises an error or return None,
        if more than one model is selected.

        Parameters
        ----------
        tag: str, optional
            Name of the model that should be returned
        raise_error: bool, optional
            If true, function will raise error, else will return
            None if not one model is selected

        Returns
        -------
        dsi: xarray Dataset
            Selected model
        """
        return self._models.get_one_model(tag, raise_error)

    def get_models(self, tag=None, always_dict=False):
        """
        Function return all GCMs in memory. If a tag is given, only return this
        one.

        Parameters
        ----------
        tag : str
            Name of the model that should be returned.
        always_dict: bool
            Force result to be a dictionary (if tag is None)

        Returns
        -------
        selected_models : GCMDatasetCollection or xarray Dataset
            All models in self._models, or only the one with the right tag.
            Will definetly be GCMDatasetCollection if always_dict=True
        """
        return self._models.get_models(tag, always_dict)

    def _replace_model(self, tag, dsi):
        """
        Add or replaces a dataset. Do some checks beforehand.

        Parameters
        ----------
        tag: str
           Tag at which the model should be stored
        dsi: xarray Dataset
           Dataset to add
        """
        if not isinstance(tag, str):
            raise ValueError("The provided tag needs to be a string.")
        if not isinstance(dsi, xarray.Dataset):
            raise ValueError(
                "The provided input dataset needs to be a xarray Dataset."
            )
        if not is_the_data_basic(dsi):
            raise ValueError("The provided input dataset is not compatible.")

        if dsi.attrs.get("p_unit") != self.p_unit:
            raise NotImplementedError(
                "Unit conversion is needed and not yet implemented"
            )
        if dsi.attrs.get("time_unit") != self.time_unit:
            raise NotImplementedError(
                "Unit conversion is needed and not yet implemented"
            )

        dsi.attrs.update({"tag": tag})
        self._models[tag] = dsi

    # ==============================================================================================
    #   Data manipulation
    # ==============================================================================================
    def add_horizontal_average(
        self,
        var_key,
        var_key_out=None,
        part="global",
        area_key="area_c",
        tag=None,
    ):
        """
        Calculate horizontal averaged quantities. Horizontal averages
        are calculated as area-weighted quantities q, so that:

            bar{q} = int{q dA}/int{dA}

        Parameters
        ----------
        var_key: str, xarray.DataArray
            The key or array of the variable quantity that should be averaged.
            If str, it will try to look up the key in the dataset.
            If DataArray, it will use this one instead.
        var_key_out: str, optional
            variable name used to store the outcome.
            If not provided, this script will just
            return the averages and not change the dataset inplace.
        part: dict or str, optional
            'global': global average
            'night': only nightside (defined around +-180,0)
            'day': only dayside (defined around 0,0)
            'morning': morning terminator (average around lon=[-100,-80])
            'evening': evening terminator (average around lon=[80,100])
            Alternatively you may specify a dict in the following way
            (example for morn. term.):
            part = {'lon': [-100,-80], 'lat':[-90,90], 'inv':False}
            The 'lon', 'lat' specify regions of lon and lat that should be used,
            whereas 'inv' (optional, default False) gives the option to invert
            the lon and lat regions (e.g., exclude instead of include for average)
        area_key: str, optional
            Variable key in the dataset for the area of grid cells
        tag : str, optional
            The tag of the dataset that should be used.
            If no tag is provided, and multiple datasets are available,
            an error is raised.

        Returns
        -------
        avg: xarray.DataArray
            Averaged data

        """
        dsi = self.get_one_model(tag)
        return mani.m_add_horizontal_average(
            dsi, var_key, var_key_out=var_key_out, part=part, area_key=area_key
        )

    def add_rcb(
        self,
        tol=0.01,
        var_key_out=None,
        area_key="area_c",
        temp_key="T",
        part="global",
        tag=None,
    ):
        """
        Calculate the radiative convective boundary (rcb) by searching
        (from the bottom upwards) for the first occurance of a deviation
        from an adiabatic temperature profile.

        Operates on and calculates the horizontal average.

        Parameters
        ----------
        tol: float
            tolerance for the relative deviation from adiabat
        var_key_out: str, optional
            variable name used to store the outcome.
            If not provided, this script will just
            return the averages and not change the dataset inplace.
        part: dict or str, optional
            'global': global average
            'night': only nightside (defined around +-180,0)
            'day': only dayside (defined around 0,0)
            'morning': morning terminator (average around lon=[-100,-80])
            'evening': evening terminator (average around lon=[80,100])
            Alternatively you may specify a dict in the following way
            (example for morn. term.):
            part = {'lon': [-100,-80], 'lat':[-90,90], 'inv':False}
            The 'lon', 'lat' specify regions of lon and lat that should be used,
            whereas 'inv' (optional, default False) gives the option to invert
            the lon and lat regions (e.g., exclude instead of include for average)
        area_key: str, optional
            Variable key in the dataset for the area of grid cells
        temp_key: str, optional
            The key to look up the temperature (needed for density calculation)
        tag : str, optional
            The tag of the dataset that should be used.
            If no tag is provided, and multiple datasets are available,
            an error is raised.

        Returns
        -------
        rcb : xarray.DataArray
            A dataArray with reduced dimensionality,
            containing the pressure of the rcb location.
        """
        dsi = self.get_one_model(tag)
        return mani.m_add_rcb(
            dsi,
            tol=tol,
            var_key_out=var_key_out,
            part=part,
            area_key=area_key,
            temp_key=temp_key,
        )

    def add_theta(self, var_key_out=None, temp_key="T", tag=None):
        """
        Convert temperature to potential temperature with respect to model boundary.

        Parameters
        ----------
        var_key_out: str, optional
            variable name used to store the outcome.
            If not provided, this script will just
            return theta and not change the dataset inplace.
        temp_key: str, optional
            The key to look up the temperature
        tag : str, optional
            The tag of the dataset that should be used.
            If no tag is provided, and multiple datasets are available,
            an error is raised.

        Returns
        -------
        theta : xarray.DataArray
            A dataArray with reduced dimensionality,
            containing the potential temperature
        """
        dsi = self.get_one_model(tag)
        return mani.m_add_theta(
            dsi, var_key_out=var_key_out, temp_key=temp_key
        )

    def add_total_energy(
        self,
        var_key_out=None,
        area_key="area_c",
        temp_key="T",
        return_all=False,
        tag=None,
    ):
        """
        Calculate the total Energy of the GCM. See e.g.,
        https://ui.adsabs.harvard.edu/abs/2014Icar..229..355P, Eq. 16

        Parameters
        ----------
        var_key_out: str, optional
            variable name used to store the outcome. If not provided,
            this script will just return the averages and not change
            the dataset inplace.
        area_key: str, optional
            Variable key in the dataset for the area of grid cells
        temp_key: str, optional
            The key to look up the temperature
        return_all: bool, optional
            Also return the partial energies
        tag : str, optional
            The tag of the dataset that should be used.
            If no tag is provided,
            and multiple datasets are available, an error is raised.

        Returns
        -------
        energy : xarray.DataArray
            A dataArray with reduced dimensionality, containing the total energy.
        therm_energy : xarray.DataArray, optional
            A dataArray with reduced dimensionality, containing the thermal energy.
        pot_energy : xarray.DataArray, optional
            A dataArray with reduced dimensionality, containing the potential energy.
        kin_energy : xarray.DataArray, optional
            A dataArray with reduced dimensionality, containing the kinetic energy.
        """

        dsi = self.get_one_model(tag)
        return mani.m_add_total_energy(
            dsi,
            var_key_out=var_key_out,
            area_key=area_key,
            temp_key=temp_key,
            return_all=return_all,
        )

    def add_total_momentum(
        self, var_key_out=None, area_key="area_c", temp_key="T", tag=None
    ):
        """
        Calculate the total angular momentum of the GCM. See e.g.,
        https://ui.adsabs.harvard.edu/abs/2014Icar..229..355P, Eq. 17

        Parameters
        ----------
        var_key_out: str, optional
            variable name used to store the outcome.
            If not provided, this script will just
            return the averages and not change the dataset inplace.
        area_key: str, optional
            Variable key in the dataset for the area of grid cells
        temp_key: str, optional
            The key to look up the temperature (needed for density calculation)
        tag : str, optional
            The tag of the dataset that should be used.
            If no tag is provided,
            and multiple datasets are available, an error is raised.

        Returns
        -------
        momentum : xarray.DataArray
            A dataArray with reduced dimensionality, containing the total momentum.
        """
        dsi = self.get_one_model(tag)
        return mani.m_add_total_momentum(
            dsi, var_key_out=var_key_out, area_key=area_key, temp_key=temp_key
        )

    def add_meridional_overturning(
        self, v_data="V", var_key_out=None, tag=None
    ):
        """
        Calculate meridional overturning streamfunction. This quantity psi is
        computed by integrating the zonal-mean meridional velocity \bar V along
        pressure, and weighting with 2*pi*R_p / g times the cosine of latitude,
        where R_p is the planetary radius and g is the surface gravity:
        bar{psi} = 2 pi R_p / g cos{lat} int{bar{V} dp'}
        (see e.g. Carone et al. (2018), Eq. 7)
    
        Parameters
        ----------
        v_data: str
            The key that holds the data to meridional velocity
        var_key_out: str, optional
            variable name used to store the outcome.
            If not provided, this script will just
            return the overturning circulation
            and not change the dataset inplace.
        tag : str, optional
            The tag of the dataset that should be used.
            If no tag is provided,
            and multiple datasets are available, an error is raised.

        Returns
        -------
        psi: xarray.DataArray
            Horizontal overturning
        """
        dsi = self.get_one_model(tag)
        return mani.m_add_meridional_overturning(
            dsi, v_data=v_data, var_key_out=var_key_out
        )

    def extend_upward(self, p_low, method=None, n_p=20, T_therm=None,
                      p_therm_high=None, tag=None):
        """
        Extend temperatures of the tagged dataset upward to the given pressure value.
        Multiple methods can be used to extend the data upwards:
            - 'isothermal' copies the temperature at lowest pressure and adds layers
              of the same temperature on top
            - 'thermosphere' imposes a simple parametrized thermosphere with a hot
              day-side and gradual decrease to an isothermal night-side
              (cfr. Baeyens+2022, Section 4.2)
        Parameters
        ----------
        p_low: float or list
            If a float is given, extend to the given low-pressure limit using an
            arbitrarily chosen pressure-spacing. If a list is given, use these as
            pressures to compute temperature extension.
        method: string
            Method by which the temperature is to be extended.
        n_p: int, optional
            Number of pressure layers that should be added (only used in case p_low
            is a single number).
        T_therm: float, optional
            Maximal day-side temperature for the upper thermosphere (required for
            the 'thermosphere' method.)
        p_therm_high: float, optional
            The high-pressure limit at which thermospheric heating should start, in
            native pressure units (required for the 'thermosphere' method).
        tag : str, optional
            The tag of the dataset that should be used. If no tag is provided,
            and multiple datasets are available, an error is raised.
        Returns
        -------
        temp_array_ext: DataArray
            The extended temperature array.
        """
        dsi = self.get_one_model(tag)
        return mani.m_extend_upward(dsi.T, p_low, method=method, n_p=n_p,
                                    T_therm=T_therm, p_therm_high=p_therm_high)

    # ======================================================
    #   Reading and writing functions
    # ======================================================
    def read_raw(
        self,
        gcm,
        data_path,
        iters="last",
        load_existing=False,
        tag=None,
        **kwargs,
    ):
        """
        General read in function for GCM data

        Parameters
        ----------
        gcm : str
            Type of GCM, must be 'MITgcm'.
        data_path : str
            Folder path to the standard output of the GCM.
        iters : list, str
            The iteration (time step) of the input files to be read.
            If None, no data will be read.
            If 'last' (default), only the last iteration will be read.
            If 'all', all iterations will be read.
        load_existing: bool
            Set to false if you want to overwrite already loaded data
            Set to true if you want to increment already loaded data
        tag : str
            Tag to reference the simulation in the collection of models.
        kwargs: dict
            Additional options passed down to read functions
        """
        return raw.m_read_raw(
            self,
            gcm,
            data_path,
            iters=iters,
            load_existing=load_existing,
            tag=tag,
            **kwargs,
        )

    def read_reduced(
        self, data_path, tag=None, time_unit_in=None, p_unit_in=None
    ):
        """
        Read in function for GCM data that has been reduced a
        nd saved according to the gcm_toolkit GCMDataset format.

        Parameters
        ----------
        data_path : str
            Folder path to the reduced (gcm_toolkit) data.
        time_unit_in: str, None
            units of time dimension in input dataset.
            If None, try to read from nc file (ds.attr.time_unit)
        p_unit_in: str, None
            units of pressure dimensions in input dataset,
            If None, try to read from nc file (ds.attr.p_unit)
        tag : str
            Tag to reference the simulation in the collection of models.
        """
        return raw.m_read_reduced(
            self,
            data_path,
            tag=tag,
            time_unit_in=time_unit_in,
            p_unit_in=p_unit_in,
        )

    def save(self, direct, method="nc", update_along_time=False, tag=None):
        """
        Save function to store current member variables.

        Parameters
        ----------
        direct : str
            directory at which the gcm_toolkit datasets should be stored.
        method : str, optional
            Datasets can be stored as '.zarr' or '.nc'.
            Decide which type you prefer. Defaults to '.nc'.
        update_along_time: bool, optional
            Decide if you want to update already saved datasets
            along the timedimension. This only works with method='zarr'.
        tag: str, optional
            tag of the model that should be loaded.
            Will save all available models by default.

        Returns
        -------
        NoneType
            None
        """
        return raw.m_save(
            self,
            direct,
            method=method,
            update_along_time=update_along_time,
            tag=tag,
        )

    def load(self, direct, method="nc", tag=None):
        """
        Load function to load stored member variables.

        Parameters
        ----------
        direct : str
            directory at which the gcm_toolkit datasets are stored
        method : str, optional
            Should be the same method with which you stored the data
        tag: str, optional
            tag of the model that should be loaded.
            Will load all available models by default.
        """
        return raw.m_load(self, direct, method=method, tag=tag)

    # =============================================================
    #   Plotting Functions
    # =============================================================

    def isobaric_slice(self, var_key, pres, tag=None, **kwargs):
        """
        Plot an isobaric slice of the given quantity at the given pressure
        level. The user can specify the DataSet to be plotted by providing the
        corresponding tag.

        Parameters
        ----------
        var_key : str
            The key of the variable quantity that should be plotted.
        pres : float
            Pressure level for the isobaric slice to be plotted, expressed in
            the units specified in the dataset attributes (e.g., init of GCMT object).
        tag : str, optional
            The tag of the dataset that should be plotted. If no tag is provided
            and multiple datasets are available, an error is raised.
        """
        # select the appropriate dataset
        dsi = self.get_one_model(tag)
        return gcmplt.isobaric_slice(dsi, var_key, pres, **kwargs)

    def time_evol(self, var_key, tag=None, **kwargs):
        """
        Function that plots the time evolution of a quantity
        in a 1D line collection plot, where the colorscale can be
        related to the time evolution.
        Note: var_key needs to contain data that is 2D in time and pressure.

        Parameters
        ----------
        dsi : DataSet
            A gcm_toolkit-compatible dataset of a 3D climate simulation.
        var_key : str
            The key of the variable quantity that should be plotted.
        tag : str, optional
            The tag of the dataset that should be plotted. If no tag is provided
            and multiple datasets are available, an error is raised.
        """
        dsi = self.get_one_model(tag)
        return gcmplt.time_evol(dsi, var_key, **kwargs)

    def zonal_mean(self, var_key, tag=None, **kwargs):
        """
        Plot a zonal mean average of a quantity for the given dataset.

        Parameters
        ----------
        dsi : DataSet
            A gcm_toolkit-compatible dataset of a 3D climate simulation.
        var_key : str
            The key of the variable quantity that should be plotted.
        tag : str, optional
            The tag of the dataset that should be plotted. If no tag is provided
            and multiple datasets are available, an error is raised.
        """
        # select the appropriate dataset
        dsi = self.get_one_model(tag)
        return gcmplt.zonal_mean(dsi, var_key, **kwargs)

    # ==============================================================================================
    #   Interfaces
    # ==============================================================================================

    def get_prt_interface(self, prt):
        """
        Constructs the interface to petitRADTRANS

        Parameters
        ----------
        prt: petitRADTRANS.Radtrans
            A fully initialized radtrans object to be used for the calculations

        Returns
        -------
        Interface: pRTInterface
            The interface object that is used to create phasecurves/spectra/etc
        """
        return PrtInterface(self, prt)
