"""
This module contains the Probe abstract base class and its subclasses,
used for probing for the different KDIGO criteria.
"""

from abc import ABC, ABCMeta
from enum import StrEnum, auto
from typing import Any

import numpy as np
import pandas as pd
from pandas import PeriodIndex

from pyaki.utils import Dataset, DatasetType, approx_gte, dataset_as_df, df_to_dataset


class Probe(ABC):
    """
    Abstract base class representing a data analysis probe.

    This class serves as an abstract base class (ABC) for data analysis probes.
    It declares the abstract method `probe()` that must be implemented by its subclasses.
    The `RESNAME` attribute can be overridden by subclasses to specify the name of the
    result column generated by the probe.

    Attributes
    ----------
    RESNAME : str
        The name of the result column generated by the probe.

    Methods
    -------
    probe()
        Abstract method to be implemented by subclasses. It performs data analysis on the
        provided datasets and returns a DataFrame with the analysis results.

    Example
    -------
    ```pycon
    >>> class MyProbe(Probe):
    ...     RESNAME = "my_result"
    ...
    ...     def probe(self, datasets: list[Dataset], **kwargs) -> pd.DataFrame:
    ...         # Implementation of the probe's analysis...
    ...
    ... my_probe = MyProbe()
    ... result_df = my_probe.probe(datasets=my_datasets, additional_arg=value)
    ```
    """

    RESNAME: str = ""  # name of the column that will be added to the dataframe

    def probe(self, datasets: list[Dataset], **kwargs: Any) -> list[Dataset]:
        """
        Abstract method to be implemented by subclasses.

        This method performs data analysis on the provided datasets and returns a DataFrame
        with the analysis results. Subclasses must override this method.

        Parameters
        ----------
        datasets : list[Dataset]
            A list of Dataset objects containing the input data.
        **kwargs
            Additional keyword arguments for the analysis.

        Returns
        -------
        pd.DataFrame
            The DataFrame containing the analysis results.
        """
        raise NotImplementedError()


class UrineOutputMethod(StrEnum):
    """
    Enumeration class representing different methods for urine output calculations

    Attributes
    ----------
    STRICT : str
        Strict method for urine output calculations. Using this method, the urine output stage is calculated based on the maximum urine output in the past 6, 12, and 24 hours.
    MEAN : str
        Mean method for urine output calculations.
    """

    STRICT = auto()
    MEAN = auto()


class UrineOutputProbe(Probe):
    """
    Subclass of Probe representing a probe calculating KDIGO stages according to urine output.

    This class specializes the abstract base class `Probe` to perform calculations of KDIGO stages based on urine output. Common KDIGO criteria apply.
    It overrides the `RESNAME` attribute to set the name of the result column.
    The `probe()` method performs urine output analysis on the provided DataFrame and returns a modified DataFrame
    with a column containing the appropriate KDIGO stage, according to urine output, added.

    Attributes
    ----------
    RESNAME : str
        The name of the result column representing urine output stage.

    Parameters
    ----------
    column : str, default: "urineoutput"
        The name of the column representing urine output in the DataFrame.
    patient_weight_column : str, default: "weight"
        The name of the column representing the patient's weight in the patient DataFrame.
    anuria_limit : float, default: 0.1
        The anuria limit for urine output calculations.

    Example
    -------
    ```pycon
    >>> probe = UrineOutputProbe(column="urineoutput", anuria_limit=0.1)
    ... result_df = probe.probe(df=my_dataframe, patient=patient_df)
    ```
    """

    RESNAME = "urineoutput_stage"

    def __init__(
        self,
        column: str = "urineoutput",
        patient_weight_column: str = "weight",
        anuria_limit: float = 0.1,
        method: UrineOutputMethod = UrineOutputMethod.MEAN,
    ) -> None:
        super().__init__()

        self._column: str = column
        self._patient_weight_column: str = patient_weight_column

        self._anuria_limit: float = anuria_limit
        self._method: UrineOutputMethod = method

    @dataset_as_df(df=DatasetType.URINEOUTPUT, patient=DatasetType.DEMOGRAPHICS)
    @df_to_dataset(DatasetType.URINEOUTPUT)
    def probe(
        self,
        df: pd.DataFrame,
        patient: pd.DataFrame,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """
        Perform urine output analysis on the provided DataFrame.

        This method calculates the KDIGO stage according to urine output based on the provided DataFrame and patient information DataFrame.
        It modifies the DataFrame by adding the urine output stage column with appropriate values based on the calculations.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the urine output data. We expect the DataFrame to contain urine output values in ml, sampled hourly.
        patient : pd.DataFrame
            The DataFrame containing patient information. Should contain the patients weight in kg.

        Returns
        -------
        pd.DataFrame
            The modified DataFrame with the urine output stage column added.
        """
        if self._patient_weight_column not in patient:
            raise ValueError("Missing weight for stay")

        df = df.copy()

        weight: pd.Series = patient[self._patient_weight_column]
        # fmt: off
        df.loc[:, self.RESNAME] = np.nan  # set all urineoutput_stage values to NaN
        df.loc[df.rolling(6).min()[self._column] >= 0, self.RESNAME] = 0

        if self._method == UrineOutputMethod.STRICT:
            df.loc[(df.rolling(6).max()[self._column] / weight) < 0.5, self.RESNAME] = 1
            df.loc[(df.rolling(12).max()[self._column] / weight) < 0.5, self.RESNAME] = 2
            df.loc[(df.rolling(24).max()[self._column] / weight) < 0.3, self.RESNAME] = 3
            df.loc[(df.rolling(12).max()[self._column] / weight) < self._anuria_limit, self.RESNAME] = 3
        elif self._method == UrineOutputMethod.MEAN:
            df.loc[(df.rolling(6).mean()[self._column] / weight) < 0.5, self.RESNAME] = 1
            df.loc[(df.rolling(12).mean()[self._column] / weight) < 0.5, self.RESNAME] = 2
            df.loc[(df.rolling(24).mean()[self._column] / weight) < 0.3, self.RESNAME] = 3
            df.loc[(df.rolling(12).mean()[self._column] / weight) < self._anuria_limit, self.RESNAME] = 3
        else:
            raise ValueError(f"Invalid method: {self._method}")
        # fmt: on

        df.loc[pd.isna(df[self._column]), self.RESNAME] = np.nan

        return df


class CreatinineBaselineMethod(StrEnum):
    """
    Enumeration class representing different methods for creatinine baseline calculations.

    This class defines the available methods for calculating creatinine values.
    It is a subclass of the `StrEnum` class, which is a string-based enumeration.
    The available methods are `MIN` and `FIRST`.

    Attributes
    ----------
    ROLLING_MIN : str
        Minimum of a rolling window following the timepoint of observation is used as baseline.
    ROLLING_FIRST : str
        First value of a rolling window following the timepoint of observation is used as baseline.
    ROLLING_MEAN : str
        Mean of a rolling window following the timepoint of observation is used as baseline.
    FIXED_MIN : str
        Minimum of the first n days of observation is used as baseline.
    FIXED_MEAN : str
        Mean of the first n days of observation is used as baseline.
    OVERALL_FIRST : str
        First observed value is used as baseline.
    OVERALL_MEAN : str
        Mean of all observed values is used as baseline.
    OVERALL_MIN : str
        Minimum of all observed values is used as baseline.
    CONSTANT : str
        A constant value is used as baseline.
    CALCULATED : str
        A calculated value is used as baseline, based off of the Cockcroft-Gault-Formula using the Adjusted Body Weight.
    """

    ROLLING_MIN = auto()
    ROLLING_FIRST = auto()
    ROLLING_MEAN = auto()
    FIXED_MIN = auto()
    FIXED_MEAN = auto()
    OVERALL_FIRST = auto()
    OVERALL_MEAN = auto()
    OVERALL_MIN = auto()
    CONSTANT = auto()
    CALCULATED = auto()


class AbstractCreatinineProbe(Probe, metaclass=ABCMeta):
    """
    Abstract base class representing a creatinine probe.

    This class serves as an abstract base class for creatinine probes.
    It extends the `Probe` class and provides common functionality and attributes
    for creatinine probe implementations.

    Parameters
    ----------
    column : str, default: "creat"
        The name of the column containing creatinine values.
    baseline_constant_column : str, default: "baseline_constant"
        The name of the column containing constant baseline values.
    patient_weight_column : str, default: "weight"
        The name of the column containing the patient's weight.
    patient_age_column : str, default: "age"
        The name of the column containing the patient's age.
    patient_height_column : str, default: "height"
        The name of the column containing the patient's height.
    patient_gender_column : str, default: "gender"
        The name of the column containing the patient's gender.
    baseline_timeframe : str, default: "7d"
        The timeframe for calculating the baseline values.
    expected_clearance : float, default: 72
        The expected creatinine clearance rate.
    method : CreatinineBaselineMethod, default: CreatinineBaselineMethod.ROLLING_MIN
        The method for calculating the creatinine baseline values.

    Example
    -------
    ```pycon
    >>> class MyCreatinineProbe(AbstractCreCreatinineProbe):
    ...     def __init__(self, column="creatinine", baseline_timeframe="7d", method=CreatinineBaselineMethod.MIN):
    ...         super().__init__(column, baseline_timeframe, method)
    ...         # Additional initialization
    ...
    ...     def probe(self, df, **kwargs):
    ...         # Probe implementation specific to the derived class
    ```
    """

    def __init__(
        self,
        column: str = "creat",
        baseline_constant_column: str = "baseline_constant",
        patient_weight_column: str = "weight",
        patient_age_column: str = "age",
        patient_height_column: str = "height",
        patient_gender_column: str = "gender",
        baseline_timeframe: str = "7d",
        expected_clearance: float = 72,
        method: CreatinineBaselineMethod = CreatinineBaselineMethod.ROLLING_MIN,
    ) -> None:
        super().__init__()

        self._column: str = column
        self._baseline_constant_column: str = baseline_constant_column
        self._patient_weight_column: str = patient_weight_column
        self._patient_age_column: str = patient_age_column
        self._patient_height_column: str = patient_height_column
        self._patient_gender_column: str = patient_gender_column

        self._baseline_timeframe: str = baseline_timeframe
        self._expected_clearance: float = expected_clearance
        self._method: CreatinineBaselineMethod = method

    def creatinine_baseline(self, df: pd.DataFrame, patient: pd.DataFrame) -> pd.Series:
        """
        Calculate the creatinine baseline values.

        This method calculates the creatinine baseline values based on the configured
        parameters and the provided DataFrame.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the creatinine data.

        Returns
        -------
        pd.Series
            The calculated creatinine baseline values.
        """
        if isinstance(df.index, PeriodIndex):
            df.index = df.index.to_timestamp()

        if self._method == CreatinineBaselineMethod.ROLLING_FIRST:
            return (
                df[df[self._column] > 0]
                .rolling(self._baseline_timeframe)
                .agg(lambda rows: rows.iloc[0])
                .resample("1h")
                .first()
                .ffill()[self._column]
            )

        if self._method == CreatinineBaselineMethod.ROLLING_MIN:
            return (
                df[df[self._column] > 0]
                .rolling(self._baseline_timeframe)
                .min()
                .resample("1h")
                .min()
                .ffill()[self._column]
            )
        if self._method == CreatinineBaselineMethod.ROLLING_MEAN:
            return (
                df[df[self._column] > 0]
                .rolling(self._baseline_timeframe)
                .mean()
                .resample("1h")
                .mean()
                .ffill()[self._column]
            )

        if self._method == CreatinineBaselineMethod.FIXED_MIN:
            values: pd.Series = (
                df[df[self._column] > 0]
                .rolling(self._baseline_timeframe)
                .min()
                .resample("1h")
                .min()
                .ffill()[self._column]
            )
            min_value: pd.DatetimeIndex = values[
                values.index <= (values.index[0] + pd.Timedelta(self._baseline_timeframe))
            ].min()  # calculate min value for first 7 days
            values[
                values.index > (values.index[0] + pd.Timedelta(self._baseline_timeframe))
            ] = min_value  # set all values after first 7 days to min value

            return values

        if self._method == CreatinineBaselineMethod.FIXED_MEAN:
            time_delta = pd.to_timedelta(self._baseline_timeframe)
            end_time = df.index[0] + time_delta
            value = df[df.index <= end_time][self._column].mean()
            values = self._to_df_length(df, value)
            return values

        if self._method == CreatinineBaselineMethod.OVERALL_FIRST:
            value = df[df[self._column] > 0].iloc[0][self._column]
            values = self._to_df_length(df, value)
            return values

        if self._method == CreatinineBaselineMethod.OVERALL_MIN:
            value = df[df[self._column] > 0][self._column].min()
            values = self._to_df_length(df, value)
            return values

        if self._method == CreatinineBaselineMethod.OVERALL_MEAN:
            value = df[df[self._column] > 0][self._column].mean()
            values = self._to_df_length(df, value)
            return values

        if self._method == CreatinineBaselineMethod.CONSTANT:
            if self._baseline_constant_column not in patient:
                raise ValueError(
                    "Baseline constant method requires baseline constant values. Please provide a pd.Series containing baseline values for creatinine."
                )

            return pd.Series(
                [patient[self._baseline_constant_column]] * len(df),
                index=df.index,
                name=self._column,
            )

        if self._method == CreatinineBaselineMethod.CALCULATED:
            columns = [
                self._patient_weight_column,
                self._patient_age_column,
                self._patient_height_column,
                self._patient_gender_column,
            ]
            for column in columns:
                if column not in patient:
                    raise ValueError(
                        f"Calculated baseline method requires patient {column}. Please provide a pd.Series containing patient {column}."
                    )

            weight = patient[self._patient_weight_column]
            height = patient[self._patient_height_column]
            gender = patient[self._patient_gender_column]
            age = patient[self._patient_age_column]

            ibw = (50.0 if gender == "M" else 45.5) + 2.3 * height / 2.54 - 60
            abw = ibw + 0.4 * (weight - ibw)

            # fmt: off
            return pd.Series(
                [
                    ((140 - age) * abw * (1 if gender == "M" else 0.85)) /
                    (70 * self._expected_clearance)
                ] * len(df),
                index=df.index,
                name=self._column,
            )
            # fmt: on

    def _to_df_length(self, df: pd.DataFrame, value: float) -> pd.Series:
        """
        Helper function to create a series, the same length as the data frame.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame to match the length of.
        value : float
            The value to fill the series with.

        Returns
        -------
        pd.Series
            The series with the same length as the DataFrame.
        """
        values = pd.Series(
            [value] * len(df),
            index=df.index,
            name=self._column,
        )
        return values


class AbsoluteCreatinineProbe(AbstractCreatinineProbe):
    """
    Probe class for absolute creatinine criterion, according to KDIGO criteria.

    This class represents a probe that calculates AKI stages according to absolute rises in creatinine, according to the KDIGO criteria.
    It extends the `AbstractCreCreatinineProbe` class.

    Attributes
    ----------
    RESNAME : str
        The name of the resulting stage column.

    Parameters
    ----------
    column : str, default: "creat"
        The name of the column containing creatinine values.
    baseline_constant_column : str, default: "baseline_constant"
        The name of the column containing constant baseline values.
    patient_weight_column : str, default: "weight"
        The name of the column containing the patient's weight.
    patient_age_column : str, default: "age"
        The name of the column containing the patient's age.
    patient_height_column : str, default: "height"
        The name of the column containing the patient's height.
    patient_gender_column : str, default: "gender"
        The name of the column containing the patient's gender.
    baseline_timeframe : str, default: "2d"
        The timeframe for calculating the baseline values.
    expected_clearance : float, default: 72
        The expected creatinine clearance rate.
    method : CreatinineBaselineMethod, default: CreatinineBaselineMethod.ROLLING_MIN
        The method for calculating the creatinine baseline values.

    Example
    -------
    ```pycon
    >>> probe = AbsoluteCreatinineProbe(column="creatinine", baseline_timeframe="7d", method=CreatinineMethod.MIN)
    ... df_result = probe.probe(df)
    ```
    """

    RESNAME = "abs_creatinine_stage"

    def __init__(
        self,
        column: str = "creat",
        baseline_constant_column: str = "baseline_constant",
        patient_weight_column: str = "weight",
        patient_age_column: str = "age",
        patient_height_column: str = "height",
        patient_gender_column: str = "gender",
        baseline_timeframe: str = "2d",
        expected_clearance: float = 72,
        method: CreatinineBaselineMethod = CreatinineBaselineMethod.ROLLING_MIN,
    ) -> None:
        super().__init__(
            column=column,
            baseline_constant_column=baseline_constant_column,
            patient_weight_column=patient_weight_column,
            patient_age_column=patient_age_column,
            patient_height_column=patient_height_column,
            patient_gender_column=patient_gender_column,
            baseline_timeframe=baseline_timeframe,
            expected_clearance=expected_clearance,
            method=method,
        )

    @dataset_as_df(df=DatasetType.CREATININE, patient=DatasetType.DEMOGRAPHICS)
    @df_to_dataset(DatasetType.CREATININE)
    def probe(
        self,
        df: pd.DataFrame,
        patient: pd.DataFrame,
        **kwargs: Any,
    ) -> pd.DataFrame:
        """
        Perform KDIGO stage calculation based on absolute creatinine elevations on the provided DataFrame.

        This method calculates the KDIGO stage based on the provided DataFrame
        and the configured baseline values. It calculates the stage according to KDIGO criteria.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the creatinine data. It should have a column
            with the name specified in the `column` attribute of the probe.
        patient : pd.DataFrame
            The DataFrame containing patient information. Should contain the patients weight in kg and the age.

        Returns
        -------
        pd.DataFrame
            The modified DataFrame with the absolute creatinine stage column added.
        """
        df = df.copy()

        baseline_values: pd.Series = self.creatinine_baseline(df, patient)

        df.loc[:, self.RESNAME] = 0
        df.loc[approx_gte((df[self._column] - baseline_values), 0.3), self.RESNAME] = 1
        df.loc[approx_gte(df[self._column], 4), self.RESNAME] = 3

        df.loc[pd.isna(df[self._column]), self.RESNAME] = np.nan

        return df


class RelativeCreatinineProbe(AbstractCreatinineProbe):
    """
    Probe class for relative creatinine measurements.

    This class represents a probe calculates KDIGO stages based on relative creatinine elevations.

    Attributes
    ----------
    RESNAME : str
        The name of the resulting stage column.

    Parameters
    ----------
    column : str, default: "creat"
        The name of the column containing creatinine values.
    baseline_constant_column : str, default: "baseline_constant"
        The name of the column containing constant baseline values.
    patient_weight_column : str, default: "weight"
        The name of the column containing the patient's weight.
    patient_age_column : str, default: "age"
        The name of the column containing the patient's age.
    patient_height_column : str, default: "height"
        The name of the column containing the patient's height.
    patient_gender_column : str, default: "gender"
        The name of the column containing the patient's gender.
    baseline_timeframe : str, default: "7d"
        The timeframe for calculating the baseline values.
    expected_clearance : float, default: 72
        The expected creatinine clearance rate.
    method : CreatinineBaselineMethod, default: CreatinineBaselineMethod.ROLLING_MIN
        The method for calculating the creatinine baseline values.

    Example
    -------
    ```pycon
    >>> probe = RelativeCreatinineProbe(column="creatinine", baseline_timeframe="7d", method=CreatinineBaselineMethod.MIN)
    ... df_result = probe.probe(df)
    ```
    """

    RESNAME = "rel_creatinine_stage"

    @dataset_as_df(df=DatasetType.CREATININE, patient=DatasetType.DEMOGRAPHICS)
    @df_to_dataset(DatasetType.CREATININE)
    def probe(self, df: pd.DataFrame, patient: pd.DataFrame, **kwargs: Any) -> pd.DataFrame:
        """
        Perform calculation of relative creatinine elevations on the provided DataFrame.

        This method calculates the relative creatinine stage based on the provided DataFrame
        and the configured baseline values. It modifies the DataFrame by adding the relative
        creatinine stage column with appropriate values based on the calculations.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the creatinine data. It should have a column
            with the name specified in the `column` attribute of the probe.
        patient : pd.DataFrame
            The DataFrame containing patient information. Should contain the patients weight in kg and the age.

        Returns
        -------
        pd.DataFrame
            The modified DataFrame with the relative creatinine stage column added.
        """
        df = df.copy()

        baseline_values: pd.Series = self.creatinine_baseline(df, patient)

        df.loc[:, self.RESNAME] = 0
        df.loc[approx_gte((df[self._column] / baseline_values), 1.5), self.RESNAME] = 1.0
        df.loc[approx_gte((df[self._column] / baseline_values), 2), self.RESNAME] = 2
        df.loc[approx_gte((df[self._column] / baseline_values), 3), self.RESNAME] = 3

        df.loc[pd.isna(df[self._column]), self.RESNAME] = np.nan

        return df


class RRTProbe(Probe):
    """
    Probe class for RRT.

    This class represents a probe that calculates RRT. It will return a KDIGO stage 3 if the patient is on RRT at any time during the ICU stay. It will return 0 otherwise.

    Attributes
    ----------
    RESNAME : str
        The name of the resulting stage column.

    Parameters
    ----------
    column : str, default: "rrt_status"
        The name of the column containing the RRT data.
    """

    RESNAME = "rrt_stage"

    def __init__(self, column: str = "rrt_status") -> None:
        super().__init__()

        self._column: str = column

    @dataset_as_df(df=DatasetType.RRT)
    @df_to_dataset(DatasetType.RRT)
    def probe(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Perform calculation of RRT on the provided DataFrame.

        Parameters
        ----------
        df : pd.DataFrame
            The DataFrame containing the RRT data. It should have a column
            with the name specified in the `column` attribute of the probe.

        Returns
        -------
        pd.DataFrame
            The modified DataFrame with the RRT stage column added.
        """
        df = df.copy()

        df.loc[:, self.RESNAME] = 0
        df.loc[df[self._column] == 1, self.RESNAME] = 3

        # transfer nans
        df.loc[pd.isna(df[self._column]), self.RESNAME] = np.nan

        return df
