from __future__ import annotations

import warnings

from typing import TYPE_CHECKING

import matplotlib.gridspec as gridspec

import matplotlib.pyplot as plt

import matplotlib.ticker as ticker

import numpy as np

import pandas as pd

from numpy import ndarray

from ehrapy.plot import scatter

if TYPE_CHECKING:

    from collections.abc import Iterable, Sequence

    from xmlrpc.client import Boolean

    from anndata import AnnData

    from lifelines import KaplanMeierFitter

    from matplotlib.axes import Axes

    from statsmodels.regression.linear_model import RegressionResults

def ols(

    adata: AnnData | None = None,

    x: str | None = None,

    y: str | None = None,

    scatter_plot: Boolean | None = True,

    ols_results: list[RegressionResults] | None = None,

    ols_color: list[str] | None | None = None,

    xlabel: str | None = None,

    ylabel: str | None = None,

    figsize: tuple[float, float] | None = None,

    lines: list[tuple[ndarray | float, ndarray | float]] | None = None,

    lines_color: list[str] | None | None = None,

    lines_style: list[str] | None | None = None,

    lines_label: list[str] | None | None = None,

    xlim: tuple[float, float] | None = None,

    ylim: tuple[float, float] | None = None,

    show: bool | None = None,

    ax: Axes | None = None,

    title: str | None = None,

    **kwds,

) -> Axes | None:

    """Plots an Ordinary Least Squares (OLS) Model result, scatter plot, and line plot.

    Args:

        adata: :class:`~anndata.AnnData` object containing all observations.

        x: x coordinate, for scatter plotting.

        y: y coordinate, for scatter plotting.

        scatter_plot: Whether to show a scatter plot.

        ols_results: List of RegressionResults from ehrapy.tl.ols. Example: [result_1, result_2]

        ols_color: List of colors for each ols_results. Example: ['red', 'blue'].

        xlabel: The x-axis label text.

        ylabel: The y-axis label text.

        figsize: Width, height in inches.

        lines: List of Tuples of (slope, intercept) or (x, y). Plot lines by slope and intercept or data points.

               Example: plot two lines (y = x + 2 and y = 2*x + 1): [(1, 2), (2, 1)]

        lines_color: List of colors for each line. Example: ['red', 'blue']

        lines_style: List of line styles for each line. Example: ['-', '--']

        lines_label: List of line labels for each line. Example: ['Line1', 'Line2']

        xlim: Set the x-axis view limits. Required for only plotting lines using slope and intercept.

        ylim: Set the y-axis view limits. Required for only plotting lines using slope and intercept.

        show: Show the plot, do not return axis.

        ax: A matplotlib axes object. Only works if plotting a single component.

        title: Set the title of the plot.

        kwds: Passed to matplotblib scatterplot.

    Examples:

        >>> import ehrapy as ep

        >>> adata = ep.dt.mimic_2(encoded=False)

        >>> co2_lm_result = ep.tl.ols(

        ...     adata, var_names=["pco2_first", "tco2_first"], formula="tco2_first ~ pco2_first", missing="drop"

        ... ).fit()

        >>> ep.pl.ols(

        ...     adata,

        ...     x="pco2_first",

        ...     y="tco2_first",

        ...     ols_results=[co2_lm_result],

        ...     ols_color=["red"],

        ...     xlabel="PCO2",

        ...     ylabel="TCO2",

        ... )

        .. image:: /_static/docstring_previews/ols_plot_1.png

        >>> import ehrapy as ep

        >>> adata = ep.dt.mimic_2(encoded=False)

        >>> ep.pl.ols(adata, x='pco2_first', y='tco2_first', lines=[(0.25, 10), (0.3, 20)],

        >>>           lines_color=['red', 'blue'], lines_style=['-', ':'], lines_label=['Line1', 'Line2'])

        .. image:: /_static/docstring_previews/ols_plot_2.png

        >>> import ehrapy as ep

        >>> ep.pl.ols(lines=[(0.25, 10), (0.3, 20)], lines_color=['red', 'blue'], lines_style=['-', ':'],

        >>>           lines_label=['Line1', 'Line2'], xlim=(0, 150), ylim=(0, 50))

        .. image:: /_static/docstring_previews/ols_plot_3.png

    """

    if ax is None:

        _, ax = plt.subplots(figsize=figsize)

    if xlim is not None:

        plt.xlim(xlim)

    if ylim is not None:

        plt.ylim(ylim)

    if ols_color is None and ols_results is not None:

        ols_color = [None] * len(ols_results)

    if lines_color is None and lines is not None:

        lines_color = [None] * len(lines)

    if lines_style is None and lines is not None:

        lines_style = [None] * len(lines)

    if lines_label is None and lines is not None:

        lines_label = [None] * len(lines)

    if adata is not None and x is not None and y is not None:

        x_processed = np.array(adata[:, x].X).astype(float)

        x_processed = x_processed[~np.isnan(x_processed)]

        if scatter_plot is True:

            ax = scatter(adata, x=x, y=y, show=False, ax=ax, **kwds)

        if ols_results is not None:

            for i, ols_result in enumerate(ols_results):

                ax.plot(x_processed, ols_result.predict(), color=ols_color[i])

    if lines is not None:

        for i, line in enumerate(lines):

            a, b = line

            if np.ndim(a) == 0 and np.ndim(b) == 0:

                line_x = np.array(ax.get_xlim())

                line_y = a * line_x + b

                ax.plot(line_x, line_y, linestyle=lines_style[i], color=lines_color[i], label=lines_label[i])

            else:

                ax.plot(a, b, lines_style[i], color=lines_color[i], label=lines_label[i])

    plt.xlabel(xlabel)

    plt.ylabel(ylabel)

    if title:

        plt.title(title)

    if lines_label is not None and lines_label[0] is not None:

        plt.legend()

    if not show:

        return ax

    else:

        return None

def kmf(

    kmfs: Sequence[KaplanMeierFitter],

    ci_alpha: list[float] | None = None,

    ci_force_lines: list[Boolean] | None = None,

    ci_show: list[Boolean] | None = None,

    ci_legend: list[Boolean] | None = None,

    at_risk_counts: list[Boolean] | None = None,

    color: list[str] | None | None = None,

    grid: Boolean | None = False,

    xlim: tuple[float, float] | None = None,

    ylim: tuple[float, float] | None = None,

    xlabel: str | None = None,

    ylabel: str | None = None,

    figsize: tuple[float, float] | None = None,

    show: bool | None = None,

    title: str | None = None,

) -> Axes | None:

    warnings.warn(

        "This function is deprecated and will be removed in the next release. Use `ep.pl.kaplan_meier` instead.",

        DeprecationWarning,

        stacklevel=2,

    )

    return kaplan_meier(

        kmfs=kmfs,

        ci_alpha=ci_alpha,

        ci_force_lines=ci_force_lines,

        ci_show=ci_show,

        ci_legend=ci_legend,

        at_risk_counts=at_risk_counts,

        color=color,

        grid=grid,

        xlim=xlim,

        ylim=ylim,

        xlabel=xlabel,

        ylabel=ylabel,

        figsize=figsize,

        show=show,

        title=title,

    )

def kaplan_meier(

    kmfs: Sequence[KaplanMeierFitter],

    *,

    display_survival_statistics: bool = False,

    ci_alpha: list[float] | None = None,

    ci_force_lines: list[Boolean] | None = None,

    ci_show: list[Boolean] | None = None,

    ci_legend: list[Boolean] | None = None,

    at_risk_counts: list[Boolean] | None = None,

    color: list[str] | None | None = None,

    grid: Boolean | None = False,

    xlim: tuple[float, float] | None = None,

    ylim: tuple[float, float] | None = None,

    xlabel: str | None = None,

    ylabel: str | None = None,

    figsize: tuple[float, float] | None = None,

    show: bool | None = None,

    title: str | None = None,

) -> Axes | None:

    """Plots a pretty figure of the Fitted KaplanMeierFitter model.

    See https://lifelines.readthedocs.io/en/latest/fitters/univariate/KaplanMeierFitter.html

    Args:

        kmfs: Iterables of fitted KaplanMeierFitter objects.

        display_survival_statistics: Whether to show survival statistics in a table below the plot.

        ci_alpha: The transparency level of the confidence interval. If more than one kmfs, this should be a list.

        ci_force_lines: Force the confidence intervals to be line plots (versus default shaded areas).

                        If more than one kmfs, this should be a list.

        ci_show: Show confidence intervals. If more than one kmfs, this should be a list.

        ci_legend: If ci_force_lines is True, this is a boolean flag to add the lines' labels to the legend.

                   If more than one kmfs, this should be a list.

        at_risk_counts: Show group sizes at time points. If more than one kmfs, this should be a list.

        color: List of colors for each kmf. If more than one kmfs, this should be a list.

        grid: If True, plot grid lines.

        xlim: Set the x-axis view limits.

        ylim: Set the y-axis view limits.

        xlabel: The x-axis label text.

        ylabel: The y-axis label text.

        figsize: Width, height in inches.

        show: Show the plot, do not return axis.

        title: Set the title of the plot.

    Examples:

        >>> import ehrapy as ep

        >>> import numpy as np

        >>> adata = ep.dt.mimic_2(encoded=False)

        # Because in MIMIC-II database, `censor_fl` is censored or death (binary: 0 = death, 1 = censored).

        # While in KaplanMeierFitter, `event_observed` is True if the the death was observed, False if the event was lost (right-censored).

        # So we need to flip `censor_fl` when pass `censor_fl` to KaplanMeierFitter

        >>> adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)

        >>> kmf = ep.tl.kaplan_meier(adata, "mort_day_censored", "censor_flg")

        >>> ep.pl.kaplan_meier(

        ...     [kmf], color=["r"], xlim=[0, 700], ylim=[0, 1], xlabel="Days", ylabel="Proportion Survived", show=True

        ... )

        .. image:: /_static/docstring_previews/kmf_plot_1.png

        >>> groups = adata[:, ["service_unit"]].X

        >>> adata_ficu = adata[groups == "FICU"]

        >>> adata_micu = adata[groups == "MICU"]

        >>> adata_sicu = adata[groups == "SICU"]

        >>> kmf_1 = ep.tl.kaplan_meier(adata_ficu, "mort_day_censored", "censor_flg", label="FICU")

        >>> kmf_2 = ep.tl.kaplan_meier(adata_micu, "mort_day_censored", "censor_flg", label="MICU")

        >>> kmf_3 = ep.tl.kaplan_meier(adata_sicu, "mort_day_censored", "censor_flg", label="SICU")

        >>> ep.pl.kaplan_meier([kmf_1, kmf_2, kmf_3], ci_show=[False,False,False], color=['k','r', 'g'],

        >>>           xlim=[0, 750], ylim=[0, 1], xlabel="Days", ylabel="Proportion Survived")

        .. image:: /_static/docstring_previews/kmf_plot_2.png

    """

    if ci_alpha is None:

        ci_alpha = [0.3] * len(kmfs)

    if ci_force_lines is None:

        ci_force_lines = [False] * len(kmfs)

    if ci_show is None:

        ci_show = [True] * len(kmfs)

    if ci_legend is None:

        ci_legend = [False] * len(kmfs)

    if at_risk_counts is None:

        at_risk_counts = [False] * len(kmfs)

    if color is None:

        color = [None] * len(kmfs)

    fig = plt.figure(constrained_layout=True, figsize=figsize)

    spec = fig.add_gridspec(2, 1) if display_survival_statistics else fig.add_gridspec(1, 1)

    ax = plt.subplot(spec[0, 0])

    for i, kmf in enumerate(kmfs):

        if i == 0:

            ax = kmf.plot_survival_function(

                ci_alpha=ci_alpha[i],

                ci_force_lines=ci_force_lines[i],

                ci_show=ci_show[i],

                ci_legend=ci_legend[i],

                at_risk_counts=at_risk_counts[i],

                color=color[i],

            )

        else:

            ax = kmf.plot_survival_function(

                ax=ax,

                ci_alpha=ci_alpha[i],

                ci_force_lines=ci_force_lines[i],

                ci_show=ci_show[i],

                ci_legend=ci_legend[i],

                at_risk_counts=at_risk_counts[i],

                color=color[i],

            )

    # Configure plot appearance

    ax.grid(grid)

    ax.set_xlim(xlim)

    ax.set_ylim(ylim)

    ax.set_xlabel(xlabel)

    ax.set_ylabel(ylabel)

    if title:

        ax.set_title(title)

    if display_survival_statistics:

        xticks = [x for x in ax.get_xticks() if x >= 0]

        xticks_space = xticks[1] - xticks[0]

        if xlabel is None:

            xlabel = "Time"

        yticks = np.arange(len(kmfs))

        ax_table = plt.subplot(spec[1, 0])

        ax_table.set_xticks(xticks)

        ax_table.set_xlim(-xticks_space / 2, xticks[-1] + xticks_space / 2)

        ax_table.set_ylim(-1, len(kmfs))

        ax_table.set_yticks(yticks)

        ax_table.set_yticklabels([kmf.label if kmf.label else f"Group {i + 1}" for i, kmf in enumerate(kmfs[::-1])])

        for i, kmf in enumerate(kmfs[::-1]):

            survival_probs = kmf.survival_function_at_times(xticks).values

            for j, prob in enumerate(survival_probs):

                ax_table.text(

                    xticks[j],  # x position

                    yticks[i],  # y position

                    f"{prob:.2f}",  # formatted survival probability

                    ha="center",

                    va="center",

                    bbox={"boxstyle": "round,pad=0.2", "edgecolor": "none", "facecolor": "lightgrey"},

                )

        ax_table.grid(grid)

        ax_table.spines["top"].set_visible(False)

        ax_table.spines["right"].set_visible(False)

        ax_table.spines["bottom"].set_visible(False)

        ax_table.spines["left"].set_visible(False)

    if not show:

        return fig, ax

    else:

        return None

def cox_ph_forestplot(

    adata: AnnData,

    *,

    uns_key: str = "cox_ph",

    labels: Iterable[str] | None = None,

    fig_size: tuple = (10, 10),

    t_adjuster: float = 0.1,

    ecolor: str = "dimgray",

    size: int = 3,

    marker: str = "o",

    decimal: int = 2,

    text_size: int = 12,

    color: str = "k",

    show: bool = None,

    title: str | None = None,

):

    """Generates a forest plot to visualize the coefficients and confidence intervals of a Cox Proportional Hazards model.

    The `adata` object must first be populated using the :func:`~ehrapy.tools.cox_ph` function. This function stores the summary table of the `CoxPHFitter` in the `.uns` attribute of `adata`.

    The summary table is created when the model is fitted using the :func:`~ehrapy.tools.cox_ph` function.

    For more information on the `CoxPHFitter`, see the `Lifelines documentation <https://lifelines.readthedocs.io/en/latest/fitters/regression/CoxPHFitter.html>`_.

    Inspired by `zepid.graphics.EffectMeasurePlot <https://readthedocs.org>`_ (zEpid Package, https://pypi.org/project/zepid/).

    Args:

        adata: :class:`~anndata.AnnData` object containing the summary table from the CoxPHFitter. This is stored in the `.uns` attribute, after fitting the model using :func:`~ehrapy.tools.cox_ph`.

        uns_key: Key in `.uns` where :func:`~ehrapy.tools.cox_ph` function stored the summary table. See argument `uns_key` in :func:`~ehrapy.tools.cox_ph`.

        labels: List of labels for each coefficient, default uses the index of the summary ta

        fig_size: Width, height in inches.

        t_adjuster: Adjust the table to the right.

        ecolor: Color of the error bars.

        size: Size of the markers.

        marker: Marker style.

        decimal: Number of decimal places to display.

        text_size: Font size of the text.

        color: Color of the markers.

        show: Show the plot, do not return figure and axis.

        title: Set the title of the plot.

    Examples:

        >>> import ehrapy as ep

        >>> adata = ep.dt.mimic_2(encoded=False)

        >>> adata_subset = adata[:, ["mort_day_censored", "censor_flg", "gender_num", "afib_flg", "day_icu_intime_num"]]

        >>> coxph = ep.tl.cox_ph(adata_subset, event_col="censor_flg", duration_col="mort_day_censored")

        >>> ep.pl.cox_ph_forestplot(adata_subset)

        .. image:: /_static/docstring_previews/coxph_forestplot.png

    """

    if uns_key not in adata.uns:

        raise ValueError(f"Key {uns_key} not found in adata.uns. Please provide a valid key.")

    coxph_fitting_summary = adata.uns[

        uns_key

    ]  # pd.Dataframe with columns: coef, exp(coef), se(coef), z, p, lower 0.95, upper 0.95

    auc_col = "coef"

    if labels is None:

        labels = coxph_fitting_summary.index

    tval = []

    ytick = []

    for row_index in range(len(coxph_fitting_summary)):

        if not np.isnan(coxph_fitting_summary[auc_col][row_index]):

            if (

                (isinstance(coxph_fitting_summary[auc_col][row_index], float))

                & (isinstance(coxph_fitting_summary["coef lower 95%"][row_index], float))

                & (isinstance(coxph_fitting_summary["coef upper 95%"][row_index], float))

            ):

                tval.append(

                    [

                        round(coxph_fitting_summary[auc_col][row_index], decimal),

                        (

                            "("

                            + str(round(coxph_fitting_summary["coef lower 95%"][row_index], decimal))

                            + ", "

                            + str(round(coxph_fitting_summary["coef upper 95%"][row_index], decimal))

                            + ")"

                        ),

                    ]

                )

            else:

                tval.append(

                    [

                        coxph_fitting_summary[auc_col][row_index],

                        (

                            "("

                            + str(coxph_fitting_summary["coef lower 95%"][row_index])

                            + ", "

                            + str(coxph_fitting_summary["coef upper 95%"][row_index])

                            + ")"

                        ),

                    ]

                )

            ytick.append(row_index)

        else:

            tval.append([" ", " "])

            ytick.append(row_index)

    x_axis_upper_bound = round(((pd.to_numeric(coxph_fitting_summary["coef upper 95%"])).max() + 0.1), 2)

    x_axis_lower_bound = round(((pd.to_numeric(coxph_fitting_summary["coef lower 95%"])).min() - 0.1), 1)

    fig = plt.figure(figsize=fig_size)

    gspec = gridspec.GridSpec(1, 6)

    plot = plt.subplot(gspec[0, 0:4])

    table = plt.subplot(gspec[0, 4:])

    plot.set_ylim(-1, (len(coxph_fitting_summary)))  # spacing out y-axis properly

    plot.axvline(1, color="gray", zorder=1)

    lower_diff = coxph_fitting_summary[auc_col] - coxph_fitting_summary["coef lower 95%"]

    upper_diff = coxph_fitting_summary["coef upper 95%"] - coxph_fitting_summary[auc_col]

    plot.errorbar(

        coxph_fitting_summary[auc_col],

        coxph_fitting_summary.index,

        xerr=[lower_diff, upper_diff],

        marker="None",

        zorder=2,

        ecolor=ecolor,

        linewidth=0,

        elinewidth=1,

    )

    # plot markers

    plot.scatter(

        coxph_fitting_summary[auc_col],

        coxph_fitting_summary.index,

        c=color,

        s=(size * 25),

        marker=marker,

        zorder=3,

        edgecolors="None",

    )

    # plot settings

    plot.xaxis.set_ticks_position("bottom")

    plot.yaxis.set_ticks_position("left")

    plot.get_xaxis().set_major_formatter(ticker.ScalarFormatter())

    plot.get_xaxis().set_minor_formatter(ticker.NullFormatter())

    plot.set_yticks(ytick)

    plot.set_xlim([x_axis_lower_bound, x_axis_upper_bound])

    plot.set_xticks([x_axis_lower_bound, 1, x_axis_upper_bound])

    plot.set_xticklabels([x_axis_lower_bound, 1, x_axis_upper_bound])

    plot.set_yticklabels(labels)

    plot.tick_params(axis="y", labelsize=text_size)

    plot.yaxis.set_ticks_position("none")

    plot.invert_yaxis()  # invert y-axis to align values properly with table

    tb = table.table(

        cellText=tval, cellLoc="center", loc="right", colLabels=[auc_col, "95% CI"], bbox=[0, t_adjuster, 1, 1]

    )

    table.axis("off")

    tb.auto_set_font_size(False)

    tb.set_fontsize(text_size)

    for _, cell in tb.get_celld().items():

        cell.set_linewidth(0)

    # remove spines

    plot.spines["top"].set_visible(False)

    plot.spines["right"].set_visible(False)

    plot.spines["left"].set_visible(False)

    if title:

        plt.title(title)

    if not show:

        return fig, plot

    else:

        return None