"""TRAM module: (T)ime (R)esolved (A)nalysis (M)odule.

For taking raw counts per second data from a Laser Ablation Inductively
Coupled Plasma Mass Spectrometry (LA-ICP-MS) experiment, choosing an
interval to be turned into a concentration, normalizing
that interval to an internal standard and outputting that value + other
metadata.
"""

from __future__ import annotations

import numpy as np
import numpy.typing as npt
import pandas as pd
from ..helpers import formatting


def _z_filter(signal: pd.Series, std_devs: float, window: int = 50) -> pd.Series:
    """Compute a rolling z-score filter for a given signal.

    Values whose z-score (relative to a rolling window) exceeds
    ``std_devs`` are replaced by the rolling mean.

    Parameters
    ----------
    signal : pandas.Series
        The timeseries signal to be filtered.
    std_devs : float
        Number of standard deviations away from the rolling mean to
        be considered an outlier.
    window : int, optional
        Number of points in the rolling average window, by default 50.

    Returns
    -------
    pandas.Series
        The despiked signal with outliers replaced by the rolling mean.
    """

    # rolling calculation setup
    roll = signal.rolling(window=window, min_periods=1, center=True)
    # rolling average
    avg = roll.mean()
    # rolling standard deviation
    std = roll.std(ddof=0)
    # z score
    z = (signal - avg) / std
    # boolean if a given z score is within the bounds set up by
    # the standard deviation limits
    m = z.between(-std_devs, std_devs)

    # original signal where z score is between bounds
    # mean of rolling window otherwise
    despiked = signal.where(m, avg)

    return despiked


class LaserTRAM:
    """Time resolved analysis of a single LA-ICP-MS spot.

    ``LaserTRAM`` contains all the information and methods related to
    reducing one individual spot analysis from raw counts-per-second
    data to internally normalised ratios. To be used in conjunction
    with :class:`lasertram.calc.LaserCalc`.

    Parameters
    ----------
    name : str
        Sample name, i.e. the value in the ``SampleLabel`` column of the
        LT-ready file.

    Examples
    --------
    >>> from lasertram import LaserTRAM, helpers
    >>> raw_data = helpers.preprocessing.load_test_rawdata()
    >>> spot = LaserTRAM(name="GSD-1G_-_1")
    >>> spot.get_data(raw_data.loc["GSD-1G_-_1", :])
    >>> spot.assign_int_std("29Si")
    >>> spot.assign_intervals(bkgd=(5, 10), keep=(25, 40))
    >>> spot.get_bkgd_data()
    >>> spot.subtract_bkgd()
    >>> spot.get_detection_limits()
    >>> spot.normalize_interval()
    >>> spot.make_output_report()
    """

    def __init__(self, name: str) -> None:
        """Initialise a LaserTRAM spot object.

        Parameters
        ----------
        name : str
            Sample name, i.e. the value in the ``SampleLabel`` column of
            the LT-ready file.
        """
        # all attributes in relative chronological order that they are created in
        # if everything is done correctly. These all will get rewritten throughout the
        # data processing pipeline but this allows us to see what all the potential attributes
        # are going to be from the beginning (PEP convention)

        # for the math involved please see:

        # name of the lasertram spot object
        self.name = name

        # boolean flag for whether or not the data have been
        # despiked
        self.despiked = False

        # list of elements that have been despiked. Also may be 'all'
        self.despiked_elements = None

        # data from a single spot to be processed. 2D pandas dataframe
        self.data = None

        # self.data but as a 2D numpy matrix. Equivalent to self.data.values
        self.data_matrix = None

        # list of analytes in the analysis
        self.analytes = None

        # datetime corresponding to the analysis
        self.timestamp = None

        # string representation internal standard analyte for the processing.
        # this is just the column header of the analyte chosen as the internal
        # standard e.g., "29Si"
        self.int_std = None

        # column number in self.data_matrix that denotes the internal standard analyte
        # data. Remember python starts counting at 0!
        self.int_std_loc = None

        # background interval start time
        self.bkgd_start = None

        # background interval stop time
        self.bkgd_stop = None

        # desired ablation interval start time
        self.int_start = None

        # desired ablation interval stop time
        self.int_stop = None

        # row in self.data corresponding to self.bkgd_start
        self.bkgd_start_idx = None

        # row in self.data corresponding to self.bkgd_stop
        self.bkgd_stop_idx = None

        # row in self.data corresponding to self.int_start
        self.int_start_idx = None

        # row in self.data corresponding to self.int_stop
        self.int_stop_idx = None

        # desired omitted region start time
        self.omit_start = None

        # desired omitted region stop time
        self.omit_stop = None

        # row in self.data corresponding to self.omit_start
        self.omit_start_idx = None
        # row in self.data corresponding to self.omit_stop
        self.omit_stop_idx = None

        #
        self.omitted_region = None

        # 1D array of median background values [self.bkgd_start - self.bkgd_stop)
        # that is len(analytes) in shape
        self.bkgd_data_median = None

        # 1D array of detection limits in counts per second
        # that is len(analytes) in shape
        self.detection_limits = None

        # 2D array of background corrected data over the self.int_start - self.int_stop
        # region
        self.bkgd_subtract_data = None

        # 2D array of background corrected data over the self.int_start - self.int_stop
        # region that is normalized to the internal standard
        self.bkgd_subtract_normal_data = None

        # 1D array of median background corrected normalized values over the self.int_start - self.int_stop
        # retion that is len(analytes) in shape
        self.bkgd_subtract_med = None

        # 1D array of 1 standard error of the mean values for each analyte over the
        # self.int_start - self.int_stop region
        self.bkgd_subtract_std_err = None

        #
        self.bkgd_subtract_std_err_rel = None

        # 1D pandas dataframe that contains many of the attributes created during the
        # LaserTRAM process:
        # |timestamp|Spot|despiked|omitted_region|bkgd_start|bkgd_stop|int_start|int_stop|norm|norm_cps|analyte vals and uncertainties -->|
        # |---------|----|--------|--------------|----------|---------|---------|--------|----|--------|----------------------------------|
        self.output_report = None

    def get_data(
        self, df: pd.DataFrame, time_units: str = "ms", verbose: bool = True
    ) -> None:
        """Assign raw counts-per-second data to the spot object.

        Parameters
        ----------
        df : pandas.DataFrame
            Raw data corresponding to the spot being processed, e.g.
            ``all_data.loc[spot, :]`` where *all_data* is the LT-ready
            DataFrame.
        time_units : str, optional
            Units for the ``Time`` column. Used to convert input time
            values to seconds. ``'ms'`` (default) or ``'s'``.
        verbose : bool, optional
            Whether to print status messages during data loading,
            by default ``True``.

        Examples
        --------
        >>> spot = LaserTRAM(name="GSD-1G_-_1")
        >>> spot.get_data(raw_data.loc["GSD-1G_-_1", :])
        """

        # TODO: get_data
        # - [x] add: check to make sure data are in right format else throw an error. do this by having a list of required columns and checking for them.

        # get data and set index to "SampleLabel" column
        self.data = df.reset_index()

        col_check, type_check = formatting.check_lt_input_format(self.data)
        if verbose:
            print(
                "checking LaserTRAM input data format for correct column headers and data types..."
            )

        if col_check is not None:
            raise ValueError(
                f"It looks like your input data are missing the following required columns: {col_check}. Please fix before continuing with processing."
            )
        if type_check is not None:
            raise TypeError(
                f"It looks like your input data have incorrect types in the following column indices: {type_check}. Please fix before continuing with processing."
            )

        if verbose:
            print("check complete...input data format looks good!")

        self.data = self.data.set_index("SampleLabel")

        # convert time units from ms --> s if applicable
        if time_units == "ms":
            self.data["Time"] = self.data["Time"] / 1000
        elif time_units == "s":
            pass

        # just numpy matrix for data
        self.data_matrix = self.data.iloc[:, 1:].to_numpy()

        # list of analytes in experiment
        self.analytes = self.data.loc[:, "Time":].columns.tolist()[1:]

        # need to add check for if this exists otherwise there is no timestamp attribute
        self.timestamp = str(self.data.loc[:, "timestamp"].unique()[0])
        self.timestamp = pd.to_datetime(self.timestamp)

    def assign_int_std(self, int_std: str) -> None:
        """Assign the internal standard analyte for normalisation.

        Parameters
        ----------
        int_std : str
            Column name for the internal standard analyte, e.g.
            ``"29Si"``.

        Examples
        --------
        >>> spot.assign_int_std("29Si")
        """

        # set the internal standard analyte
        self.int_std = int_std

        # get the internal standard array index
        self.int_std_loc = np.where(np.array(self.analytes) == self.int_std)[0][0]

    def assign_intervals(
        self,
        bkgd: tuple[float, float],
        keep: tuple[float, float],
        omit: tuple[float, float] | None = None,
    ) -> None:
        """Assign background, signal, and optional omission intervals.

        Parameters
        ----------
        bkgd : tuple of float
            ``(start, stop)`` pair of values (in seconds) for the
            background signal region.
        keep : tuple of float
            ``(start, stop)`` pair of values (in seconds) for the
            ablation signal region used in concentration calculations.
        omit : tuple of float or None, optional
            ``(start, stop)`` pair of values (in seconds) to be omitted
            from the ``keep`` interval, by default ``None``.

        Examples
        --------
        >>> spot.assign_intervals(bkgd=(5, 10), keep=(25, 40))

        With a region to omit:

        >>> spot.assign_intervals(bkgd=(5, 10), keep=(23, 40), omit=(30, 33))
        """

        # set background and interval times in s
        self.bkgd_start = bkgd[0]
        self.bkgd_stop = bkgd[1]
        self.int_start = keep[0]
        self.int_stop = keep[1]

        # equivalent background and interval times but as indices
        # in their respective arrays
        self.bkgd_start_idx = np.where(self.data["Time"] > self.bkgd_start)[0][0]
        self.bkgd_stop_idx = np.where(self.data["Time"] > self.bkgd_stop)[0][0]
        self.int_start_idx = np.where(self.data["Time"] > self.int_start)[0][0]
        self.int_stop_idx = np.where((self.data["Time"] > self.int_stop))[0][0]

        # boolean whether or not there is an omitted region
        self.omitted_region = False
        # if omission is true, set those start and stop times like above
        if omit:
            self.omit_start = omit[0]
            self.omit_stop = omit[1]
            self.omit_start_idx = (
                np.where(self.data["Time"] > self.omit_start)[0][0] - self.int_start_idx
            )
            self.omit_stop_idx = (
                np.where(self.data["Time"] > self.omit_stop)[0][0] - self.int_start_idx
            )

            self.omitted_region = True

    def get_bkgd_data(self) -> None:
        """Calculate median background values for all analytes.

        Uses the intervals assigned in :meth:`assign_intervals` to take
        the median value of all analytes within the background range.
        These are later subtracted from the ablation signal in
        :meth:`subtract_bkgd`.
        """
        # median background data values
        self.bkgd_data_median = np.median(
            self.data_matrix[self.bkgd_start_idx : self.bkgd_stop_idx, 1:], axis=0
        )

    def get_detection_limits(self) -> None:
        """Calculate detection limits in counts per second.

        Defined as the median background plus three standard deviations
        of the background signal for each analyte.
        """

        self.detection_limits = np.std(
            self.data_matrix[self.bkgd_start_idx : self.bkgd_stop_idx, 1:], axis=0
        ) * 3 + np.median(
            self.data_matrix[self.bkgd_start_idx : self.bkgd_stop_idx, 1:], axis=0
        )

    def subtract_bkgd(self) -> None:
        """Subtract the median background from the ablation signal.

        Removes the median background values calculated in
        :meth:`get_bkgd_data` from the signal in the *keep* interval
        established in :meth:`assign_intervals`.
        """
        self.bkgd_subtract_data = (
            self.data_matrix[self.int_start_idx : self.int_stop_idx, 1:]
            - self.bkgd_data_median
        )

    def normalize_interval(self) -> None:
        """Normalise the ablation interval to the internal standard.

        Divides all analytes in the *keep* portion of the signal by the
        internal standard analyte. Also calculates the median normalised
        value, its standard error of the mean, and relative standard
        error of the mean.
        """

        # set the detection limit thresholds to be checked against
        # with the interval data. This basically takes the detection limits
        threshold = self.detection_limits - self.bkgd_data_median

        # if there's an omitted region, remove it from the data to be further processed
        # for the chosen interval
        if self.omitted_region is True:
            self.bkgd_subtract_normal_data = np.delete(
                self.bkgd_subtract_data,
                np.arange(self.omit_start_idx, self.omit_stop_idx),
                axis=0,
            ) / np.delete(
                self.bkgd_subtract_data[:, self.int_std_loc][:, None],
                np.arange(self.omit_start_idx, self.omit_stop_idx),
                axis=0,
            )

        else:
            self.bkgd_subtract_normal_data = (
                self.bkgd_subtract_data
                / self.bkgd_subtract_data[:, self.int_std_loc][:, None]
            )

        # get background corrected and normalized median values for an interval
        self.bkgd_subtract_med = np.median(self.bkgd_subtract_normal_data, axis=0)
        self.bkgd_subtract_med[
            np.median(self.bkgd_subtract_data, axis=0) <= threshold
        ] = -9999
        self.bkgd_subtract_med[np.median(self.bkgd_subtract_data, axis=0) == 0] = -9999

        # standard error of the mean for the interval region
        self.bkgd_subtract_std_err = self.bkgd_subtract_normal_data.std(
            axis=0
        ) / np.sqrt(abs(self.int_stop_idx - self.int_start_idx))

        self.bkgd_subtract_std_err_rel = 100 * (
            self.bkgd_subtract_std_err / self.bkgd_subtract_med
        )

    def make_output_report(self) -> None:
        """Create an output report summarising the spot processing.

        Generates a single-row :class:`pandas.DataFrame` stored in
        ``self.output_report`` with the following columns:

        ``timestamp | Spot | despiked | omitted_region | bkgd_start |
        bkgd_stop | int_start | int_stop | norm | norm_cps |``
        followed by normalised analyte values and their standard errors.

        Examples
        --------
        >>> spot.make_output_report()
        >>> spot.output_report.head()
        """
        if self.despiked is True:
            despike_col = self.despiked_elements
        else:
            despike_col = "None"

        if self.omitted_region is True:
            omitted_col = (
                self.data["Time"].iloc[self.omit_start_idx + self.int_start_idx],
                self.data["Time"].iloc[self.omit_stop_idx + self.int_start_idx],
            )
        else:
            omitted_col = "None"

        spot_data = pd.DataFrame(
            [
                self.timestamp,
                self.name,
                despike_col,
                omitted_col,
                self.data["Time"].iloc[self.bkgd_start_idx],
                self.data["Time"].iloc[self.bkgd_stop_idx],
                self.data["Time"].iloc[self.int_start_idx],
                self.data["Time"].iloc[self.int_stop_idx],
                self.int_std,
                np.median(self.bkgd_subtract_data[:, self.int_std_loc]),
            ]
        ).T
        spot_data.columns = [
            "timestamp",
            "Spot",
            "despiked",
            "omitted_region",
            "bkgd_start",
            "bkgd_stop",
            "int_start",
            "int_stop",
            "norm",
            "norm_cps",
        ]
        spot_data["timestamp"] = pd.to_datetime(spot_data["timestamp"])
        spot_data = pd.concat(
            [
                spot_data,
                pd.DataFrame(
                    self.bkgd_subtract_med[np.newaxis, :], columns=self.analytes
                ),
                pd.DataFrame(
                    self.bkgd_subtract_std_err_rel[np.newaxis, :],
                    columns=[f"{analyte}_se" for analyte in self.analytes],
                ),
            ],
            axis="columns",
        )

        for col in ["bkgd_start", "bkgd_stop", "int_start", "int_stop", "norm_cps"]:
            spot_data[col] = spot_data[col].astype(np.float64)

        self.output_report = spot_data

    def despike_data(
        self,
        analyte_list: str | list[str] = "all",
        std_devs: int = 4,
        window: int = 25,
    ) -> None:
        """Despike normalised data using a rolling z-score filter.

        Applies a z-score filter to the internally normalised data to
        remove analytical spikes. Must be called **after**
        :meth:`normalize_interval` and **before**
        :meth:`make_output_report`.

        Parameters
        ----------
        analyte_list : str or list of str, optional
            Analytes to despike. Accepts a single analyte
            (e.g., ``"29Si"``), a list (e.g., ``["7Li", "29Si"]``),
            or ``"all"`` to despike every analyte, by default ``"all"``.
        std_devs : int, optional
            Number of standard deviations from the rolling mean to be
            considered an outlier, by default 4.
        window : int, optional
            Size of the rolling average window, by default 25.

        Examples
        --------
        Despike all analytes with default settings:

        >>> spot.despike_data()

        Despike a single analyte with custom parameters:

        >>> spot.despike_data(analyte_list="208Pb", std_devs=3, window=30)
        """

        assert (
            self.bkgd_subtract_normal_data is not None
        ), "please normalize your data prior to despiking"

        self.despiked = True

        if analyte_list == "all":
            filter_list = self.analytes
        else:
            if isinstance(analyte_list, list):
                pass
            else:
                analyte_list = [analyte_list]

            filter_list = analyte_list

        self.despiked_elements = filter_list

        df = pd.DataFrame(self.bkgd_subtract_normal_data, columns=self.analytes)

        for analyte in filter_list:

            filtered = _z_filter(df[analyte], window=window, std_devs=std_devs)

            # replaces data with despiked data
            df[analyte] = filtered

        self.bkgd_subtract_normal_data = df.loc[:, self.analytes].values

        # now recalculate uncertainties after despiking
        # standard error of the mean for the interval region
        self.bkgd_subtract_std_err = self.bkgd_subtract_normal_data.std(
            axis=0
        ) / np.sqrt(abs(self.int_stop_idx - self.int_start_idx))

        self.bkgd_subtract_std_err_rel = 100 * (
            self.bkgd_subtract_std_err / self.bkgd_subtract_med
        )
