# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import numpy.testing as npt
import pytest

from pandas.testing import assert_frame_equal
from lifelines import statistics as stats
from lifelines import CoxPHFitter, KaplanMeierFitter, WeibullFitter
from lifelines.exceptions import StatisticalWarning
from lifelines.datasets import load_waltons, load_g3, load_lymphoma, load_dd, load_regression_dataset, load_leukemia


def test_sample_size_necessary_under_cph():
    assert stats.sample_size_necessary_under_cph(0.8, 1, 0.8, 0.2, 0.139) == (14, 14)
    assert stats.sample_size_necessary_under_cph(0.8, 1, 0.5, 0.5, 1.2) == (950, 950)
    assert stats.sample_size_necessary_under_cph(0.8, 1.5, 0.5, 0.5, 1.2) == (1231, 821)
    assert stats.sample_size_necessary_under_cph(0.8, 1.5, 0.5, 0.5, 1.2, alpha=0.01) == (1832, 1221)


def test_power_under_cph():
    assert abs(stats.power_under_cph(12, 12, 0.8, 0.2, 0.139) - 0.744937) < 10e-6
    assert abs(stats.power_under_cph(12, 20, 0.8, 0.2, 1.2) - 0.05178317) < 10e-6


def test_unequal_intensity_with_random_data():
    data1 = np.random.exponential(5, size=(2000, 1))
    data2 = np.random.exponential(1, size=(2000, 1))
    test_result = stats.logrank_test(data1, data2)
    assert test_result.p_value < 0.05


def test_logrank_test_output_against_R_1():
    df = load_g3()
    ix = df["group"] == "RIT"
    d1, e1 = df.loc[ix]["time"], df.loc[ix]["event"]
    d2, e2 = df.loc[~ix]["time"], df.loc[~ix]["event"]

    expected = 0.0138
    result = stats.logrank_test(d1, d2, event_observed_A=e1, event_observed_B=e2)
    assert abs(result.p_value - expected) < 0.0001


def test_logrank_test_output_against_R_2():
    # from https://stat.ethz.ch/education/semesters/ss2011/seminar/contents/presentation_2.pdf
    control_T = [1, 1, 2, 2, 3, 4, 4, 5, 5, 8, 8, 8, 8, 11, 11, 12, 12, 15, 17, 22, 23]
    control_E = np.ones_like(control_T)

    treatment_T = [6, 6, 6, 7, 10, 13, 16, 22, 23, 6, 9, 10, 11, 17, 19, 20, 25, 32, 32, 34, 25]
    treatment_E = [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

    result = stats.logrank_test(control_T, treatment_T, event_observed_A=control_E, event_observed_B=treatment_E)
    expected_p_value = 4.17e-05

    assert abs(result.p_value - expected_p_value) < 0.0001
    assert abs(result.test_statistic - 16.8) < 0.1


def test_logrank_test_with_t_0():
    control_T = [1, 1, 2, 2, 3, 4, 4, 5, 5, 8, 8, 8, 8, 11, 11, 12, 12, 15, 17, 22, 23]
    control_E = np.ones_like(control_T)

    treatment_T = [6, 6, 6, 7, 10, 13, 16, 22, 23, 6, 9, 10, 11, 17, 19, 20, 25, 32, 32, 34, 25]
    treatment_E = [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

    result = stats.logrank_test(control_T, treatment_T, event_observed_A=control_E, event_observed_B=treatment_E, t_0=10)


def test_rank_test_output_against_R_no_censorship():
    """
    > time <- c(10,20,30,10,20,50)
    > status <- c(1,1,1,1,1,1)
    > treatment <- c(1,1,1,0,0,0)
    > survdiff(Surv(time, status) ~ treatment)
    """
    result = stats.multivariate_logrank_test([10, 20, 30, 10, 20, 50], [1, 1, 1, 0, 0, 0])
    r_p_value = 0.614107
    r_stat = 0.254237
    assert abs(result.p_value - r_p_value) < 10e-6
    assert abs(result.test_statistic - r_stat) < 10e-6


def test_load_lymphoma_logrank():
    # from https://www.statsdirect.com/help/content/survival_analysis/logrank.htm
    df_ = load_lymphoma()
    results = stats.multivariate_logrank_test(df_["Time"], df_["Stage_group"], df_["Censor"])
    assert abs(results.test_statistic - 6.70971) < 1e-4
    assert abs(results.p_value - 0.0096) < 1e-4


def test_multivariate_logrank_on_dd_dataset():
    """
    library('survival')
    dd = read.csv('~/code/lifelines/lifelines/datasets/dd.csv')
    results = survdiff(Surv(duration, observed)~regime, data=dd, rho=0)
    results[5]
    """
    dd = load_dd()
    results = stats.multivariate_logrank_test(dd["duration"], dd["regime"], dd["observed"])
    assert abs(results.test_statistic - 322.5991) < 0.0001


def test_rank_test_output_against_R_with_censorship():
    """
    > time <- c(10,20,30,10,20,50)
    > status <- c(1,0,1,1,0,1)
    > treatment <- c(1,1,1,0,0,0)
    > survdiff(Surv(time, status) ~ treatment)
    """
    result = stats.multivariate_logrank_test([10, 20, 30, 10, 20, 50], [1, 1, 1, 0, 0, 0], [1, 0, 1, 1, 0, 1])
    r_p_value = 0.535143
    r_stat = 0.384615
    assert abs(result.p_value - r_p_value) < 10e-6
    assert abs(result.test_statistic - r_stat) < 10e-6


def test_wilcoxon_weighted_logrank_on_leukemia_dataset():
    """
    Test against result from "Survival Analysis: A Self-learning Text" by Kleinbaum & Klein, 3rd edition, 2012.
    """
    data = load_leukemia()
    group_1 = data[data["Rx"] == 0]
    group_2 = data[data["Rx"] == 1]

    result = stats.logrank_test(group_1["t"], group_2["t"], group_1["status"], group_2["status"], weightings="wilcoxon")

    assert abs(result.test_statistic - 13.457852) < 10e-6
    assert result.test_name == "Wilcoxon_test"


def test_tarone_ware_weighted_logrank_on_leukemia_dataset():
    """
    Test against result from "Survival Analysis: A Self-learning Text" by Kleinbaum & Klein, 3rd edition, 2012.
    """
    data = load_leukemia()
    group_1 = data[data["Rx"] == 0]
    group_2 = data[data["Rx"] == 1]

    result = stats.logrank_test(group_1["t"], group_2["t"], group_1["status"], group_2["status"], weightings="tarone-ware")

    assert abs(result.test_statistic - 15.123575) < 10e-6
    assert result.test_name == "Tarone-Ware_test"


def test_peto_weighted_logrank_on_leukemia_dataset():
    """
    Test against result from "Survival Analysis: A Self-learning Text" by Kleinbaum & Klein, 3rd edition, 2012.
    """
    data = load_leukemia()
    group_1 = data[data["Rx"] == 0]
    group_2 = data[data["Rx"] == 1]

    result = stats.logrank_test(group_1["t"], group_2["t"], group_1["status"], group_2["status"], weightings="peto")

    assert abs(result.test_statistic - 14.084139) < 10e-6
    assert result.test_name == "Peto_test"


def test_unequal_intensity_event_observed():
    data1 = np.random.exponential(5, size=(2000, 1))
    data2 = np.random.exponential(1, size=(2000, 1))
    eventA = np.random.binomial(1, 0.5, size=(2000, 1))
    eventB = np.random.binomial(1, 0.5, size=(2000, 1))
    result = stats.logrank_test(data1, data2, event_observed_A=eventA, event_observed_B=eventB)
    assert result.p_value < 0.05


def test_integer_times_logrank_test():
    data1 = np.random.exponential(5, size=(2000, 1)).astype(int)
    data2 = np.random.exponential(1, size=(2000, 1)).astype(int)
    result = stats.logrank_test(data1, data2)
    assert result.p_value < 0.05


def test_equal_intensity_with_negative_data():
    data1 = np.random.normal(0, size=(2000, 1))
    data1 -= data1.mean()
    data1 /= data1.std()
    data2 = np.random.normal(0, size=(2000, 1))
    data2 -= data2.mean()
    data2 /= data2.std()
    result = stats.logrank_test(data1, data2)
    assert result.p_value > 0.05


def test_unequal_intensity_with_negative_data():
    data1 = np.random.normal(-5, size=(2000, 1))
    data2 = np.random.normal(5, size=(2000, 1))
    result = stats.logrank_test(data1, data2)
    assert result.p_value < 0.05


def test_log_rank_test_on_waltons_dataset():
    df = load_waltons()
    ix = df["group"] == "miR-137"
    waltonT1 = df.loc[ix]["T"]
    waltonT2 = df.loc[~ix]["T"]
    result = stats.logrank_test(waltonT1, waltonT2)
    assert result.p_value < 0.05


def test_log_rank_test_on_waltons_dataset_with_case_weights():
    df = load_waltons()
    ix = df["group"] == "miR-137"
    waltonT1 = df.loc[ix]["T"]
    waltonT2 = df.loc[~ix]["T"]
    result = stats.logrank_test(waltonT1, waltonT2)
    print(result)

    dfw = df.groupby(["T", "E", "group"]).size().reset_index().rename(columns={0: "weights"})
    ixw = dfw["group"] == "miR-137"
    waltonT1w = dfw.loc[ixw]["T"]
    waltonT2w = dfw.loc[~ixw]["T"]

    weightsA = dfw.loc[ixw]["weights"]
    weightsB = dfw.loc[~ixw]["weights"]

    resultw = stats.logrank_test(waltonT1w, waltonT2w, weights_A=weightsA, weights_B=weightsB)
    assert_frame_equal(resultw.summary, result.summary)


def test_logrank_test_is_symmetric():
    data1 = np.random.exponential(5, size=(2000, 1)).astype(int)
    data2 = np.random.exponential(1, size=(2000, 1)).astype(int)
    result1 = stats.logrank_test(data1, data2)
    result2 = stats.logrank_test(data2, data1)
    assert abs(result1.p_value - result2.p_value) < 10e-8


def test_multivariate_unequal_intensities():
    T = np.random.exponential(10, size=300)
    g = np.random.binomial(2, 0.5, size=300)
    T[g == 1] = np.random.exponential(1, size=(g == 1).sum())
    result = stats.multivariate_logrank_test(T, g)
    assert result.p_value < 0.05


def test_pairwise_waltons_dataset_is_significantly_different():
    waltons_dataset = load_waltons()
    R = stats.pairwise_logrank_test(waltons_dataset["T"], waltons_dataset["group"])
    assert R.summary.loc[("control", "miR-137")]["p"] < 0.05


def test_pairwise_allows_dataframes_and_gives_correct_counts():
    N = 100
    N_groups = 5
    df = pd.DataFrame(np.empty((N, 3)), columns=["T", "C", "group"])
    df["T"] = np.random.exponential(1, size=N)
    df["C"] = np.random.binomial(1, 0.6, size=N)
    df["group"] = np.tile(np.arange(N_groups), 20)
    R = stats.pairwise_logrank_test(df["T"], df["group"], event_observed=df["C"])
    assert R.summary.shape[0] == N_groups * (N_groups - 1) / 2


def test_log_rank_returns_None_if_equal_arrays():
    T = np.random.exponential(5, size=200)
    result = stats.logrank_test(T, T)
    assert result.p_value > 0.05

    C = np.random.binomial(1, 0.8, size=200)
    result = stats.logrank_test(T, T, C, C)
    assert result.p_value > 0.05


def test_multivariate_log_rank_is_identital_to_log_rank_for_n_equals_2():
    N = 200
    T1 = np.random.exponential(5, size=N)
    T2 = np.random.exponential(5, size=N)
    C1 = np.random.binomial(1, 0.9, size=N)
    C2 = np.random.binomial(1, 0.9, size=N)
    result = stats.logrank_test(T1, T2, C1, C2)

    T = np.r_[T1, T2]
    C = np.r_[C1, C2]
    G = np.array([1] * 200 + [2] * 200)
    result_m = stats.multivariate_logrank_test(T, G, C)
    assert result.p_value == result_m.p_value


def test_StatisticalResult_kwargs():

    sr = stats.StatisticalResult(0.05, 5.0, kw="some_value", test_name="test")
    assert hasattr(sr, "kw")
    assert getattr(sr, "kw") == "some_value"
    assert "some_value" in sr.to_ascii(2)


def test_StatisticalResult_can_be_added():

    sr1 = stats.StatisticalResult(0.01, 1.0, name=["1"], kw1="some_value1")
    sr2 = stats.StatisticalResult([0.02], [2.0], name=["2"], kw2="some_value2")
    sr3 = stats.StatisticalResult([0.03, 0.04], [3.3, 4.4], name=["3", "4"], kw3=3)
    sr = sr1 + sr2 + sr3

    assert sr.summary.shape[0] == 4
    assert sr.summary.index.tolist() == ["1", "2", "3", "4"]
    assert "kw3" in sr._kwargs


def test_StatisticalResult_can_be_printed():

    sr = stats.StatisticalResult(0.01, 1.0, name=["1"], kw1="some_value1", test_name="test")
    sr.print_summary(decimals=3, style="ascii")
    sr.print_summary(decimals=3, style="latex")
    sr.print_summary(decimals=3, style="html")


def test_proportional_hazard_test():
    """
    c = coxph(formula=Surv(T, E) ~ var1 + var2 + var3, data=df)
    cz = cox.zph(c, transform='rank')
    cz
    """
    cph = CoxPHFitter()
    df = load_regression_dataset()
    cph.fit(df, "T", "E")
    results = stats.proportional_hazard_test(cph, df)
    npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 1.4938293, rtol=1e-3)
    npt.assert_allclose(results.summary.loc["var2"]["test_statistic"], 0.8792998, rtol=1e-3)
    npt.assert_allclose(results.summary.loc["var3"]["test_statistic"], 2.2686088, rtol=1e-3)
    npt.assert_allclose(results.summary.loc["var3"]["p"], 0.1320184, rtol=1e-3)


def test_proportional_hazard_test_with_log_transform():
    cph = CoxPHFitter()
    df = load_regression_dataset()
    cph.fit(df, "T", "E")

    results = stats.proportional_hazard_test(cph, df, time_transform="log")
    npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 2.227627, rtol=1e-3)
    npt.assert_allclose(results.summary.loc["var2"]["test_statistic"], 0.714427, rtol=1e-3)
    npt.assert_allclose(results.summary.loc["var3"]["test_statistic"], 1.466321, rtol=1e-3)
    npt.assert_allclose(results.summary.loc["var3"]["p"], 0.225927, rtol=1e-3)


def test_proportional_hazard_test_with_weights():
    """

    library(survival)
    df <- data.frame(
      "var1" = c(0.209325, 0.693919, 0.443804, 0.065636, 0.386294),
      "T" = c(5.269797, 6.601666, 7.335846, 11.684092, 12.678458),
      "E" = c(1, 1, 1, 1, 1),
      "w" = c(1, 0.5, 2, 1, 1)
    )

    c = coxph(formula=Surv(T, E) ~ var1 , data=df, weights=w)
    cox.zph(c, transform='rank')
    """

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "w": [1, 0.5, 2, 1, 1],
        }
    )
    df["E"] = True

    cph = CoxPHFitter()
    cph.fit(df, "T", "E", weights_col="w")

    results = stats.proportional_hazard_test(cph, df, time_transform=["km", "rank", "log", "identity"])
    results.print_summary(5)
    npt.assert_allclose(results.summary.loc["var1", "rank"]["test_statistic"], 0.108, rtol=1e-2)


def test_proportional_hazard_test_with_strata_weights_and_strata():
    """
    library(survival)
    df <- data.frame(
      "var1" = c(0.209325, 0.693919, 0.443804, 0.065636, 0.386294),
      "T" = c(5.269797, 6.601666, 7.335846, 11.684092, 12.678458),
      "E" = c(1, 1, 1, 1, 1),
      "w" = c(1, 0.5, 2, 1, 1),
      "s" = c(1, 1, 0, 0, 0)
    )

    c = coxph(formula=Surv(T, E) ~ var1 + strata(s), data=df, weights=w)
    cz = cox.zph(c, transform='identity')
    cz_km = cox.zph(c, transform='km')

    """

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "w": [1, 0.5, 2, 1, 1],
            "s": [1, 1, 0, 0, 0],
        }
    )
    df["E"] = True

    cph = CoxPHFitter()
    cph.fit(df, "T", "E", weights_col="w", strata="s")

    results = stats.proportional_hazard_test(cph, df, time_transform="identity")
    npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 0.0283, rtol=1e-3)

    results = stats.proportional_hazard_test(cph, df, time_transform="km")
    npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 0.0434, rtol=1e-1)


def test_proportional_hazard_test_with_kmf():
    """
    library(survival)
    df <- data.frame(
      "var1" = c(0.209325, 0.693919, 0.443804, 0.065636, 0.386294),
      "T" = c(5.269797, 6.601666, 7.335846, 11.684092, 12.678458),
      "E" = c(1, 1, 1, 1, 1)
    )

    c = coxph(formula=Surv(T, E) ~ var1 , data=df)
    cox.zph(c, transform='km')
    """

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "E": [1, 1, 1, 1, 1],
        }
    )

    cph = CoxPHFitter()
    cph.fit(df, "T", "E")

    results = stats.proportional_hazard_test(cph, df)
    npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 0.00971, rtol=1e-3)


def test_proportional_hazard_test_with_kmf_with_some_censorship():
    """

    library(survival)
    df <- data.frame(
      "var1" = c(0.209325, 0.693919, 0.443804, 0.065636, 0.386294),
      "T" = c(5.269797, 6.601666, 7.335846, 11.684092, 12.678458),
      "E" = c(1, 1, 1, 0, 1)
    )

    c = coxph(formula=Surv(T, E) ~ var1 , data=df)
    cox.zph(c, transform='km')
    """

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "E": [1, 1, 1, 0, 1],
        }
    )

    cph = CoxPHFitter()
    cph.fit(df, "T", "E")

    results = stats.proportional_hazard_test(cph, df)
    npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 1.013802, rtol=1e-3)


def test_proportional_hazard_test_with_kmf_with_some_censorship_and_weights():
    """

    library(survival)
    df <- data.frame(
      "var1" = c(0.209325, 0.693919, 0.443804, 0.065636, 0.386294),
      "T" = c(5.269797, 6.601666, 7.335846, 11.684092, 12.678458),
      "E" = c(1, 1, 1, 0, 1),
      "w" = c(1, 0.5, 2, 1, 1),
    )

    c = coxph(formula=Surv(T, E) ~ var1 , data=df, weights=w)
    cox.zph(c, transform='km')
    """

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "E": [1, 1, 1, 0, 1],
            "w": [1, 0.5, 5, 1, 1],
        }
    )

    cph = CoxPHFitter()
    with pytest.warns(StatisticalWarning, match="weights are not integers"):
        cph.fit(df, "T", "E", weights_col="w")
        results = stats.proportional_hazard_test(cph, df)
        npt.assert_allclose(results.summary.loc["var1"]["test_statistic"], 0.916, rtol=1e-2)


def test_proportional_hazard_test_with_all():

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "E": [1, 1, 1, 0, 1],
        }
    )

    cph = CoxPHFitter()
    cph.fit(df, "T", "E")
    results = stats.proportional_hazard_test(cph, df, time_transform="all")
    assert results.summary.shape[0] == 1 * 4


def test_proportional_hazard_test_with_list():

    df = pd.DataFrame(
        {
            "var1": [0.209325, 0.693919, 0.443804, 0.065636, 0.386294],
            "var2": [1, 0, 1, 0, 1],
            "T": [5.269797, 6.601666, 7.335846, 11.684092, 12.678458],
            "E": [1, 1, 1, 0, 1],
        }
    )

    cph = CoxPHFitter()
    cph.fit(df, "T", "E")
    results = stats.proportional_hazard_test(cph, df, time_transform=["rank", "km"])
    assert results.summary.shape[0] == 2 * 2


def test_survival_difference_at_fixed_point_in_time_test_nonparametric():
    df = load_waltons()
    ix = df["group"] == "miR-137"
    kmf1 = KaplanMeierFitter().fit(df.loc[ix]["T"], df.loc[ix]["E"])
    kmf2 = KaplanMeierFitter().fit(df.loc[~ix]["T"], df.loc[~ix]["E"])
    result = stats.survival_difference_at_fixed_point_in_time_test(10, kmf1, kmf2)
    assert result.p_value < 0.05


def test_survival_difference_at_fixed_point_in_time_test_parametric():
    df = load_waltons()
    ix = df["group"] == "miR-137"
    wf1 = WeibullFitter().fit(df.loc[ix]["T"], df.loc[ix]["E"])
    wf2 = WeibullFitter().fit(df.loc[~ix]["T"], df.loc[~ix]["E"])
    result = stats.survival_difference_at_fixed_point_in_time_test(10, wf1, wf2)
    assert result.p_value < 0.05


def test_survival_difference_at_fixed_point_in_time_test_left_censoring():
    T1 = np.random.exponential(1e-6, size=1000)
    T2 = np.random.exponential(1e-6, size=1000)
    E = T1 > T2
    T = np.maximum(T1, T2)
    kmf1 = KaplanMeierFitter().fit_left_censoring(T)
    kmf2 = KaplanMeierFitter().fit_left_censoring(2 * T)
    result = stats.survival_difference_at_fixed_point_in_time_test(T.mean(), kmf1, kmf2)
    assert result.p_value < 0.05


def test_survival_difference_at_fixed_point_in_time_test_interval_censoring():
    T1 = np.random.exponential(1e-6, size=1000)
    T2 = np.random.exponential(1e-6, size=1000)
    E = T1 > T2
    T = np.maximum(T1, T2)
    wf1 = WeibullFitter().fit_interval_censoring(T, T)
    wf2 = WeibullFitter().fit_interval_censoring(2 * T, 2 * T)
    result = stats.survival_difference_at_fixed_point_in_time_test(T.mean(), wf1, wf2)
    assert result.p_value < 0.05
