"""
This module contains the main class that describes an isotherm through discrete points.
"""

import textwrap
import typing as t

import numpy
import pandas

from pygaps import logger
from pygaps.core.baseisotherm import BaseIsotherm
from pygaps.units.converter_mode import c_loading
from pygaps.units.converter_mode import c_material
from pygaps.units.converter_mode import c_pressure
from pygaps.utilities.exceptions import CalculationError
from pygaps.utilities.exceptions import ParameterError
from pygaps.utilities.isotherm_interpolator import IsothermInterpolator

class PointIsotherm(BaseIsotherm):
    """
    Class which contains the points from an adsorption isotherm.

    This class is designed to be a complete description of a discrete isotherm.
    It extends the BaseIsotherm class, which contains all the description of the
    isotherm parameters, but also holds the datapoints recorded during an
    experiment or simulation.

    The minimum arguments required to instantiate the class, besides those
    required for the parent BaseIsotherm, is the actual data, specified either
    as `pressure` + `loading` arrays or as `isotherm_data` (a
    `pandas.DataFrame`) + keys for the columns of the dataframe which have the
    loading and the pressure data.

    Parameters
    ----------
    pressure : list
        Create an isotherm directly from an array. Values for pressure.
        If the ``isotherm_data`` dataframe is specified, these values are ignored.
    loading : list
        Create an isotherm directly from an array. Values for loading.
        If the ``isotherm_data`` dataframe is specified, these values are ignored.
    isotherm_data : `pandas.DataFrame`
        Pure-component adsorption isotherm data.
    pressure_key : str
        The title of the pressure data in the DataFrame provided.
    loading_key : str
        The title of the loading data in the DataFrame provided.
    branch : ['guess', ads', 'des', iterable], optional
        The branch of the isotherm. The code will automatically attempt to
        guess if there's an adsorption and desorption branch.
        The user can instead tell the framework that all points are
        part of an adsorption ('ads') or desorption ('des') curve.
        Alternatively, an iterable can be passed which contains
        detailed info for each data point if adsorption points ('False')
        or desorption points ('True'). eg: [False, False, True, True...]
        or as a column of the isotherm_data.
    material : str
        Name of the material on which the isotherm is measured.
    adsorbate : str
        Isotherm adsorbate.
    temperature : float
        Isotherm temperature.

    Other Parameters
    ----------------
    pressure_mode : str, optional
        The pressure mode, either 'absolute' pressure or 'relative'
        ('relative%') in the form of p/p0.
    pressure_unit : str, optional
        Unit of pressure, if applicable.
    loading_basis : str, optional
        Whether the adsorbed amount is in terms of either 'volume_gas'
        'volume_liquid', 'molar', 'mass', or a fractional/percent basis.
    loading_unit : str, optional
        Unit in which the loading basis is expressed.
    material_basis : str, optional
        Whether the underlying material is in terms of 'per volume'
        'per molar amount' or 'per mass' of material.
    material_unit : str, optional
        Unit in which the material basis is expressed.

    Notes
    -----
    This class assumes that the datapoints do not contain noise.
    Detection of adsorption/desorption branches will not work if
    data is noisy.

    """

    _reserved_params = BaseIsotherm._reserved_params + [
        'data_raw',
        'l_interpolator',
        'p_interpolator',
        'loading_key',
        'pressure_key',
        'other_keys',
    ]

    ##########################################################
    #   Instantiation and classmethods

    def __init__(
        self,
        pressure: t.List[float] = None,
        loading: t.List[float] = None,
        isotherm_data: pandas.DataFrame = None,
        pressure_key: str = None,
        loading_key: str = None,
        branch: t.Union[str, t.List[bool]] = 'guess',
        **other_properties
    ):
        """
        Instantiation is done by passing the discrete data as a pandas
        DataFrame, the column keys as string as well as the parameters
        required by parent class.
        """
        # Run base class constructor
        super().__init__(**other_properties)

        # Checks
        if isotherm_data is not None:
            if None in [pressure_key, loading_key]:
                raise ParameterError(
                    "Pass loading_key and pressure_key, the names of the loading and"
                    " pressure columns in the DataFrame, to the constructor."
                )

            # Save column names
            # Name of column in the dataframe that contains adsorbed amount.
            self.loading_key = loading_key

            # Name of column in the dataframe that contains pressure.
            self.pressure_key = pressure_key

            # Pandas DataFrame that stores the data.
            columns = [self.pressure_key, self.loading_key]
            if not all(a in isotherm_data.columns for a in columns):
                raise ParameterError(
                    "Could not find columns "
                    f"({[a for a in columns if a not in isotherm_data.columns]})"
                    " in the adsorption DataFrame."
                )
            if 'branch' not in isotherm_data.columns:
                columns.append('branch')
            other_keys = [c for c in isotherm_data.columns if c not in columns]
            columns = columns + sorted(other_keys)
            self.data_raw = isotherm_data.reindex(columns=columns)

        elif pressure is not None or loading is not None:
            if pressure is None or loading is None:
                raise ParameterError(
                    "If you've chosen to pass loading and pressure directly as"
                    " arrays, make sure both are specified!"
                )
            if len(pressure) != len(loading):
                raise ParameterError("Pressure and loading arrays are not equal!")

            # Standard column names
            self.pressure_key = 'pressure'
            self.loading_key = 'loading'

            # DataFrame creation
            self.data_raw = pandas.DataFrame({
                self.pressure_key: pressure,
                self.loading_key: loading
            })
        else:
            raise ParameterError(
                "Pass either the isotherm data in a pandas.DataFrame as ``isotherm_data``"
                " or directly ``pressure`` and ``loading`` as arrays."
            )

        # Deal with the isotherm branches
        if isotherm_data is not None and 'branch' in isotherm_data.columns:
            pass
        elif isinstance(branch, str):
            if branch == 'guess':
                # Split the data in adsorption/desorption
                self.data_raw['branch'] = self._splitdata(self.data_raw, self.pressure_key)
            elif branch == 'ads':
                self.data_raw['branch'] = 0
            elif branch == 'des':
                self.data_raw['branch'] = 1
            else:
                raise ParameterError(
                    "Isotherm branch parameter must be 'guess ,'ads' or 'des'"
                    " or an array of booleans."
                )
        else:
            try:
                self.data_raw['branch'] = branch
            except Exception as e_info:
                raise ParameterError(e_info)

        # The internal interpolator for loading given pressure.
        self.l_interpolator = None

        # The internal interpolator for pressure given loading.
        self.p_interpolator = None

    @classmethod
    def from_isotherm(
        cls,
        isotherm: BaseIsotherm,
        pressure: t.List[float] = None,
        loading: t.List[float] = None,
        isotherm_data: pandas.DataFrame = None,
        pressure_key: str = None,
        loading_key: str = None,
    ):
        """
        Construct a point isotherm using a parent isotherm as the template for
        all the parameters.

        Parameters
        ----------
        isotherm : Isotherm
            An instance of the Isotherm parent class.
        pressure : list
            Create an isotherm directly from an array. Values for pressure.
            If the ``isotherm_data`` dataframe is specified, these values are ignored.
        loading : list
            Create an isotherm directly from an array. Values for loading.
            If the ``isotherm_data`` dataframe is specified, these values are ignored.
        isotherm_data : pandas.DataFrame
            Pure-component adsorption isotherm data.
        loading_key : str
            Column of the pandas DataFrame where the loading is stored.
        pressure_key : str
            Column of the pandas DataFrame where the pressure is stored.
        """
        # get isotherm parameters as a dictionary
        iso_params = isotherm.to_dict()
        # add pointisotherm values to dict
        iso_params['pressure'] = pressure
        iso_params['loading'] = loading
        iso_params['isotherm_data'] = isotherm_data
        iso_params['pressure_key'] = pressure_key
        iso_params['loading_key'] = loading_key

        return cls(**iso_params)

    @classmethod
    def from_modelisotherm(
        cls,
        modelisotherm,
        pressure_points: t.List[float] = None,
        loading_points: t.List[float] = None,
    ):
        """
        Construct a PointIsotherm from a ModelIsothem class.

        This class method allows for the model to be converted into
        a list of points calculated by using the model in the isotherm.

        Parameters
        ----------
        modelisotherm : ModelIsotherm
            The isotherm containing the model.
        pressure_points : None or List or PointIsotherm
            How the pressure points should be chosen for the resulting PointIsotherm.

            - If ``None``, the PointIsotherm returned has a fixed number of
              equidistant points
            - If an array, the PointIsotherm returned has points at each of the
              values of the array
            - If a PointIsotherm is passed, the values will be calculated at
              each of the pressure points in the passed isotherm. This is useful
              for comparing a model overlap with the real isotherm.
        """
        if pressure_points is not None and loading_points is not None:
            raise ParameterError("""Cannot specify both pressure and loading points.""")

        pressure = None
        loading = None
        if modelisotherm.model.calculates == "loading":
            # The user may request loading even if the model calculates pressure
            if loading_points is None:
                if pressure_points is None:
                    pressure = modelisotherm.pressure()
                elif isinstance(pressure_points, PointIsotherm):
                    pressure = pressure_points.pressure(branch=modelisotherm.branch)
                else:
                    pressure = pressure_points
                loading = modelisotherm.loading_at(pressure)
            else:
                loading = loading_points
                pressure = modelisotherm.pressure_at(loading_points)
        elif modelisotherm.model.calculates == "pressure":
            # The user may request pressure even if the model calculates loading
            if pressure_points is None:
                if loading_points is None:
                    loading = modelisotherm.loading()
                elif isinstance(loading_points, PointIsotherm):
                    loading = loading_points.loading(branch=modelisotherm.branch)
                else:
                    loading = loading_points
                pressure = modelisotherm.pressure_at(loading)
            else:
                pressure = pressure_points
                loading = modelisotherm.loading_at(pressure)

        return PointIsotherm(
            pressure=pressure,
            loading=loading,
            model_from=modelisotherm.model.name,
            **modelisotherm.to_dict()
        )

    ##########################################################
    #   Conversion functions

    def convert(
        self,
        pressure_mode: str = None,
        pressure_unit: str = None,
        loading_basis: str = None,
        loading_unit: str = None,
        material_basis: str = None,
        material_unit: str = None,
        verbose: bool = False,
    ):
        """
        Convenience function for permanently converting any isotherm
        mode/basis/units.

        Parameters
        ----------
        pressure_mode : {'absolute', 'relative', 'relative%'}
            The mode in which the isotherm should be converted.
        pressure_unit : str
            The unit into which the internal pressure should be converted to.
            Only makes sense if converting to absolute pressure.
        loading_basis : {'mass', 'molar', 'volume_gas', 'volume_liquid', 'percent', 'fraction'}
            The basis in which the isotherm should be converted.
        loading_unit : str
            The unit into which the internal loading should be converted to.
        material_basis : {'mass', 'molar', 'volume'}
            The basis in which the isotherm should be converted.
        material_unit : str
            The unit into which the material should be converted to.
        verbose : bool
            Print out steps taken.

        """
        if pressure_mode or pressure_unit:
            self.convert_pressure(
                mode_to=pressure_mode,
                unit_to=pressure_unit,
                verbose=verbose,
            )

        if material_basis or material_unit:
            self.convert_material(
                basis_to=material_basis,
                unit_to=material_unit,
                verbose=verbose,
            )

        if loading_basis or loading_unit:
            self.convert_loading(
                basis_to=loading_basis,
                unit_to=loading_unit,
                verbose=verbose,
            )

    def convert_pressure(
        self,
        mode_to: str = None,
        unit_to: str = None,
        verbose: bool = False,
    ):
        """
        Convert isotherm pressure from one unit to another
        and the pressure mode from absolute to relative.

        Only applicable in the case of isotherms taken below critical
        point of adsorbate.

        Parameters
        ----------
        mode_to : {'absolute', 'relative', 'relative%'}
            The mode in which the isotherm should be converted.
        unit_to : str
            The unit into which the internal pressure should be converted to.
            Only makes sense if converting to absolute pressure.
        verbose : bool
            Print out steps taken.

        """
        if not mode_to:
            mode_to = self.pressure_mode

        if mode_to == self.pressure_mode and unit_to == self.pressure_unit:
            if verbose:
                logger.info("Mode and units are the same, no changes made.")
            return

        self.data_raw[self.pressure_key] = c_pressure(
            self.data_raw[self.pressure_key],
            mode_from=self.pressure_mode,
            mode_to=mode_to,
            unit_from=self.pressure_unit,
            unit_to=unit_to,
            adsorbate=self.adsorbate,
            temp=self.temperature
        )

        if mode_to != self.pressure_mode:
            self.pressure_mode = mode_to
        if unit_to != self.pressure_unit and mode_to == 'absolute':
            self.pressure_unit = unit_to
        else:
            self.pressure_unit = None

        # Reset interpolators
        self.l_interpolator = None
        self.p_interpolator = None

        if verbose:
            logger.info(f"Changed pressure to mode '{mode_to}', unit '{unit_to}'.")

    def convert_loading(
        self,
        basis_to: str = None,
        unit_to: str = None,
        verbose: bool = False,
    ):
        """
        Convert isotherm loading from one unit to another
        and the basis of the isotherm loading to be
        either 'mass', 'molar' or 'percent'/'fraction'.

        Parameters
        ----------
        basis_to : {'mass', 'molar', 'volume_gas', 'volume_liquid', 'percent', 'fraction'}
            The basis in which the isotherm should be converted.
        unit_to : str
            The unit into which the internal loading should be converted to.
        verbose : bool
            Print out steps taken.

        """
        if not basis_to:
            basis_to = self.loading_basis

        if basis_to == self.loading_basis and unit_to == self.loading_unit:
            if verbose:
                logger.info("Basis and units are the same, no changes made.")
            return

        if self.loading_basis in ['percent', 'fraction']:
            # TODO this is
            if basis_to == self.loading_basis and unit_to != self.loading_unit:
                if verbose:
                    logger.info("There are no loading units in this mode.")
                return

        self.data_raw[self.loading_key] = c_loading(
            self.data_raw[self.loading_key],
            basis_from=self.loading_basis,
            basis_to=basis_to,
            unit_from=self.loading_unit,
            unit_to=unit_to,
            adsorbate=self.adsorbate,
            temp=self.temperature,
            basis_material=self.material_basis,
            unit_material=self.material_unit,
        )

        if basis_to != self.loading_basis:
            self.loading_basis = basis_to
        if basis_to in ['percent', 'fraction']:
            self.loading_unit = None
        else:
            self.loading_unit = unit_to

        # Reset interpolators
        self.l_interpolator = None
        self.p_interpolator = None

        if verbose:
            logger.info(f"Changed loading to basis '{basis_to}', unit '{unit_to}'.")

    def convert_material(
        self,
        basis_to: str = None,
        unit_to: str = None,
        verbose: bool = False,
    ):
        """
        Convert the material of the isotherm from one unit to another and the
        basis of the isotherm loading to be either 'per mass' or 'per volume' or
        'per mole' of material.

        Only applicable to materials that have been loaded in memory with a
        'density' or 'molar mass' property respectively.

        Parameters
        ----------
        basis : {'mass', 'molar', 'volume'}
            The basis in which the isotherm should be converted.
        unit_to : str
            The unit into which the material should be converted to.
        verbose : bool
            Print out steps taken.

        """
        if not basis_to:
            basis_to = self.material_basis

        if basis_to == self.material_basis and unit_to == self.material_unit:
            if verbose:
                logger.info("Basis and units are the same, no changes made.")
            return

        if (
            self.loading_basis in ['percent', 'fraction'] and basis_to == self.material_basis
            and unit_to != self.material_unit
        ):
            # We "virtually" change the unit without any conversion
            self.material_unit = unit_to
            if verbose:
                logger.info("There are no material units in this mode.")
            return

        self.data_raw[self.loading_key] = c_material(
            self.data_raw[self.loading_key],
            basis_from=self.material_basis,
            basis_to=basis_to,
            unit_from=self.material_unit,
            unit_to=unit_to,
            material=self.material
        )

        # A special case is when conversion is performed from
        # a "fractional" basis to another "fractional" basis.
        # Here, the loading must be simultaneously converted.
        # e.g.: wt% = g/g -> cm3/cm3 = vol%
        if self.loading_basis in ['percent', 'fraction']:
            if basis_to == 'volume':
                _basis_to = 'volume_liquid'
            else:
                _basis_to = basis_to
            if self.material_basis == 'volume':
                _basis_from = 'volume_liquid'
            else:
                _basis_from = self.material_basis
            self.data_raw[self.loading_key] = c_loading(
                self.data_raw[self.loading_key],
                basis_from=_basis_from,
                basis_to=_basis_to,
                unit_from=self.material_unit,
                unit_to=unit_to,
                adsorbate=self.adsorbate,
                temp=self.temperature,
            )
            if verbose:
                logger.info(f"Changed loading to basis '{basis_to}', unit '{unit_to}'.")

        if unit_to != self.material_unit:
            self.material_unit = unit_to
        if basis_to != self.material_basis:
            self.material_basis = basis_to

        # Reset interpolators
        self.l_interpolator = None
        self.p_interpolator = None

        if verbose:
            logger.info(f"Changed material to basis '{basis_to}', unit '{unit_to}'.")

    ###########################################################
    #   Info functions

    def print_info(self, **plot_iso_args):
        """
        Print a short summary of all the isotherm parameters and a graph.

        Parameters
        ----------
        show : bool, optional
            Specifies if the graph is shown automatically or not.

        Other Parameters
        ----------------
        plot_iso_args : dict
            options to be passed to pygaps.plot_iso()

        Returns
        -------
        axes : matplotlib.axes.Axes or numpy.ndarray of them

        """
        print(self)
        return self.plot(**plot_iso_args)

    def plot(self, **plot_iso_args):
        """
        Plot the isotherm using pygaps.plot_iso().

        Parameters
        ----------
        show : bool, optional
            Specifies if the graph is shown automatically or not.

        Other Parameters
        ----------------
        plot_iso_args : dict
            options to be passed to pygaps.plot_iso()

        Returns
        -------
        axes : matplotlib.axes.Axes or numpy.ndarray of them

        """
        plot_dict = dict(
            y2_data=self.other_keys[0] if self.other_keys else None,
            material_basis=self.material_basis,
            material_unit=self.material_unit,
            loading_basis=self.loading_basis,
            loading_unit=self.loading_unit,
            pressure_unit=self.pressure_unit,
            pressure_mode=self.pressure_mode,
        )
        plot_dict.update(plot_iso_args)

        from pygaps.graphing.isotherm_graphs import plot_iso
        return plot_iso(self, **plot_dict)

    ##########################################################
    #   Functions that return part of the isotherm data

    def data(self, branch: str = None) -> pandas.DataFrame:
        """
        Return underlying isotherm data.

        Parameters
        ----------
        branch : {None, 'ads', 'des'}
            The branch of the isotherm to return. If ``None``, returns entire
            dataset.

        Returns
        -------
        DataFrame
            The pandas DataFrame containing all isotherm data.

        """
        if branch is None or branch.startswith('all'):
            return self.data_raw
        if branch == 'ads':
            return self.data_raw.loc[self.data_raw['branch'] == 0]
        if branch == 'des':
            return self.data_raw.loc[self.data_raw['branch'] == 1]
        raise ParameterError('Bad branch specification.')

    def pressure(
        self,
        branch: str = None,
        pressure_unit: str = None,
        pressure_mode: str = None,
        limits: t.Tuple[float, float] = None,
        indexed: bool = False,
    ) -> t.Union[numpy.ndarray, pandas.Series]:
        """
        Return pressure points as an array.

        Parameters
        ----------
        branch : {None, 'ads', 'des'}
            The branch of the pressure to return. If ``None``, returns entire
            dataset.
        pressure_unit : str, optional
            Unit in which the pressure should be returned. If ``None``
            it defaults to which pressure unit the isotherm is currently in.
        pressure_mode : {None, 'absolute', 'relative', 'relative%'}
            The mode in which to return the pressure, if possible. If ``None``,
            returns mode the isotherm is currently in.
        limits : [float, float], optional
            Minimum and maximum pressure limits.
            Put None or -+np.inf for no limit.
        indexed : bool, optional
            If this is specified to true, then the function returns an indexed
            pandas.Series instead of an array.

        Returns
        -------
        array or Series
            The pressure slice corresponding to the parameters passed.

        """
        ret = self.data(branch=branch).loc[:, self.pressure_key]

        if not ret.empty:
            # Convert if needed
            if pressure_mode or pressure_unit:
                # If pressure mode not given, try current
                if not pressure_mode:
                    pressure_mode = self.pressure_mode
                # If pressure unit not given, try current
                if not pressure_unit:
                    pressure_unit = self.pressure_unit

                ret = c_pressure(
                    ret,
                    mode_from=self.pressure_mode,
                    mode_to=pressure_mode,
                    unit_from=self.pressure_unit,
                    unit_to=pressure_unit,
                    adsorbate=self.adsorbate,
                    temp=self.temperature
                )

            # Select required points
            if limits and any(limits):
                ret = ret.loc[ret.between(
                    -numpy.inf if limits[0] is None else limits[0],
                    numpy.inf if limits[1] is None else limits[1]
                )]

        if indexed:
            return ret
        return ret.values

    def loading(
        self,
        branch: str = None,
        loading_unit: str = None,
        loading_basis: str = None,
        material_unit: str = None,
        material_basis: str = None,
        limits: t.Tuple[float, float] = None,
        indexed: bool = False
    ) -> t.Union[numpy.ndarray, pandas.Series]:
        """
        Return loading points as an array.

        Parameters
        ----------
        branch : {None, 'ads', 'des'}
            The branch of the loading to return. If ``None``, returns entire
            dataset.
        loading_unit : str, optional
            Unit in which the loading should be returned. If ``None``
            it defaults to which loading unit the isotherm is currently in.
        loading_basis : {None, 'mass', 'volume_gas', 'volume_liquid', 'molar'}
            The basis on which to return the loading, if possible. If ``None``,
            returns on the basis the isotherm is currently in.
        material_unit : str, optional
            Unit in which the material should be returned. If ``None``
            it defaults to which loading unit the isotherm is currently in.
        material_basis : {None, 'mass', 'volume', 'molar'}
            The basis on which to return the material, if possible. If ``None``,
            returns on the basis the isotherm is currently in.
        limits : [float, float], optional
            Minimum and maximum loading limits.
            Put None or -+np.inf for no limit.
        indexed : bool, optional
            If this is specified to true, then the function returns an indexed
            pandas.Series instead of an array.

        Returns
        -------
        Array or Series
            The loading slice corresponding to the parameters passed.

        """
        ret = self.data(branch=branch).loc[:, self.loading_key]

        if not ret.empty:
            # Convert if needed

            # First adsorbent is converted
            if material_basis or material_unit:
                if not material_basis:
                    material_basis = self.material_basis

                ret = c_material(
                    ret,
                    basis_from=self.material_basis,
                    basis_to=material_basis,
                    unit_from=self.material_unit,
                    unit_to=material_unit,
                    material=self.material
                )

            # Then loading
            if loading_basis or loading_unit:
                if not loading_basis:
                    loading_basis = self.loading_basis

                # These must be specified
                # in the case of fractional conversions
                if not material_basis:
                    material_basis = self.material_basis
                if not material_unit:
                    material_unit = self.material_unit

                ret = c_loading(
                    ret,
                    basis_from=self.loading_basis,
                    basis_to=loading_basis,
                    unit_from=self.loading_unit,
                    unit_to=loading_unit,
                    adsorbate=self.adsorbate,
                    temp=self.temperature,
                    basis_material=material_basis,
                    unit_material=material_unit,
                )

            # Select required points
            if limits and any(limits):
                ret = ret.loc[ret.between(
                    -numpy.inf if limits[0] is None else limits[0],
                    numpy.inf if limits[1] is None else limits[1]
                )]

        if indexed:
            return ret
        return ret.values

    @property
    def other_keys(self):
        """
        Return column names of any supplementary data points.
        """
        return [
            c for c in self.data_raw.columns
            if c not in (self.pressure_key, self.loading_key, 'branch')
        ]

    def other_data(
        self,
        key: str,
        branch: str = None,
        limits: t.Tuple[float, float] = None,
        indexed: bool = False,
    ) -> t.Union[numpy.ndarray, pandas.Series]:
        """
        Return supplementary data points as an array.

        Parameters
        ----------
        key : str
            Key in the isotherm DataFrame containing the data to select.
        branch : {None, 'ads', 'des'}
            The branch of the data to return. If ``None``, returns entire
            dataset.
        limits : [float, float], optional
            Minimum and maximum data limits.
            Put None or -+np.inf for no limit.
        indexed : bool, optional
            If this is specified to true, then the function returns an indexed
            pandas.Series instead of an array.

        Returns
        -------
        array or Series
            The data slice corresponding to the parameters passed.

        """
        if key in self.other_keys:
            ret = self.data(branch=branch).loc[:, key]

            if not ret.empty:
                # Select required points
                if limits and any(limits):
                    ret = ret.loc[ret.between(
                        -numpy.inf if limits[0] is None else limits[0],
                        numpy.inf if limits[1] is None else limits[1]
                    )]

            if indexed:
                return ret
            return ret.values

        raise ParameterError(f"Isotherm does not contain any {key} data.")

    def has_branch(self, branch: str) -> bool:
        """
        Check if the isotherm has an specific branch.

        Parameters
        ----------
        branch : {None, 'ads', 'des'}
            The branch of the data to check for.

        Returns
        -------
        bool
            Whether the data exists or not.

        """
        return not self.data(branch=branch).empty

    ##########################################################
    #   Functions that interpolate values of the isotherm data

    def pressure_at(
        self,
        loading: t.List[float],
        branch: str = 'ads',
        interpolation_type: str = 'linear',
        interp_fill: t.Union[float, t.Tuple[float, float], str] = None,
        pressure_unit: str = None,
        pressure_mode: str = None,
        loading_unit: str = None,
        loading_basis: str = None,
        material_unit: str = None,
        material_basis: str = None,
    ) -> numpy.ndarray:
        """
        Interpolate isotherm to compute pressure at any loading given.

        Parameters
        ----------
        loading : float
            Loading at which to compute pressure.
        branch : {'ads', 'des'}
            The branch of the use for calculation. Defaults to adsorption.
        interpolation_type : str
            The type of scipy.interp1d used: `linear`, `nearest`, `zero`,
            `slinear`, `quadratic`, `cubic`. It defaults to `linear`.
        interp_fill : array-like or (array-like, array_like) or “extrapolate”, optional
            Parameter to determine what to do outside data bounds.
            Passed to the scipy.interpolate.interp1d function as ``fill_value``.
            If blank, interpolation will not predict outside the bounds of data.

        pressure_unit : str
            Unit the pressure is returned in. If ``None``, it defaults to
            internal isotherm units.
        pressure_mode : str
            The mode the pressure is returned in. If ``None``, it defaults to
            internal isotherm mode.

        loading_unit : str
            Unit the loading is specified in. If ``None``, it defaults to
            internal isotherm units.
        loading_basis : {None, 'mass', 'molar', 'volume_gas', 'volume_liquid'}
            The basis the loading is specified in. If ``None``,
            assumes the basis the isotherm is currently in.
        material_unit : str, optional
            Unit in which the material is passed in. If ``None``
            it defaults to which loading unit the isotherm is currently in
        material_basis : str
            The basis the loading is passed in. If ``None``, it defaults to
            internal isotherm basis.

        Returns
        -------
        float
            Predicted pressure at loading specified.

        """
        # Convert to numpy array just in case
        loading = numpy.asarray(loading)

        # Check if interpolator is applicable
        if (
            self.p_interpolator is None or self.p_interpolator.interp_branch != branch
            or self.p_interpolator.interp_kind != interpolation_type
            or self.p_interpolator.interp_fill != interp_fill
        ):
            self.p_interpolator = IsothermInterpolator(
                self.loading(branch=branch),
                self.pressure(branch=branch),
                interp_branch=branch,
                interp_kind=interpolation_type,
                interp_fill=interp_fill
            )

        # Ensure loading is in correct units and basis for the internal model
        if material_basis or material_unit:
            if not material_basis:
                material_basis = self.material_basis
            if not material_unit:
                raise ParameterError(
                    "Must specify an material unit if the input is in another basis."
                )

            loading = c_material(
                loading,
                basis_from=material_basis,
                basis_to=self.material_basis,
                unit_from=material_unit,
                unit_to=self.material_unit,
                material=self.material
            )

        if loading_basis or loading_unit:
            if not loading_basis:
                loading_basis = self.loading_basis
            if not loading_unit:
                raise ParameterError(
                    "Must specify a loading unit if the input is in another basis."
                )

            loading = c_loading(
                loading,
                basis_from=loading_basis,
                basis_to=self.loading_basis,
                unit_from=loading_unit,
                unit_to=self.loading_unit,
                adsorbate=self.adsorbate,
                temp=self.temperature,
                basis_material=self.material_basis,
                unit_material=self.material_unit,
            )

        # Interpolate using the internal interpolator
        pressure = self.p_interpolator(loading)

        # Ensure pressure is in correct units and mode requested
        if pressure_mode or pressure_unit:
            if not pressure_mode:
                pressure_mode = self.pressure_mode

            pressure = c_pressure(
                pressure,
                mode_from=self.pressure_mode,
                mode_to=pressure_mode,
                unit_from=self.pressure_unit,
                unit_to=pressure_unit,
                adsorbate=self.adsorbate,
                temp=self.temperature
            )

        return pressure

    def loading_at(
        self,
        pressure: t.List[float],
        branch: str = 'ads',
        interpolation_type: str = 'linear',
        interp_fill: t.Union[float, t.Tuple[float, float], str] = None,
        pressure_unit: str = None,
        pressure_mode: str = None,
        loading_unit: str = None,
        loading_basis: str = None,
        material_unit: str = None,
        material_basis: str = None,
    ) -> numpy.ndarray:
        """
        Interpolate isotherm to compute loading at any pressure given.

        Parameters
        ----------
        pressure : float or array
            Pressure at which to compute loading.
        branch : {'ads', 'des'}
            The branch the interpolation takes into account.
        interpolation_type : str
            The type of scipy.interp1d used: `linear`, `nearest`, `zero`,
            `slinear`, `quadratic`, `cubic`. It defaults to `linear`.
        interp_fill : array-like or (array-like, array_like) or “extrapolate”, optional
            Parameter to determine what to do outside data bounds.
            Passed to the scipy.interpolate.interp1d function as ``fill_value``.
            If blank, interpolation will not predict outside the bounds of data.

        pressure_unit : str
            Unit the pressure is specified in. If ``None``, it defaults to
            internal isotherm units.
        pressure_mode : str
            The mode the pressure is passed in. If ``None``, it defaults to
            internal isotherm mode.

        loading_unit : str, optional
            Unit in which the loading should be returned. If ``None``
            it defaults to which loading unit the isotherm is currently in.
        loading_basis : {None, 'mass', 'molar', 'volume_gas', 'volume_liquid'}
            The basis on which to return the loading, if possible. If ``None``,
            returns on the basis the isotherm is currently in.
        material_unit : str, optional
            Material unit in which the data should be returned. If ``None``
            it defaults to which loading unit the isotherm is currently in.
        material_basis : {None, 'mass', 'volume', 'molar'}
            Material basis on which to return the data, if possible. If ``None``,
            returns on the basis the isotherm is currently in.

        Returns
        -------
        float or array
            Predicted loading at pressure P.

        """
        # Convert to a numpy array just in case
        pressure = numpy.asarray(pressure)

        # Check if interpolator is applicable
        if (
            self.l_interpolator is None or self.l_interpolator.interp_branch != branch
            or self.l_interpolator.interp_kind != interpolation_type
            or self.l_interpolator.interp_fill != interp_fill
        ):
            self.l_interpolator = IsothermInterpolator(
                self.pressure(branch=branch),
                self.loading(branch=branch),
                interp_branch=branch,
                interp_kind=interpolation_type,
                interp_fill=interp_fill
            )

        # Ensure pressure is in correct units and mode for the internal model
        if pressure_mode or pressure_unit:
            if not pressure_mode:
                pressure_mode = self.pressure_mode
            if pressure_mode == 'absolute' and not pressure_unit:
                raise ParameterError(
                    "Must specify a pressure unit if the input is in an absolute mode."
                )

            pressure = c_pressure(
                pressure,
                mode_from=pressure_mode,
                mode_to=self.pressure_mode,
                unit_from=pressure_unit,
                unit_to=self.pressure_unit,
                adsorbate=self.adsorbate,
                temp=self.temperature
            )

        # Interpolate using the internal interpolator
        loading = self.l_interpolator(pressure)

        # Ensure loading is in correct units and basis requested
        if material_basis or material_unit:

            if not material_basis:
                material_basis = self.material_basis

            loading = c_material(
                loading,
                basis_from=self.material_basis,
                basis_to=material_basis,
                unit_from=self.material_unit,
                unit_to=material_unit,
                material=self.material
            )

        if loading_basis or loading_unit:
            if not loading_basis:
                loading_basis = self.loading_basis

            loading = c_loading(
                loading,
                basis_from=self.loading_basis,
                basis_to=loading_basis,
                unit_from=self.loading_unit,
                unit_to=loading_unit,
                adsorbate=self.adsorbate,
                temp=self.temperature,
                basis_material=self.material_basis,
                unit_material=self.material_unit,
            )

        return loading

    def spreading_pressure_at(
        self,
        pressure: t.List[float],
        branch: str = 'ads',
        pressure_unit: str = None,
        pressure_mode: str = None,
        loading_unit: str = None,
        loading_basis: str = None,
        material_unit: str = None,
        material_basis: str = None,
        interp_fill: t.Union[float, t.Tuple[float, float], str] = None,
    ) -> numpy.ndarray:
        r"""
        Calculate reduced spreading pressure at a bulk adsorbate pressure P.

        Use numerical quadrature on isotherm data points to compute the reduced
        spreading pressure via the integral:

        .. math::

            \Pi(p) = \int_0^p \frac{q(\hat{p})}{ \hat{p}} d\hat{p}.

        In this integral, the isotherm :math:`q(\hat{p})` is represented by a
        linear interpolation of the data.

        For in-detail explanations, check reference [#]_.

        Parameters
        ----------
        pressure : float
            Pressure (in corresponding units as data in instantiation).
        branch : {'ads', 'des'}
            The branch of the use for calculation. Defaults to adsorption.
        loading_unit : str
            Unit the loading is specified in. If ``None``, it defaults to
            internal isotherm units.
        pressure_unit : str
            Unit the pressure is returned in. If ``None``, it defaults to
            internal isotherm units.
        material_basis : str
            The basis the loading is passed in. If ``None``, it defaults to
            internal isotherm basis.
        pressure_mode : str
            The mode the pressure is returned in. If ``None``, it defaults to
            internal isotherm mode.
        interp_fill : array-like or (array-like, array_like) or “extrapolate”, optional
            Parameter to determine what to do outside data bounds.
            Passed to the scipy.interpolate.interp1d function as ``fill_value``.
            If blank, interpolation will not predict outside the bounds of data.

        Returns
        -------
        float
            Spreading pressure, :math:`\Pi`.

        References
        ----------
        .. [#] C. Simon, B. Smit, M. Haranczyk. pyIAST: Ideal Adsorbed Solution
           Theory (IAST) Python Package. Computer Physics Communications.

        """
        # Get all data points
        pressures = self.pressure(
            branch=branch, pressure_unit=pressure_unit, pressure_mode=pressure_mode
        )
        loadings = self.loading(
            branch=branch,
            loading_unit=loading_unit,
            loading_basis=loading_basis,
            material_unit=material_unit,
            material_basis=material_basis
        )

        # throw exception if interpolating outside the range.
        if (self.l_interpolator is not None and self.l_interpolator.interp_fill is None) & \
                (pressure > pressures.max() or pressure < pressures.min()):
            raise CalculationError(
                textwrap.dedent(
                    f"""
                To compute the spreading pressure at this bulk adsorbate pressure,
                we would need to extrapolate the isotherm since this pressure ({pressure:.3g} {self.pressure_unit})
                is outside the range of the highest pressure in your pure-component
                isotherm data ({pressures.max()} {self.pressure_unit}).

                At present, the PointIsotherm class is set to throw an exception
                when this occurs, as we do not have data outside this pressure range
                to characterize the isotherm at higher pressures.

                Option 1: fit an analytical model to extrapolate the isotherm
                Option 2: pass a `interp_fill` to the spreading pressure function of the
                    PointIsotherm object. Then, that PointIsotherm will
                    assume that the uptake beyond {pressures.max()} {self.pressure_unit} is given by
                    `interp_fill`. This is reasonable if your isotherm data exhibits
                    a plateau at the highest pressures.
                Option 3: Go back to the lab or computer to collect isotherm data
                    at higher pressures. (Extrapolation can be dangerous!)
                """
                )
            )

        # approximate loading up to first pressure point with Henry's law
        # loading = henry_const * P
        # henry_const is the initial slope in the adsorption isotherm
        henry_const = loadings[0] / pressures[0]

        # get how many of the points are less than pressure P
        n_points = numpy.sum(pressures < pressure)

        if n_points == 0:
            # if this pressure is between 0 and first pressure point...
            # \int_0^P henry_const P /P dP = henry_const * P ...
            return henry_const * pressure

        # P > first pressure point
        area = loadings[0]  # area of first segment \int_0^P_1 n(P)/P dP

        # get area between P_1 and P_k, where P_k < P < P_{k+1}
        for i in range(n_points - 1):
            # linear interpolation of isotherm data
            slope = (loadings[i + 1] - loadings[i]) / (pressures[i + 1] - pressures[i])
            intercept = loadings[i] - slope * pressures[i]
            # add area of this segment
            area += slope * (pressures[i + 1] - pressures[i]) + intercept * \
                numpy.log(pressures[i + 1] / pressures[i])

        # finally, area of last segment
        slope = (
            self.loading_at(
                pressure,
                branch=branch,
                pressure_unit=pressure_unit,
                pressure_mode=pressure_mode,
                loading_unit=loading_unit,
                loading_basis=loading_basis,
                material_unit=material_unit,
                material_basis=material_basis,
                interp_fill=interp_fill
            ) - loadings[n_points - 1]
        ) / (pressure - pressures[n_points - 1])

        intercept = loadings[n_points - 1] - \
            slope * pressures[n_points - 1]
        area += slope * (pressure - pressures[n_points - 1]) + intercept * \
            numpy.log(pressure / pressures[n_points - 1])

        return area


# coding: utf-8

""" Transform a proper motion to/from Galactic to/from ICRS coordinates """

from __future__ import division, print_function

# Standard library
import warnings

# Third-party
import numpy as np
import astropy.coordinates as coord

__all__ = ['transform_proper_motion', 'pm_gal_to_icrs', 'pm_icrs_to_gal']

def transform_proper_motion(coordinate, pm, new_frame):
    """
    Transform the proper motion vector ``pm`` at the sky coordinate
    ``coordinate`` to the desired coordinate frame ``new_frame``.

    Parameters
    ----------
    coordinate : :class:`~astropy.coordinates.SkyCoord`, :class:`~astropy.coordinates.BaseCoordinateFrame`
        An instance of an Astropy coordinate object.
    pm : :class:`~astropy.units.Quantity`, iterable
        Proper motion components (longitude, latitude) in the same frame
        as the input coordinates. Can either be an iterable of two
        :class:`~astropy.units.Quantity` objects or a single
        :class:`~astropy.units.Quantity` with shape (2,N).
        The proper motion in longitude is assumed to already include
        the cos(latitude) term.
    new_frame : :class:`~astropy.coordinates.BaseCoordinateFrame` subclass
        The desired coordinate frame of the proper motion.

    Returns
    -------
    new_pm : tuple
        A length-2 tuple containing the proper motion components
        (longitude, latitude) in the new frame. The longitude component is
        includes the cos(latitude) term.

    Examples
    --------

        >>> import astropy.units as u
        >>> import astropy.coordinates as coord
        >>> c = coord.SkyCoord(ra=196.5*u.degree, dec=-10.33*u.deg, distance=16.2*u.kpc)
        >>> pm = [-1.53, 3.5]*u.mas/u.yr
        >>> transform_proper_motion(c, pm, coord.Galactic) # doctest: +FLOAT_CMP
        <Quantity [-1.19944367, 3.62660101] mas / yr>

    """
    url = "http://docs.astropy.org/en/stable/coordinates/velocities.html"
    warnings.warn("This function is now deprecated. Use the velocity "
                  "transformation functionality in Astropy instead. For more "
                  "information, see: {0}".format(url), DeprecationWarning)

    if hasattr(coordinate, 'frame'):
        coordinate = coordinate.frame

    frame_cls = coordinate.__class__
    c = frame_cls(coordinate.data.with_differentials(
        coord.UnitSphericalCosLatDifferential(*pm)))
    new_c = c.transform_to(new_frame)
    diff = new_c.data.differentials['s']
    pm = np.vstack((diff.d_lon_coslat.value, diff.d_lat.value)) * diff.d_lat.unit

    return pm.reshape((2,) + c.shape)

# ----------------------------------------------------------------------------
# Deprecated:
#

def pm_gal_to_icrs(coordinate, pm):
    r"""
    Convert proper motion in Galactic coordinates (l,b) to
    ICRS coordinates (RA, Dec).

    Parameters
    ----------
    coordinate : :class:`~astropy.coordinates.SkyCoord`, :class:`~astropy.coordinates.BaseCoordinateFrame`
        An instance of an Astropy coordinate object. Can be in any
        frame that is transformable to ICRS coordinates.
    pm : :class:`~astropy.units.Quantity`, iterable
        Full description of proper motion in Galactic longitude and
        latitude. Can either be a tuple of two
        :class:`~astropy.units.Quantity` objects or a single
        :class:`~astropy.units.Quantity` with shape (2,N).
        The proper motion in longitude is assumed to be multipled by
        cosine of Galactic latitude, :math:`\mu_l\cos b`.

    Returns
    -------
    pm : :class:`~astropy.units.Quantity`
        An astropy :class:`~astropy.units.Quantity` object specifying the
        proper motion vector array in ICRS coordinates. Will have shape
        (2,N).

    Examples
    --------

        >>> import astropy.units as u
        >>> import astropy.coordinates as coord
        >>> c = coord.SkyCoord(ra=196.5*u.degree, dec=-10.33*u.deg, distance=16.2*u.kpc)
        >>> pm = [-1.53, 3.5]*u.mas/u.yr
        >>> pm_gal_to_icrs(c, pm) # doctest: +FLOAT_CMP
        <Quantity [-1.84741767, 3.34334366] mas / yr>

    """

    # Note: Deprecation warning gets emitted from calling
    #       transform_proper_motion below.
    g = coordinate.transform_to(coord.Galactic)
    return transform_proper_motion(g, pm, coord.ICRS)

def pm_icrs_to_gal(coordinate, pm):
    r"""
    Convert proper motion in ICRS coordinates (RA, Dec) to
    Galactic coordinates (l,b).

    Parameters
    ----------
    coordinate : :class:`~astropy.coordinates.SkyCoord`, :class:`~astropy.coordinates.BaseCoordinateFrame`
        An instance of an Astropy coordinate object. Can be in any
        frame that is transformable to ICRS coordinates.
    pm : :class:`~astropy.units.Quantity`, iterable
        Full description of proper motion in Right ascension (RA) and
        declination (Dec). Can either be a tuple of two
        :class:`~astropy.units.Quantity` objects or a single
        :class:`~astropy.units.Quantity` with shape (2,N).
        The proper motion in RA is assumed to be multipled by
        cosine of declination, :math:`\mu_\alpha\cos\delta`.

    Returns
    -------
    pm : :class:`~astropy.units.Quantity`
        An astropy :class:`~astropy.units.Quantity` object specifying the
        proper motion vector array in Galactic coordinates. Will have shape
        (2,N).

    Examples
    --------

        >>> import astropy.units as u
        >>> import astropy.coordinates as coord
        >>> c = coord.SkyCoord(ra=196.5*u.degree, dec=-10.33*u.deg, distance=16.2*u.kpc)
        >>> pm = [-1.84741767, 3.34334366]*u.mas/u.yr
        >>> pm_icrs_to_gal(c, pm) # doctest: +FLOAT_CMP
        <Quantity [-1.52999988, 3.49999973] mas / yr>

    """

    # Note: Deprecation warning gets emitted from calling
    #       transform_proper_motion below.
    i = coordinate.transform_to(coord.ICRS)
    return transform_proper_motion(i, pm, coord.Galactic)


#!/usr/bin/python
#-*- coding: utf-8 -*-

# >.>.>.>.>.>.>.>.>.>.>.>.>.>.>.>.
# Licensed under the Apache License, Version 2.0 (the "License")
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0

# --- File Name: tpl_score.py
# --- Creation Date: 13-10-2020
# --- Last Modified: Wed 18 Nov 2020 18:42:10 AEDT
# --- Author: Xinqi Zhu
# .<.<.<.<.<.<.<.<.<.<.<.<.<.<.<.<
"""Implementation of the TPL score.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
import sys
import pickle
sys.path.insert(
    0,
    os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))),
                 'disentanglement_lib'))
sys.path.insert(
    0,
    os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))),
                 'stylegan2'))
import dnnlib.tflib
from absl import logging
from disentanglement_lib.evaluation.metrics import utils
import numpy as np
from six.moves import range
from scipy import stats
import gin.tf

@gin.configurable("tpl_score",
                  blacklist=[
                      "ground_truth_data", "generator_function",
                      "random_state", "activation_str", "latent_size",
                      "artifact_dir"
                  ])
def compute_tpl_score(ground_truth_data,
                      generator_function,
                      random_state,
                      activation_str,
                      latent_size,
                      artifact_dir=None,
                      batch_size=gin.REQUIRED,
                      num_traversals=gin.REQUIRED,
                      num_samples_per_dim=gin.REQUIRED,
                      traversal_bound=gin.REQUIRED,
                      active_thresh=gin.REQUIRED):
    """Computes the FactorVAE disentanglement metric.

  Args:
    ground_truth_data: GroundTruthData to be sampled from.
    generator_function: Function that takes latent code as input and
      outputs an image.
    random_state: Numpy random state used for randomness.
    artifact_dir: Optional path to directory where artifacts can be saved.
    latent_size: Latent code size.
    batch_size: Number of points to be used to in a GPU.
    num_traversals: Number of traversals to sample.
    num_samples_per_dim: Number of samples per dim.
    traversal_bound: The bound for traversal.
    active_thresh: Threshold for determining active dims.

  Returns:
    Dictionary with scores:
      avg_tpl_dim: Mean TPL score for each dim.
      avg_tpl: Overall mean TPL score.
      active_mask: Mask of which dims are active.
      active_distances: Latent distances of active dims.
      active_stds: Latent stds of active dims.
      n_active_dims: Number of active dims.
  """
    del ground_truth_data
    del artifact_dir
    dnnlib.tflib.init_tf()
    distance_measure = load_pkl(
        '.stylegan2-cache/vgg16_zhang_perceptual.pkl'
        # 'http://d36zk2xti64re0.cloudfront.net/stylegan1/networks/metrics/vgg16_zhang_perceptual.pkl'
    )
    if activation_str == "'logits'":
        activation = sigmoid
    elif activation_str == "'tanh'":
        activation = tanh
    else:
        raise ValueError(
            "Activation function  could not be infered from gin config.")
    tpl_dim_ls = []
    for i in range(num_traversals):
        sample = random_state.normal(size=(1, latent_size))
        # factor_index = random_state.randint(ground_truth_data.num_factors)
        tpl_dim = compute_tpl_for(
            sample, generator_function, batch_size, num_samples_per_dim,
            latent_size, activation,
            distance_measure, traversal_bound)  # np.array of [latent_size]
        tpl_dim_ls.append(tpl_dim)
    tpl_dim_np = np.array(tpl_dim_ls)  # np.array of [n_trav, latent_size]
    avg_tpl_dim = np.mean(tpl_dim_np, axis=0)  # np.array of [latent_size]
    std_tpl_dim = np.std(tpl_dim_np, axis=0)  # np.array of [latent_size]

    active_mask = avg_tpl_dim > active_thresh
    active_distances = np.extract(active_mask, avg_tpl_dim)
    active_stds = np.extract(active_mask, std_tpl_dim)
    avg_tpl = np.sum(active_distances)
    n_active_dims = active_mask.astype(int).sum()

    scores_dict = {}
    scores_dict['avg_tpl_dim'] = avg_tpl_dim.tolist()
    scores_dict['avg_tpl'] = avg_tpl
    scores_dict['active_mask'] = active_mask.tolist()
    scores_dict['active_distances'] = active_distances.tolist()
    scores_dict['active_stds'] = active_stds.tolist()
    scores_dict['n_active_dims'] = n_active_dims
    # print('scores_dict:', scores_dict)
    return scores_dict

def compute_tpl_for(sample, generator_function, batch_size,
                    num_samples_per_dim, latent_size, activation,
                    distance_measure, traversal_bound):
    '''
    Return: np.array of [latent_size]
    '''
    tpl_sample_dim_ls = []
    for i in range(latent_size):
        samples_traversal = np.tile(sample, [num_samples_per_dim, 1])
        samples_traversal[:, i] = np.linspace(-traversal_bound, traversal_bound, num=num_samples_per_dim)
        raw_imgs_traversal = generator_function(
            samples_traversal)  # size: [b, h, w, c]
        imgs_traversal = activation(raw_imgs_traversal)
        j = 0
        traversal_score = 0
        while j < num_samples_per_dim:
            b_j = min(batch_size, num_samples_per_dim - j)
            cur_imgs_traversal = imgs_traversal[j:j + b_j, ...]
            j += b_j
            traversal_score += measure_distance(cur_imgs_traversal,
                                                distance_measure)
        tpl_sample_dim_ls.append(traversal_score)
    return np.array(tpl_sample_dim_ls)

def measure_distance(imgs_traversal, distance_measure):
    images = np.transpose(imgs_traversal, [0, 3, 1, 2])  # bhwc -> bchw
    images = images * 255  # [0, 1] -> [0, 255]
    v = get_return_v(distance_measure.run(images[:-1, ...], images[1:, ...]),
                     1)
    return v.sum()

def sigmoid(x):
    return stats.logistic.cdf(x)

def tanh(x):
    return np.tanh(x) / 2. + .5

def get_return_v(x, topk=1):
    if (not isinstance(x, tuple)) and (not isinstance(x, list)):
        return x if topk == 1 else tuple([x] + [None] * (topk - 1))
    if topk > len(x):
        return tuple(list(x) + [None] * (topk - len(x)))
    else:
        if topk == 1:
            return x[0]
        else:
            return tuple(x[:topk])

def load_pkl(file_or_url):
    with open(file_or_url, 'rb') as file:
        return pickle.load(file, encoding='latin1')


# -*- coding: utf-8 -*-
'''
Created on Nov 24, 2014

@author: pascale
'''

import rmodel
import numpy as np
import logging
import rttov
import copy
import matplotlib.pyplot as plt
import os
import h5py
import shutil
from r1Dvar import r1dvarObjects

q_mixration_to_ppmv = 1.60771704e+6

def vector2profile(vector, exprofile):
    """ convert vector to profile """
    """ take exprofile to begin with """
    """ a vector contains only T values, """
    """ lnq bottom lnq vales , Tsurf, lnq surf and Tskin"""
    """ temperature, and ln q are stored from top of atmosphere to ground """

    nlevels = vector.shape[0] - 29 - 3
    if nlevels != exprofile['NLEVELS']:
        logging.error("vector and profile are not of the same dimension")
        return None

    profile = copy.deepcopy(exprofile)

    for i in range(nlevels):
        profile["T"][i] = vector[i]

    for i in range(29):
        profile["Q"][nlevels - 29 + i] = vector[nlevels + i]

    profile["S2M"]["T"] = vector[nlevels + 29]
    profile["S2M"]["Q"] = vector[nlevels + 29 + 1]
    profile["SKIN"]["T"] = vector[nlevels + 29 + 2]

    return profile

class Project1dvar(object):

    def __init__(self, pTrue, pBg, satellite, instrument,
                 matrixBfile=None, channel_list=None,
                 changeTheBackground=False):
        """ initialize a 1DvarProject with a True and a Background project """
        """ if channel_list is present : we work on a restricted list """
        """ of channel """
        """ but in fact channel_list is set from the Coefficient """
        """ object of pBg """
        """ so we use this list to extract a sub-matrix from the Rmatrix """
        """ satellite must be = platform-num (ex : metop-1 , noaa-19) """
        """ if changeTheBackground=False we don not change the Background """
        """ if changeTheBackground=False at each new step the previously  """
        """ retrieved profile becomes the new background """

        self.changeTheBackground = changeTheBackground
        self.ENV = {}
        self.pTrue = pTrue
        self.pBg = pBg
        # set distinct filenames for the true and the background project
        # (we will run rttov with different files)
        self.pTrue.setFileNameMark("True")
        self.pBg.setFileNameMark("Bg")
        logging.info("Initialize a Project1dvar object")
        logging.info("are pBg coefficients loaded ? " +
                     str(self.pBg.myCoeffs.loadCoeffs))
        self.satellite = satellite
        self.instrument = instrument

        self.BgInitialProfile = copy.deepcopy(self.pBg.myProfile)
        self.BmatrixFileName = None
        self.RmatrixFilename = None

        self._SetConfig()
        self.factorB = 1.
        self.factorR = 1.
        self.MaxNoise = 1.
        self.retrievedVectors = []
        self.retrievedProfiles = []
        self.Rmatrix = Rmatrix()
        self.Bmatrix = Bmatrix()
        # read R matrix
        self.initCoeffsandReadRmatrix()
        # Bmatrix
        if matrixBfile:
            self.BmatrixFileName = matrixBfile
        if not os.path.exists(self.BmatrixFileName):
            logging.error(
                "Sorry something is wrong : I cannot find the Bmatrix :"
                " check the presence of the nwpsaf-1dvar files ")
            raise IOError
        # read the Bmatrix
        B = Bmatrix()
        B.read_matrices(self.BmatrixFileName)
        # The file contains 4 matrix but according to Peter the 4 matrix are
        # identical
        self.Bmat = B.matrices[1]
        self.step_counter = 0

    def initCoeffsandReadRmatrix(self):
        """ initialise myCoeffs for pBg project and read the R matrix """
        """ the R matrix is computed according to the list of channels """
        if not self.pBg.myCoeffs.loadCoeffs:
            logging.info("load coefficients for pBg")
            err = self.pBg.loadCoefficients()
            if err != 0:
                raise RuntimeError("Cannot load coefficient in r1dvar project")
        nbChannels = self.pBg.myCoeffs.nchannels
        logging.info("nbChannels: {}".format(nbChannels))
        # initialise the  1dvarproject channel_list used by readRmatrix
        my_chan_list = self.pBg.myCoeffs.getFF_ORI_CHN()
        # must start at 0
        self.channel_list = [my_chan_list[i] -
                             1 for i in range(len(my_chan_list))]
        # read the R matrix

        self.readRmatrix(self.satellite, self.instrument)

    def reinitCoeff(self, satellite, instrument, coeffs):
        """ must be called if we change the coefficients """
        """ (select channel or change instrument ) """
        """  satellite must be = platform-num (ex : metop-1 , noaa-19) """
        self.pBg.myCoeffs = coeffs
        self.satellite = satellite
        self.instrument = instrument

        self.initCoeffsandReadRmatrix()

    def _SetConfig(self):
        try:
            self.ENV["NWPSAF_1DVAR_HOME"] = os.environ("NWPSAF_1DVAR_HOME")
        except:
            logging.warning(
                "The environment variable NWPSAF_1DVAR_HOME is not set :"
                " no worries I can use my matrix ..")
            self.ENV["NWPSAF_1DVAR_HOME"] = os.environ[
                "RTTOV_GUI_PREFIX"] + "/r1Dvar/data"
        self.BmatrixFileName = self.ENV[
            "NWPSAF_1DVAR_HOME"] + "/Sample_Bmatrices/Bmatrix_54L"

        self.RmatrixFilenames = {
            "AIRS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "AIRS_COEFFS_DIR/Rmatrix_orig"),
            "ATMS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "ATMS_COEFFS_DIR/Rmatrix_orig"),
            "AMSU-A": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                   "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "AMSU-B": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                   "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "HIRS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "MHS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "mhs": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "IASI": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "IASI_COEFFS_DIR/Rmatrix_orig"),
            "CRIS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "CrIS_COEFFS_DIR/Rmatrix_orig"),
            "SSMIS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                  "SSMIS_COEFFS_DIR/Rmatrix_orig"),
            "CrIS": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "CrIS_COEFFS_DIR/Rmatrix_orig"),
            "amsua": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                  "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "amsub": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                  "ATOVS_COEFFS_DIR/Rmatrix_orig"),
            "hirs": os.path.join(self.ENV["NWPSAF_1DVAR_HOME"],
                                 "ATOVS_COEFFS_DIR/Rmatrix_orig")
        }

    def readRmatrix(self, satellite, instrument, filename=None):
        """ read the observation error matrix file and init self.Rmat """
        """ if filename is specified then it has the priority """
        """ first step : determine the file name from the instrument """
        """ (stored in rge Rmatrixfilenames dictionary"""
        """ create a Rmatrix objet which will contains all the Rmatrix """
        """ read in a file """
        """ then determine the actual Rmat from satellite and instrument """
        """ ATOVS case is particular : must retrieve HIRS, AMSU-A and MHS"""
        """ for some other files it is trivial : they contains just one """
        """ matrix """
        """  satellite must be = platform-num (ex : metop-1 , noaa-19) """
        # when this method is call self.channel_list exit and is ready (no need
        # to shift indices)
        logging.info("read R matrix for " + satellite + " " + instrument)
        isATOV = False
        if instrument in ("hirs", "amsua", "amsub", "mhs"):
            isATOV = True
        else:
            if instrument == "cris":
                instrument = "CrIS"
            else:
                instrument = instrument.upper()
        if satellite == "metop-2":
            satellite = "MetOp-A"
        else:
            if satellite == "metop-1":
                satellite = "MetOp-B"
            else:
                satellite = satellite.upper()

        logging.info("instrument :" + str(instrument) + " " +
                     str(satellite) + " is ATOV ?: " + str(isATOV))
        if filename is not None:
            self.RmatrixFilename = filename
        else:
            self.instrument = instrument
            self.satellite = satellite

            self.RmatrixFilename = self.RmatrixFilenames[instrument]
            try:
                self.RmatrixFilename = self.RmatrixFilenames[instrument]
                logging.info(" RmatrixFilename : " + self.RmatrixFilename +
                             " for " +
                             instrument + " " + satellite +
                             " is ATOV ?: " + str(isATOV))
            except (KeyError):
                raise ValueError("wrong instrument" + str(instrument))

            self.Rmatrix = Rmatrix()
            logging.info("Read the Rmatrix")
            self.Rmatrix.read(self.RmatrixFilename)
            logging.info("R matrix successfully read")

            try:
                if len(self.Rmatrix.Rmatrix.keys()) == 1:
                    self.Rmat = self.Rmatrix.Rmatrix.values()[0]["matrix"]
                else:
                    if isATOV:

                        if self.instrument == "amsua":
                            self.Rmat = self.Rmatrix.Rmatrix[
                                satellite + " " + "ATOVS"]["matrix"][
                                                                20:35, 20:35]
                        else:
                            if (
                                    self.instrument == "amsub" or
                                    self.instrument == "MHS" or
                                    self.instrument == "mhs"):
                                self.Rmat = self.Rmatrix.Rmatrix[
                                    satellite + " " + "ATOVS"]["matrix"][
                                    35:40, 35:40]
                                print self.Rmat
                            else:
                                if self.instrument == "hirs":
                                    # HIRS : 1 vis channel, 19 IR channels :
                                    # take only the IR channels !
                                    # in this case we start at 1 because the
                                    # first channel is vis !!!!
                                    self.Rmat = self.Rmatrix.Rmatrix[
                                        satellite + " " + "ATOVS"]["matrix"][
                                        1:20, 1:20]
                                else:
                                    raise ValueError(
                                        "wrong satellite and instrument" +
                                        str(satellite) + " " + str(instrument))
                    else:
                        raise ValueError(
                            "wrong satellite and instrument " +
                            str(satellite) +
                            " " + str(instrument))
            except (KeyError):

                raise ValueError("wrong satellite and instrument " +
                                 str(satellite) + " " + str(instrument))
        # keep only selected channels

        if self.channel_list is None:
            logging.info("no channel list take all of the matrix")
            self.channel_list = range(0, self.Rmat.shape[0])
        else:
            # this channel list already begin with 0
            foo = self.Rmat[self.channel_list, :]
            bar = foo[:, self.channel_list]
            self.Rmat = bar

    def plotRmat(self):
        plt.imshow(self.Rmat, origin="lower", interpolation="nearest")
        plt.title("Rmatrix for " + self.instrument + " " + self.satellite)
        plt.colorbar()
        plt.show()

    def setFactorB(self, factor):
        self.factorB = factor

    def setFactorR(self, factor):
        self.factorR = factor

    def setMaxNoise(self, maxNoise):
        self.MaxNoise = maxNoise

    def stepRetrieve1d(self):
        """ perform a step of retrieval
          the retrieved vector becomes the new background profile"""
        if self.Xr is not None:
            if self.changeTheBackground:
                # make a new profile object with the retrieved vector and the
                # former background profile
                newBgProfile = vector2profile(self.Xr, self.pBg.myProfile)
                self.pBg.myProfile = newBgProfile
                self.pBg.myProfile = self.retrievedProfiles[-1]
                logging.info("new profile background initialised " +
                             self.pBg.profileFileName)
                self.retrieve1d()
            else:
                # we do not change the Background
                self.retrieve1d()

    def set1dvarRequiredOptionsAndValues(self, project):
        """ set the required option and valuesfor the 1dvar algorithm """
        """ this options are CFRACTION=0 """
        """ Nadir condition """
        logging.info(
            "set 1dvar required options and values (CFRACTION=0, no gas,"
            " calculation at nadir etc) ")
        project.myOption["SWITCHRAD"] = True
        # No gas
        project.myOption["CO2_DATA"] = False
        project.myOption["CH4_DATA"] = False
        project.myOption["OZONE_DATA"] = False
        project.myOption["CO_DATA"] = False
        project.myOption["N2O_DATA"] = False
        # no cloud
        project.myOption["ADDCLOUDS"] = False
        project.myProfile["CFRACTION"] = 0
        # no aerosols
        project.myOption["ADDAEROSL"] = False
        # no solar
        project.myOption["ADDSOLAR"] = False

        # calculation at nadir
        project.myProfile["ZENANGLE"] = 0
        project.myProfile["SUNZENANGLE"] = 0

    def assumptionOnSurfaceParameters(self, profileTrue, profileBg):
        """ restrictions the surface parameters of the Bg are applied
             on the True except for T skin and T 2m  """
        profileBg["SKIN"]['SURFTYPE'] = 1  # (0=Land, 1=Sea, 2=sea-ice)
        for item in ["LATITUDE", "LONGITUDE", "ELEVATION", "AZANGLE", "BE",
                     "COSBK", "SNOW_FRAC", "SOIL_MOISTURE"]:
            profileTrue[item] = profileBg[item]
        for item in ["Q", "O", "P", "U", "V", "WFETC"]:
            profileTrue["S2M"][item] = profileBg["S2M"][item]
        for item in ["SURFTYPE", "WATERTYPE", "SALINITY", "FASTEM"]:
            profileTrue["SKIN"][item] = profileBg["SKIN"][item]

    def retrieve1d(self):
        """ perform the basic retrieval algorithm """
        """ must put switchrad a true et cfraction a 0 """
        # load coefficients
        self.step_counter = self.step_counter + 1
        self.pTrue.myCoeffs = self.pBg.myCoeffs
        self.pTrue.myOption = self.pBg.myOption
        # we are making assumption on surface parameters some must be the same
        self.assumptionOnSurfaceParameters(
            self.pTrue.myProfile, self.pBg.myProfile)

        # at this point the coefficient must be loaded
        if not self.pBg.myCoeffs.loadCoeffs:
            logging.warning(
                "strange : coefficients would have been "
                "loaded at this point ...")
            err = self.pBg.loadCoefficients()
            if err != 0:
                raise RuntimeError("Cannot load coefficient in r1dvar project")

        if len(self.channel_list) != self.pTrue.myCoeffs.nchannels:
            logging.warning("wrong channel list : take all")
            self.channel_list = range(1, self.pTrue.myCoeffs.nchannels + 1)
            self.readRmatrix(self.satellite, self.instrument)

        nbChannels = self.pTrue.myCoeffs.nchannels
        logging.info("nbChannels: {}".format(nbChannels))

        logging.debug("coefficients file: " +
                      self.pBg.myCoeffs.fileName["standard"])
        if self.step_counter > 1:
            # test if we have a different number of channels
            if len(self.Y_Xt) == 1:
                previousNumberOfChannels = 1
            else:
                previousNumberOfChannels = self.Y_Xt.shape[0]
            notSameNumberOfChannels = (
                self.pBg.myCoeffs.nchannels != previousNumberOfChannels)
            if notSameNumberOfChannels:
                logging.info("not the same number of channels")

        # Run direct on True

        logging.info(">>>>>>>>> Run direct on True")
        # must put switchrad a true et cfraction a 0 and other things on
        # options
        self.set1dvarRequiredOptionsAndValues(self.pTrue)
        err = self.pTrue.runDirect()
        if (err != 0):
            raise RuntimeError("Run Direct failed on true")
        # put radiance to 0 in the radiance file for plotting
        self.changeRadianceFile(self.pTrue.radianceFileName)
        # 1) Y(Xt) brightness temperatures from Xt
        YXt = rttov.radiance.Radiance()
        logging.info(">>>>>>>>>>>reading BT from : ")
        YXt.read(self.pTrue.radianceFileName)
        self.Y_Xt = YXt["BT_CLEAR"]  # No clouds so BT_CLEAR=BT
        Xt = self.pTrue.myProfile.to1DvarVect()
        self.Xt = Xt

        # Run direct on Background
        logging.info(">>>>>>>> Run direct on Background")
        # must put switchrad a true et cfraction a 0 and other things on
        # options
        self.set1dvarRequiredOptionsAndValues(self.pBg)
        err = self.pBg.runDirect()
        if (err != 0):
            raise RuntimeError("Run Direct failed on background")
        # put radiance to 0 in the radiance file for plotting
        self.changeRadianceFile(self.pBg.radianceFileName)

        logging.info(">>>>>>>> Run K on Background")
        err = self.pBg.runK()
        if (err != 0):
            raise RuntimeError("Run K failed on background")
        else:
            logging.info("run K on background OK")

        # 2) Y((Xb) brightness temperarures from Xb
        YXb = rttov.radiance.Radiance()
        logging.info(">>>>>>>>>>reading BT " + self.pBg.radianceFileName)
        YXb.read(self.pBg.radianceFileName)
        Y_Xb = YXb["BT_CLEAR"]

        # 3) compute jacobian matrix K(Xb) and transpose KT(Xb)
        KXb = r1dvarObjects.R1dvarKmatrix()
        logging.info("reading kmatrix " + self.pBg.KMatrixFileName)
        KXb.read(self.pBg.KMatrixFileName)

        logging.info(">>>>>>>KmatrixVectorAndTranspose ")
        # extract the jacobian matrix from the KMatrix Object
        K_Xb = KXb.toKmatrix1dvar()
        # transpose this matrix
        KT_Xb = K_Xb.T

        # 4) apply scaling factor to background errors B --> fb x B
        # the Bmatrix has been read in init and has been put in self.Bmat
        myBmat = self.Bmat * self.factorB

        # 5) apply scaling factor to observation errors R --> fr x R
        # The R matrix has been read in init and has been put in self.Rmat

        myRmat = self.Rmat * self.factorR

        # 6) compute add random noise to Y(Xt)

#
        # Noise in BT = +/- MaxNoise degree
        logging.info("Max Noise : " + str(self.MaxNoise))
        self.N = self.MaxNoise * (np.random.rand(nbChannels) * 2 - 1)
        self.Y_Xt = self.Y_Xt + self.N

        # save the file with True BT+Noise in order to be able to plot it after
        new_name = self.pTrue.radianceFileName.replace(".h5", "_Noise.h5")
        self.changeRadianceFile(self.pTrue.radianceFileName, self.N, new_name)

        # 7) compute linear 1DVAR weights W where
        # W = B . KT . [ K .B . KT + R ](-1)
        # W =  fact1 . [ foo1 .KT  + R ](-1)
        # W =  fact1 .  [  foo2   +  R ](-1)
        # W =  fact1 .  [       foo3   ](-1)
        # W =  fact1 . Matrice_inverse
        fact1 = myBmat.dot(KT_Xb)
        foo1 = K_Xb.dot(myBmat)
        foo2 = foo1.dot(KT_Xb)
        foo3 = foo2 + myRmat
        # matrix inversion
        logging.info("matrix inversion")
        Matrice_inverse = np.linalg.inv(foo3)
        logging.info("compute the linear 1DVAR weight ")
        W = fact1.dot(Matrice_inverse)

        # 8) compute linear 1Dvar retrieved profile (Xr)
        Xb = self.pBg.myProfile.to1DvarVect()
        # we have already added N to Y_Xt earlier
        YtminusYb = self.Y_Xt - Y_Xb

        Xr = Xb + W.dot(YtminusYb)
        # if Q <0 alors Q=0
        for k in range(Xr.shape[0] - 29 - 3, (Xr.shape[0] - 3)):
            if Xr[k] <= 0:
                # if find something negative for Q : stick to the background
                # ....
                Xr[k] = Xb[k]
        k = Xr.shape[0] - 2
        if Xr[k] <= 0:
            # if find something negative for Q : stick to the background ....
            Xr[k] = Xb[k]

        self.Xr = Xr
        self.Y_Xb = Y_Xb
        self.Xb = Xb
        # append the retrieved Vector and the retrieved profile to list
        self.retrievedVectors.append(Xr)
        retrievedProfile = vector2profile(self.Xr, self.pBg.myProfile)
        self.retrievedProfiles.append(retrievedProfile)
        return Xr

    def changeRadianceFile(self, radianceFileName, Noise=None, NewName=None):
        """   change the radiance file set radiance to 0 add Noise to BT) """
        if NewName:
            shutil.copyfile(radianceFileName, NewName)
            rf = h5py.File(radianceFileName, 'r+')
        else:
            rf = h5py.File(radianceFileName, 'r+')
        group = rf["RADIANCE"]
        bt = rf["RADIANCE/BT"][:]
        bt_clear = rf["RADIANCE/BT_CLEAR"][:]
        rad_clear = rf["RADIANCE/CLEAR"]
        a = rad_clear[:]

        a_new = np.zeros_like(a)
        del rf["RADIANCE/CLEAR"]
        del rf["RADIANCE/CLOUDY"]
        del rf["RADIANCE/TOTAL"]
        del rf["RADIANCE/OVERCAST"]
        group.create_dataset("CLEAR", a_new.shape, data=a_new)
        group.create_dataset("CLOUDY", a_new.shape, data=a_new)
        group.create_dataset("TOTAL", a_new.shape, data=a_new)
        group.create_dataset("OVERCAST", a_new.shape, data=a_new)
        if Noise is not None:
            bt_new = bt + Noise
            bt_clear_new = bt_clear + Noise
            del rf["RADIANCE/BT"]
            del rf["RADIANCE/BT_CLEAR"]
            group.create_dataset("BT", bt_new.shape, data=bt_new)
            group.create_dataset("BT_CLEAR", bt_new.shape, data=bt_clear_new)
        rf.close()

    def plot(self):
        colors = {0: "red",
                  1: "orange",
                  2: "yellow",
                  3: "magenta",
                  4: "brown"
                  }
        ax = plt.subplot(2, 2, 1)
        Xb0 = self.BgInitialProfile.to1DvarVect()
        Y = self.pTrue.myProfile["P"]
        pression = Y
        print ("pression")
        for k in range(0, pression.shape[0]):
            print ("%.2f" % pression[k])
        ax.set_ylim((pression[-1] + 30, pression[0]))
        ax.set_yscale("log")
        ax.plot(self.Xt[:54], Y, color="black", label="T true")
        ax.plot(Xb0[:54], Y, color="blue", label="T background")
        for k in range(len(self.retrievedVectors)):
            ax.plot(self.retrievedVectors[k][
                    :54], Y, color="red", label="T retrieved step " + str(k))
        ax.plot(self.Xt[85], Y[-1] + 20, color="black",
                marker="*", label="T skin True")
        ax.plot(Xb0[85], Y[-1] + 20, color="blue",
                marker="*", label="T skin background")
        for k in range(len(self.retrievedVectors)):
            ax.plot(self.retrievedVectors[k][85],
                    Y[-1] + 20, color="red", marker="*",
                    label="T skin retrieved step" + str(k))
        ax.legend(prop={'size': 10}, shadow=True, fancybox=True, loc='best')
        ax = plt.subplot(2, 2, 2)

        pression = Y[:29]
        print(pression.shape)

        ax.set_ylim((pression[-1], pression[0]))
        ax.set_yscale("log")

        ax.plot(self.Xt[54:83], pression, color="black", label="Q true")
        ax.plot(Xb0[54:83], pression, color="blue", label="Q background")
        print "q true"
        print self.Xt[54:83]
        print "Q initial"
        print Xb0[54:83]
        for k in range(len(self.retrievedVectors)):
            ax.plot(self.retrievedVectors[k][54:83], pression, color=colors[
                    k], label="Q retrieved step " + str(k))
            print "Q retrieved k", k
            print self.retrievedVectors[k][54:83]
        ax.legend(prop={'size': 10}, shadow=True, fancybox=True, loc='best')

        ax = plt.subplot(2, 2, 3)
        if len(self.Y_Xb.shape) == 0:
            nbchannels = 1
        else:
            nbchannels = self.Y_Xb.shape[0]

        channels = np.arange(1, nbchannels + 1)

        print(self.Y_Xb.shape)
        print(channels.shape)
        ax.plot(channels, self.Y_Xt, color="black", label="BT True")
        ax.plot(channels, self.Y_Xb, color="blue", label="BT Background")
        ax.plot(channels, self.Y_Xt + self.N,
                color="yellow", label="BT True+Noise")
        ax.legend(prop={'size': 10}, shadow=True, fancybox=True, loc='best')
        plt.title(self.satellite + " " + self.instrument)
        plt.show()

class Bmatrix(object):
    """ Class for dealing with the background error covariance
        matrices used by the 1DVar sheme """

    def __init__(self):
        self.matrices = {}
        self.header = {}
        self.dimension = {}

    def _read(self, f, dim=86):

        data = np.zeros((dim * dim), dtype=np.float64)
        indice = 0
        while indice < dim * dim:
            lin = f.readline().split()
            for value_str in lin:
                data[indice] = np.float64(value_str)
                indice = indice + 1
        return data.reshape(dim, dim)

    def read_matrices(self, filename):
        f = open(filename)
        for j in (1, 2, 3, 4):
            self.header[j] = []
            for i in (0, 1, 3):
                k = f.readline()
                self.header[j].append(k)
            self.dimension[j] = int(k)
            self.matrices[j] = self._read(f, dim=self.dimension[j])

            # convert lnq g/kg in ppmv for Q
            dim = self.dimension[j]
            for k in range(dim - 29 - 3, dim - 3):
                for l in range(dim - 29 - 3, dim - 3):
                    self.matrices[j][k, l] = q_mixration_to_ppmv * \
                        np.exp(self.matrices[j][k, l]) * 0.001
            self.matrices[j][dim - 2, dim - 2] = q_mixration_to_ppmv * \
                np.exp(self.matrices[j][dim - 2, dim - 2]) * 0.001

        f.close()

    def plot(self, j):
        print "Bmatrice :", j
        print self.matrices[j]

        plt.imshow(self.matrices[j], origin="lower", interpolation="nearest")
        plt.title(self.header[j][0])
        plt.colorbar()
        plt.show()

class RmatrixBase(object):
    """ Class for dealing with mesurement error covariance matrix
        used by the 1DVar scheme """

    def __init__(self):
        pass

    def read(self, filename):
        pass

class Rmatrix(RmatrixBase):

    def read(self, filename):
        """ Rmatrix fila structure is explained  here : """
        """ https://nwpsaf.eu/deliverables/nwpsaf_1dvar/
            nwpsaf-mo-ud-032_NWPSAF_1DVar_Manual.html#aux """
        """ read the entire file : populate a dictionary named Rmatrx with :"""
        """ nchannels : number of channels"""
        """ nband : number of band """
        """ inverse : """
        """ type :"""
        """ data : the data"""
        """ finaly with the __mkMatrix method is called to make the Matrix """
        """ dictionary which will contain numpy arrays """

        self.Rmatrix = {}
        logging.debug("open " + filename)
        f = open(filename)
        cond = True
        while cond:
            l = f.readline()
            if len(l) == 0:
                cond = False
                continue
            satint = l.replace("\n", "")
            logging.debug("read" + str(satint))
            self.Rmatrix[satint] = {}
            line = f.readline()
            self.Rmatrix[satint]["type"], self.Rmatrix[satint][
                "nchannel"], self.Rmatrix[
                    satint]["nband"], self.Rmatrix[satint]["inverse"] = [
                        int(x) for x in line.split()]
            logging.debug(str(self.Rmatrix[satint]["type"]) +
                          str(self.Rmatrix[satint]["nchannel"]) +
                          str(self.Rmatrix[satint]["nband"]) +
                          str(self.Rmatrix[satint]["inverse"]))
            if self.Rmatrix[satint]["type"] == 2:
                self.Rmatrix[satint]["channel"] = []
                self.Rmatrix[satint]["data"] = {}
                while (len(self.Rmatrix[satint]["channel"]) <
                       self.Rmatrix[satint]["nchannel"]):
                    l = f.readline().split()
                    for x in l:
                        self.Rmatrix[satint]["channel"].append(int(x))
                for band in range(self.Rmatrix[satint]["nband"]):
                    self.Rmatrix[satint]["data"][band] = []
                    while (len(self.Rmatrix[satint]["data"][band]) <
                           self.Rmatrix[satint]["nchannel"]):
                        l = f.readline().split()
                        for x in l:
                            self.Rmatrix[satint]["data"][band].append(float(x))
                logging.debug("end")
            else:
                logging.debug('cannot read matrix')
                cond = False
        self._mkMatrix()

    def _mkMatrix(self):
        """ create the matrix from the data read """
        logging.debug("create the matrix")
        for satint in self.Rmatrix.keys():
            logging.debug("create the matrix for satellite" + satint +
                          str(self.Rmatrix[satint]["type"]) +
                          str(self.Rmatrix[satint]["nband"]))
            self.Rmatrix[satint]["matrix"] = np.zeros(
                (self.Rmatrix[satint]["nchannel"],
                 self.Rmatrix[satint]["nchannel"]), dtype=float)
            if (
                    self.Rmatrix[satint]["type"] == 2 and
                    self.Rmatrix[satint]["nband"] == 1):
                for i in range(self.Rmatrix[satint]["nchannel"]):
                    self.Rmatrix[satint]["matrix"][
                        i, i] = self.Rmatrix[satint]["data"][0][i]

            if (
                    self.Rmatrix[satint]["type"] == 2 and
                    self.Rmatrix[satint]["nband"] != 1):
                # diagonal case
                for i in range(self.Rmatrix[satint]["nchannel"]):
                    self.Rmatrix[satint]["matrix"][
                        i, i] = self.Rmatrix[satint]["data"][0][i]

                for band in range(1, self.Rmatrix[satint]["nband"]):
                    logging.debug("band" + str(band))

                    for i in range(0, self.Rmatrix[satint]["nchannel"] - band):
                        self.Rmatrix[satint]["matrix"][
                            i + band, i] = self.Rmatrix[satint]["data"][band][
                            i]
                        self.Rmatrix[satint]["matrix"][
                            i, i + band] = self.Rmatrix[satint]["data"][band][
                            i]

    def plot_matrix(self):
        for satint in self.Rmatrix.keys():
            plt.imshow(self.Rmatrix[satint]["matrix"],
                       origin="upper", interpolation="nearest")
            plt.title(satint)
            plt.colorbar()
            plt.show()

class Nwpsaf1dvarProfile(object):

    def __init__(self):
        self.data = {}

    def toRttovGuiProfile(self):
        """ convert Background to profile """

        self.profile = rmodel.project.pProfile()
        self.option = rmodel.project.pOption()
        self.option.default()
        self.profile.setDefaultAttributes()
        self.profile["NLEVELS"] = self.data["P"][::-1].shape[0]
        self.profile["LAYERS"] = self.profile["NLEVELS"] - 1
        self.profile["P"] = self.data["P"][::-1]
        self.profile["T"] = self.data["T"][::-1]
        self.profile["Q"] = self.data["Q"][::-1]
        self.profile["O3"] = self.data["O3"][::-1]
        self.profile["S2M"]["T"] = self.data["Surface Temperature (K)"]
        self.profile["SKIN"]["T"] = self.data["Skin Temperature (K)"]
        for k, v in self.data.items():
            if k.startswith("Surface Humidity"):
                self.profile["S2M"]["Q"] = v
        self.profile["S2M"]["U"] = self.data["10m U-Wind (m/s)"]
        self.profile["S2M"]["V"] = self.data["10m U-Wind (m/s)"]
        self.profile["S2M"]["P"] = self.data["Surface Pressure (hPa)"]
        self.profile.setDefaultProfileAsciiInput()
        return self.profile

class RetrievedProfile(object):
    """ class for dealing with 1D var ascii Retrieved_Profile files
        from the NWPDAF 1DVAR software """
    """ read them and transform it in profile object """

    def __init__(self):
        self.retrieved = Nwpsaf1dvarProfile()
        self.background = Nwpsaf1dvarProfile()
        self.header = {}

    def read(self, filename):
        data = np.genfromtxt(filename, skip_header=3, skip_footer=8)
        self.retrieved.data["P"] = data[:, 0]
        self.retrieved.data["T"] = data[:, 1]
        self.retrieved.data["Q"] = data[:, 2]
        self.retrieved.data["O3"] = data[:, 3]
        self.background.data["P"] = data[:, 0]
        self.background.data["T"] = data[:, 4]
        self.background.data["Q"] = data[:, 5]
        self.background.data["O3"] = data[:, 5]

    def plot(self):
        colors = {0: "red",
                  1: "orange",
                  2: "yellow",
                  3: "magenta",
                  4: "brown"
                  }
        ax = plt.subplot(2, 2, 1)
        Xb = self.background
        Xr = self.retrieved

        Y = self.background.data["P"][:]
        pression = Y
        print ("pression")
        for k in range(0, pression.shape[0]):
            print ("%.2f" % pression[k])
        ax.set_ylim((pression[0] + 30, pression[-1]))
        ax.set_yscale("log")
        ax.plot(Xb.data["T"], Y, color="blue", label="T background")
        ax.plot(Xr.data["T"], Y, color="red", label="T retrieved")

        ax.legend(prop={'size': 10}, shadow=True, fancybox=True, loc='best')
        ax = plt.subplot(2, 2, 2)
        ax.set_ylim((pression[0] + 30, pression[-1]))

        ax.set_yscale("log")

        ax.plot(Xb.data["Q"], pression, color="blue", label="Q background")

        ax.plot(Xr.data["Q"], pression, color="red", label="Q retrieved  ")

        ax.legend(prop={'size': 10}, shadow=True, fancybox=True, loc='best')

        plt.show()

class Background(object):
    """ class for dealing with 1D var ascii Backround files from the
        NWPDAF 1DVAR software """
    """ read them and transform it in profile object """

    def __init__(self):
        self.prof = Nwpsaf1dvarProfile()
        self.data = {}
        self.header = {}

    def read(self, filename):
        data = np.genfromtxt(filename, skip_header=16, skip_footer=6)

        self.data["P"] = data[:, 0]
        self.data["T"] = data[:, 1]
        self.data["Q"] = data[:, 2]
        self.data["O3"] = data[:, 3]
        self.nlevels = data.shape[0]

        labels = np.genfromtxt(
            filename, usecols=0, dtype=str, skip_header=16 + data.shape[0],
            delimiter=":")
        raw_data = np.genfromtxt(
            filename, skip_header=16 + data.shape[0], delimiter=":")[:, 1:]

        self.data_surf = {label: row for label,
                          row in zip(labels, raw_data[:, 0])}
        for key in ("P", "Q", "T", "O3"):
            self.prof.data[key] = self.data[key]
        self.prof.data["Surface Temperature (K)"] = self.data_surf[
            "Surface Temperature (K)"]
        self.prof.data["Skin Temperature (K)"] = self.data_surf[
            "Skin Temperature (K)"]
        self.prof.data["10m U-Wind (m/s)"] = self.data_surf["10m U-Wind (m/s)"]
        self.prof.data["10m U-Wind (m/s)"] = self.data_surf["10m U-Wind (m/s)"]
        self.prof.data["Surface Pressure (hPa)"] = self.data_surf[
            "Surface Pressure (hPa)"]

    def toProfile(self):
        return self.prof.toRttovGuiProfile()

    def print_data(self):
        print ("P")
        print (self.data["P"])
        print ("T")
        print (self.data["T"])
        print ("Q")
        print (self.data["Q"])
        print ("O3")
        print (self.data["P"])
        print (self.data_surf)

if __name__ == '__main__':

    rmatrixfilename = "./data/IASI_COEFFS_DIR/Rmatrix_orig"

    print ("rmatrixfilename")
    R = Rmatrix()
    R.read(rmatrixfilename)
    R.plot_matrix()