import numpy as np
import pytest
import statsmodels
from lifelines import (
    CoxPHFitter,
    KaplanMeierFitter,
    LogLogisticAFTFitter,
    NelsonAalenFitter,
    WeibullAFTFitter,
    WeibullFitter,
)

import ehrapy as ep


@pytest.fixture
def mimic_2_sa():
    adata = ep.dt.mimic_2(encoded=False)
    adata[:, ["censor_flg"]].X = np.where(adata[:, ["censor_flg"]].X == 0, 1, 0)
    adata = adata[:, ["mort_day_censored", "censor_flg"]].copy()
    duration_col, event_col = "mort_day_censored", "censor_flg"

    return adata, duration_col, event_col


class TestSA:
    def test_ols(self):
        adata = ep.dt.mimic_2(encoded=False)
        formula = "tco2_first ~ pco2_first"
        var_names = ["tco2_first", "pco2_first"]
        ols = ep.tl.ols(adata, var_names, formula, missing="drop")
        s = ols.fit().params.iloc[1]
        i = ols.fit().params.iloc[0]
        assert isinstance(ols, statsmodels.regression.linear_model.OLS)
        assert 0.18857179158259973 == pytest.approx(s)
        assert 16.210859352601442 == pytest.approx(i)

    def test_glm(self):
        adata = ep.dt.mimic_2(encoded=False)
        formula = "day_28_flg ~ age"
        var_names = ["day_28_flg", "age"]
        family = "Binomial"
        glm = ep.tl.glm(adata, var_names, formula, family, missing="drop", as_continuous=["age"])
        Intercept = glm.fit().params.iloc[0]
        age = glm.fit().params.iloc[1]
        assert isinstance(glm, statsmodels.genmod.generalized_linear_model.GLM)
        assert 5.778006344870297 == pytest.approx(Intercept)
        assert -0.06523274132877163 == pytest.approx(age)

    @pytest.mark.parametrize("weightings", ["wilcoxon", "tarone-ware", "peto", "fleming-harrington"])
    def test_calculate_logrank_pvalue(self, weightings):
        durations_A = [1, 2, 3]
        event_observed_A = [1, 1, 0]
        durations_B = [1, 2, 3, 4]
        event_observed_B = [1, 0, 0, 1]

        kmf1 = KaplanMeierFitter()
        kmf1.fit(durations_A, event_observed_A)

        kmf2 = KaplanMeierFitter()
        kmf2.fit(durations_B, event_observed_B)

        results_pairwise = ep.tl.test_kmf_logrank(kmf1, kmf2)
        p_value_pairwise = results_pairwise.p_value
        assert 0 < p_value_pairwise < 1

    def test_anova_glm(self):
        adata = ep.dt.mimic_2(encoded=False)
        formula = "day_28_flg ~ age"
        var_names = ["day_28_flg", "age"]
        family = "Binomial"
        age_glm = ep.tl.glm(adata, var_names, formula, family, missing="drop", as_continuous=["age"])
        age_glm_result = age_glm.fit()
        formula = "day_28_flg ~ age + service_unit"
        var_names = ["day_28_flg", "age", "service_unit"]
        ageunit_glm = ep.tl.glm(adata, var_names, formula, family="Binomial", missing="drop", as_continuous=["age"])
        ageunit_glm_result = ageunit_glm.fit()
        dataframe = ep.tl.anova_glm(
            age_glm_result, ageunit_glm_result, "day_28_flg ~ age", "day_28_flg ~ age + service_unit"
        )

        assert len(dataframe) == 2
        assert dataframe.shape == (2, 6)
        assert dataframe.iloc[1, 4] == 2
        assert pytest.approx(dataframe.iloc[1, 5], 0.1) == 0.103185

    def _sa_function_assert(self, model, model_class, adata=None):
        assert isinstance(model, model_class)
        assert len(model.durations) == 1776
        assert sum(model.event_observed) == 497

        if adata is not None:  # doing it disway, due to legacy kmf function
            model_summary = adata.uns.get("test")
            assert model_summary is not None
            if isinstance(model, KaplanMeierFitter) or isinstance(
                model, NelsonAalenFitter
            ):  # kmf and nelson_aalen have event_table
                assert model_summary.equals(model.event_table)
            else:
                assert model_summary.equals(model.summary)

    def _sa_func_test(self, sa_function, sa_class, mimic_2_sa):
        adata, duration_col, event_col = mimic_2_sa
        sa = sa_function(adata, duration_col=duration_col, event_col=event_col, uns_key="test")

        self._sa_function_assert(sa, sa_class, adata)

    def test_kmf(self, mimic_2_sa):
        # check for deprecation warning
        with pytest.warns(DeprecationWarning):
            adata, _, _ = mimic_2_sa
            kmf = ep.tl.kmf(adata[:, ["mort_day_censored"]].X, adata[:, ["censor_flg"]].X)
            self._sa_function_assert(kmf, KaplanMeierFitter)

    def test_kaplan_meier(self, mimic_2_sa):
        self._sa_func_test(ep.tl.kaplan_meier, KaplanMeierFitter, mimic_2_sa)

    def test_cox_ph(self, mimic_2_sa):
        self._sa_func_test(ep.tl.cox_ph, CoxPHFitter, mimic_2_sa)

    def test_nelson_aalen(self, mimic_2_sa):
        self._sa_func_test(ep.tl.nelson_aalen, NelsonAalenFitter, mimic_2_sa)

    def test_weibull(self, mimic_2_sa):
        self._sa_func_test(ep.tl.weibull, WeibullFitter, mimic_2_sa)

    def test_weibull_aft(self, mimic_2_sa):
        self._sa_func_test(ep.tl.weibull_aft, WeibullAFTFitter, mimic_2_sa)

    def test_log_logistic(self, mimic_2_sa):
        self._sa_func_test(ep.tl.log_logistic_aft, LogLogisticAFTFitter, mimic_2_sa)