import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression

from causalml.dataset import synthetic_data
from causalml.inference.meta import BaseXLearner
from causalml.metrics.sensitivity import Sensitivity
from causalml.metrics.sensitivity import (
    SensitivityPlaceboTreatment,
    SensitivityRandomCause,
)
from causalml.metrics.sensitivity import (
    SensitivityRandomReplace,
    SensitivitySelectionBias,
)
from causalml.metrics.sensitivity import (
    one_sided,
    alignment,
    one_sided_att,
    alignment_att,
)

from .const import TREATMENT_COL, SCORE_COL, OUTCOME_COL, NUM_FEATURES


def test_Sensitivity():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )

    # generate the dataset format for sensitivity analysis
    INFERENCE_FEATURES = ["feature_" + str(i) for i in range(NUM_FEATURES)]
    df = pd.DataFrame(X, columns=INFERENCE_FEATURES)
    df[TREATMENT_COL] = treatment
    df[OUTCOME_COL] = y
    df[SCORE_COL] = e

    # calling the Base XLearner class and return the sensitivity analysis summary report
    learner = BaseXLearner(LinearRegression())
    sens = Sensitivity(
        df=df,
        inference_features=INFERENCE_FEATURES,
        p_col=SCORE_COL,
        treatment_col=TREATMENT_COL,
        outcome_col=OUTCOME_COL,
        learner=learner,
    )

    # check the sensitivity summary report
    sens_summary = sens.sensitivity_analysis(
        methods=[
            "Placebo Treatment",
            "Random Cause",
            "Subset Data",
            "Random Replace",
            "Selection Bias",
        ],
        sample_size=0.5,
    )

    print(sens_summary)


def test_SensitivityPlaceboTreatment():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )

    # generate the dataset format for sensitivity analysis
    INFERENCE_FEATURES = ["feature_" + str(i) for i in range(NUM_FEATURES)]
    df = pd.DataFrame(X, columns=INFERENCE_FEATURES)
    df[TREATMENT_COL] = treatment
    df[OUTCOME_COL] = y
    df[SCORE_COL] = e

    # calling the Base XLearner class and return the sensitivity analysis summary report
    learner = BaseXLearner(LinearRegression())
    sens = SensitivityPlaceboTreatment(
        df=df,
        inference_features=INFERENCE_FEATURES,
        p_col=SCORE_COL,
        treatment_col=TREATMENT_COL,
        outcome_col=OUTCOME_COL,
        learner=learner,
    )

    sens_summary = sens.summary(method="Random Cause")
    print(sens_summary)


def test_SensitivityRandomCause():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )

    # generate the dataset format for sensitivity analysis
    INFERENCE_FEATURES = ["feature_" + str(i) for i in range(NUM_FEATURES)]
    df = pd.DataFrame(X, columns=INFERENCE_FEATURES)
    df[TREATMENT_COL] = treatment
    df[OUTCOME_COL] = y
    df[SCORE_COL] = e

    # calling the Base XLearner class and return the sensitivity analysis summary report
    learner = BaseXLearner(LinearRegression())
    sens = SensitivityRandomCause(
        df=df,
        inference_features=INFERENCE_FEATURES,
        p_col=SCORE_COL,
        treatment_col=TREATMENT_COL,
        outcome_col=OUTCOME_COL,
        learner=learner,
    )

    sens_summary = sens.summary(method="Random Cause")
    print(sens_summary)


def test_SensitivityRandomReplace():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )

    # generate the dataset format for sensitivity analysis
    INFERENCE_FEATURES = ["feature_" + str(i) for i in range(NUM_FEATURES)]
    df = pd.DataFrame(X, columns=INFERENCE_FEATURES)
    df[TREATMENT_COL] = treatment
    df[OUTCOME_COL] = y
    df[SCORE_COL] = e

    # calling the Base XLearner class and return the sensitivity analysis summary report
    learner = BaseXLearner(LinearRegression())
    sens = SensitivityRandomReplace(
        df=df,
        inference_features=INFERENCE_FEATURES,
        p_col=SCORE_COL,
        treatment_col=TREATMENT_COL,
        outcome_col=OUTCOME_COL,
        learner=learner,
        sample_size=0.9,
        replaced_feature="feature_0",
    )

    sens_summary = sens.summary(method="Random Replace")
    print(sens_summary)


def test_SensitivitySelectionBias():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )

    # generate the dataset format for sensitivity analysis
    INFERENCE_FEATURES = ["feature_" + str(i) for i in range(NUM_FEATURES)]
    df = pd.DataFrame(X, columns=INFERENCE_FEATURES)
    df[TREATMENT_COL] = treatment
    df[OUTCOME_COL] = y
    df[SCORE_COL] = e

    # calling the Base XLearner class and return the sensitivity analysis summary report
    learner = BaseXLearner(LinearRegression())
    sens = SensitivitySelectionBias(
        df,
        INFERENCE_FEATURES,
        p_col=SCORE_COL,
        treatment_col=TREATMENT_COL,
        outcome_col=OUTCOME_COL,
        learner=learner,
        confound="alignment",
        alpha_range=None,
    )

    lls_bias_alignment, partial_rsqs_bias_alignment = sens.causalsens()
    print(lls_bias_alignment, partial_rsqs_bias_alignment)

    # Plot the results by confounding vector and plot Confidence Intervals for ATE
    sens.plot(lls_bias_alignment, ci=True)


def test_one_sided():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )
    alpha = np.quantile(y, 0.25)
    adj = one_sided(alpha, e, treatment)

    assert y.shape == adj.shape


def test_alignment():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )
    alpha = np.quantile(y, 0.25)
    adj = alignment(alpha, e, treatment)

    assert y.shape == adj.shape


def test_one_sided_att():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )
    alpha = np.quantile(y, 0.25)
    adj = one_sided_att(alpha, e, treatment)

    assert y.shape == adj.shape


def test_alignment_att():
    y, X, treatment, tau, b, e = synthetic_data(
        mode=1, n=100000, p=NUM_FEATURES, sigma=1.0
    )
    alpha = np.quantile(y, 0.25)
    adj = alignment_att(alpha, e, treatment)

    assert y.shape == adj.shape
