"""Tables module."""

import numpy as np
from astropy.io import fits

class Tables:
    r"""Class for handling the tables of the emulator.

    It handles both the training and validation sets, in the form of tables
    containing the different contributions to :math:`P_\mathrm{gg}(k,\mu)`.
    Provides routine to convert the raw tables into their transformed versions,
    which span a shorter dynamical range (to increase the accuracy of the
    emulator), to train the emulator, and to transform back the emulator
    predictions in the original metric.
    """

    def __init__(self, params, real_space=False):
        r"""Class constructor.

        Parameters
        ----------
        params: list
            List of parameters used to train the emulator.
        real_space: bool
            Flag to determine if the tables are in real or redshift space.
        """
        self.params = params
        self.n_params = len(params)
        self.param_ranges = None
        self.model = None
        self.model_transformed = None

        self.ells = [0, 2, 4] if not real_space else [0]

        self.n_diagrams = 19
        self.names_diagrams = ['P0L_b1b1', 'PNL_b1', 'PNL_id', 'P1L_b1b1',
                               'P1L_b1b2', 'P1L_b1g2', 'P1L_b1g21', 'P1L_b2b2',
                               'P1L_b2g2', 'P1L_g2g2', 'P1L_b2', 'P1L_g2',
                               'P1L_g21', 'Pctr_c0', 'Pctr_c2', 'Pctr_c4',
                               'Pctr_b1b1cnlo', 'Pctr_b1cnlo', 'Pctr_cnlo']

    def set_param_ranges(self, ranges):
        r"""Set ranges for the parameters of the emulator.

        Sets the internal class attribute defining the prior for the
        parameters of the emulator.

        Parameters
        ----------
        ranges: dict
            Dictionary containing the parameter ranges. Keys correspond to the
            name of the parameters, while values are list with two entries,
            which correspond to the minimum and maximum value of the parameter.
        """
        if self.param_ranges is None:
            self.param_ranges = {}
        for p in self.params:
            self.param_ranges[p] = ranges[p]

    def assign_samples(self, samples_hdu):
        r"""Read parameter sample from HDU object.

        Reads parameter sample from a Header Data Unit object, and stores it
        as class attribute. If the sample is meant for validation, determines
        the parameters from the header of the HDU object, otherwise uses the
        parameters defined as class attributes.

        Parameters
        ----------
        samples_hdu: astropy.io.fits.BinTableHDU
            Header Data Unit containing the parameter sample.
        """
        self.n_samples = samples_hdu.header['NAXIS2']
        self.samples = np.zeros([self.n_samples, self.n_params])
        for i, p in enumerate(self.params):
            self.samples[:, i] = samples_hdu.data[p]

    def assign_table(self, table_hdu, nk, nkloop):
        r"""Read model table from HDU object.

        Reads a model table from a Header Data Unit object, and stores it as
        class attribute. If the sample is a training one, additionally calls
        the routine to convert the model table into its transformed version.

        Parameters
        ----------
        table_hdu: astropy.io.fits.BinTableHDU
            Header Data Unit containing the model table.
        nk: int
            Number of :math:`k` bins of the table.
        nkloop: int
            Number of :math:`k` bins of the table corresponding to the loop
            predictions.
        """
        self.nk = nk
        self.nkloop = nkloop

        if table_hdu.header['EXTNAME'] in ['MODEL_SHAPE', 'MODEL_LINEAR']:
            if self.model is None:
                self.model = {}
            for TYPE in [table_hdu.header['TTYPE{}'.format(i+1)]
                         for i in range(table_hdu.header['TFIELDS'])]:
                self.model[TYPE] = table_hdu.data[TYPE]
                if self.model[TYPE].ndim == 1:
                    self.model[TYPE] = self.model[TYPE][:, None]
            self.transform_emulator_data(data_type=list(self.model.keys()))
        elif 'MODEL_FULL' == table_hdu.header['EXTNAME']:
            if self.model is None:
                self.model = {}
            self.model['PL'] = table_hdu.data['PL']
            for ell in self.ells:
                for diagram in self.names_diagrams:
                    diagram_full = '{}_ell{}'.format(diagram, ell)
                    if diagram == 'PNL_b1':
                        # combine tree-level, one-loop
                        # and IR k^2 correction terms
                        self.model[diagram_full] = (
                            table_hdu.data['P0L_b1_ell{}'.format(ell)] +
                            table_hdu.data['P1L_b1_ell{}'.format(ell)] +
                            table_hdu.data['Pk2corr_b1_ell{}'.format(ell)])
                    elif diagram == 'PNL_id':
                        # combine tree-level, one-loop
                        # and IR k^2 correction terms
                        self.model[diagram_full] = (
                            table_hdu.data['P0L_id_ell{}'.format(ell)] +
                            table_hdu.data['P1L_id_ell{}'.format(ell)] +
                            table_hdu.data['Pk2corr_id_ell{}'.format(ell)])
                    elif diagram == 'P1L_b1b1':
                        # combine one-loop and IR k^2 correction terms
                        self.model[diagram_full] = (
                            table_hdu.data['P1L_b1b1_ell{}'.format(ell)] +
                            table_hdu.data['Pk2corr_b1b1_ell{}'.format(ell)])
                    else:
                        self.model[diagram_full] = table_hdu.data[diagram_full]
            self.transform_emulator_data()
        else:
            raise KeyError('HDU table does not contain valid identifiers.')

    def get_flip_and_offset(self, table):
        r"""Compute flip and offset to rescale an input model table.

        Parameters
        ----------
        table: numpy.ndarray
            Model table containing a given term as a function of :math:`k`
            (first index) and the different parameter sample (second index).

        Returns
        -------
        flip: list
            Flip value for each :math:`k` entry.
        max_offset: list
            Maximum offset for each :math:`k` entry.
        """
        flip = []
        offset_list = []
        for i in range(table.shape[0]):
            idmax = np.abs(table[i, :]).argmax()
            flip.append(np.sign(table[i, idmax]))
            offset_list.append(np.abs(np.amin(flip[-1]*table[i, :])))
        if len(offset_list) == 0:
            max_offset = 0.0
        else:
            max_offset = max(offset_list)*1.1
        return flip, max_offset

    def transform(self, table, data_type):
        r"""Rescale the dynamical range of a table for the emulation.

        Applies flip, offset, and rescales the model table. This step is
        performed by first taking the logarithm of the model table, and then
        subtracting its mean and dividing by its standard deviation.

        Parameters
        ----------
        table: numpy.ndarray
            Model table containing a given term as a function of :math:`k`
            (first index) and the different parameter sample (second index).
        data_type: str
            Type of the table that is passed as input. This is used as a
            keyword for the dictionary where the values of flip, offset,
            mean and standard deviation are stored.

        Returns
        -------
        resc_table: numpy.ndarray
            Rescaled table.
        """
        if data_type not in ['PL', 'PNW', 's12', 'sv']:
            self.flip[data_type], self.offset[data_type] = \
                self.get_flip_and_offset(table.T)
        else:
            self.flip[data_type] = np.ones(table.shape[1])
            self.offset[data_type] = 0.0
        temp = np.log10(self.flip[data_type]*table+self.offset[data_type])
        self.mean[data_type] = np.mean(temp, axis=0)
        self.std[data_type] = np.std(temp, axis=0)
        return (temp-self.mean[data_type])/self.std[data_type]

    def transform_inv(self, table, data_type):
        r"""Rescale back the dynamical range of a table.

        Performs the inverse operations defined in the **transform** method.
        Needed to obtain predictions from the emulator in the original metric.

        Parameters
        ----------
        table: numpy.ndarray
            Table containing a given term as a function of :math:`k` (first
            index) and the different parameter sample (second index).
        data_type: str
            Type of the table that is passed as input. This is used as a
            keyword for the dictionary where the values of flip, offset,
            mean and standard deviation are stored.

        Returns
        -------
        resc_table: numpy.ndarray
            Rescaled table.
        """
        return (10**(table*self.std[data_type] + self.mean[data_type]) -
                self.offset[data_type]) * self.flip[data_type]

    def transform_emulator_data(self, data_type=None):
        r"""Rescale class attribute tables.

        Transforms the class attribute model tables from the original metric
        into the emulator one.

        Parameters
        ----------
        data_type: str or list, optional
            Types of the tables which have to be rescaled. If *None*, trasform
            the full set of model tables. Defaults to **None**.
        """
        if data_type is not None:
            data_type = ([data_type] if not isinstance(data_type, list)
                         else data_type)
            for dt in data_type:
                if self.model_transformed is None:
                    self.model_transformed = {}
                    self.mean = {}
                    self.std = {}
                    self.flip = {}
                    self.offset = {}
                self.model_transformed[dt] = self.transform(self.model[dt],
                                                            data_type=dt)
        else:
            self.model_transformed = {}
            self.mean = {}
            self.std = {}
            self.flip = {}
            self.offset = {}
            for ell in self.ells:
                temp = np.zeros([self.n_samples, 9*self.nk + 10*self.nkloop])
                cnt = 0
                # only include the cell_ell counterterm
                # (the others are negligible without AP)
                for diagram in ['P0L_b1b1', 'PNL_b1', 'PNL_id',
                                'Pctr_c0', 'Pctr_c2', 'Pctr_c4',
                                'Pctr_b1b1cnlo', 'Pctr_b1cnlo', 'Pctr_cnlo']:
                    diagram_full = '{}_ell{}'.format(diagram, ell)
                    temp[:, cnt*self.nk:(cnt+1)*self.nk] = \
                        self.model[diagram_full]/self.model['PL']
                    cnt += 1
                cnt = 0
                for diagram in ['P1L_b1b1', 'P1L_b1b2', 'P1L_b1g2',
                                'P1L_b1g21', 'P1L_b2b2', 'P1L_b2g2',
                                'P1L_g2g2', 'P1L_b2', 'P1L_g2', 'P1L_g21']:
                    diagram_full = '{}_ell{}'.format(diagram, ell)
                    temp[:,
                         9*self.nk + cnt*self.nkloop:9*self.nk +
                         (cnt+1)*self.nkloop] = (
                            self.model[diagram_full][:,
                                                     (self.nk-self.nkloop):] /
                            self.model['PL'][:, (self.nk-self.nkloop):])
                    cnt += 1

                self.model_transformed[ell] = self.transform(temp,
                                                             data_type=ell)
