import numpy as np

random_state_0, random_state_1, random_state_2, random_state_3 = np.random.RandomState(42).randint(low=0, high=2**32 - 1, size=4)

y_true = np.random.RandomState(random_state_0).normal(size=1_000)

y_score_1 = np.exp(3 + y_true) + np.random.RandomState(random_state_1).normal(size=len(y_true))
y_score_2 = 3 * y_true + np.random.RandomState(random_state_2).normal(size=len(y_true))
y_score_3 = np.random.RandomState(random_state_3).normal(size=1_000)

SCORES = dict(zip(['y_score_1', 'y_score_2', 'y_score_3'], [y_score_1, y_score_2, y_score_3]))

from sklearn import metrics

for score_name, y_score in SCORES.items():
    rmse = np.sqrt(metrics.mean_squared_error(y_true=y_true, y_pred=y_score))
    print(f"RMSE for {score_name}: {rmse:6.3f}")

from scipy import stats

for score_name, y_score in SCORES.items():
    spearman = stats.spearmanr(y_true, y_score).statistic
    print(f"Spearman for {score_name}: {spearman:7.5f}")

for score_name, y_score in SCORES.items():
    kendall = stats.kendalltau(y_true, y_score).statistic
    print(f"kendall's tau for {score_name}: {kendall:7.5f}")

from sklearn import utils

def regression_roc_auc(y_true, y_score):
    """Compute the generalized ROC AUC for regression tasks.

    This function calculates the probability that the predicted values maintain
    the correct order relative to the true values, specifically for pairs where
    y_true[i] > y_true[j].

    Parameters
    ----------
    y_true : array-like of shape (n_samples,)
        True continuous target values.

    y_score : array-like of shape (n_samples,)
        Predicted continuous target values.

    Returns
    -------
    score : float
        The computed generalized ROC AUC score for regression.
    """
    y_true = utils.check_array(y_true, ensure_2d=False, dtype=None)
    y_score = utils.check_array(y_score, ensure_2d=False)

    total_pairs = 0
    correct_orderings = 0

    # Efficiently compute the metric without explicit loops
    # Create a mask for all pairs where y_true[i] > y_true[j]
    diff_matrix = y_true[:, None] - y_true[None, :]
    valid_pairs = diff_matrix > 0

    # Count total valid pairs
    total_pairs = np.sum(valid_pairs)

    if total_pairs == 0:
        # If no valid pairs, return 0.5 (equivalent to random ordering)
        return 0.5

    # Compare predictions for valid pairs
    pred_diff = y_score[:, None] - y_score[None, :]
    correct_orderings = np.sum((pred_diff > 0) & valid_pairs)

    score = correct_orderings / total_pairs
    return score

for score_name, y_score in SCORES.items():
    roc_auc = regression_roc_auc(y_true, y_score)
    print(f"ROCAUC for {score_name}: {roc_auc:7.5f}")

def ranking_curve(y_true, y_score, n_buckets=10, statistic='mean'):
    """Compute the ranking curve for a regression task.

    Calculates statistics of `y_true` values across `n_buckets`  of `y_score`
    values, ordered by the predicted scores. It can be used to  assess the
    distribution or trends of true values as a function of predicted scores.

    Parameters
    ----------
    y_true : array-like of shape (n_samples,)
        True continuous target values.

    y_score : array-like of shape (n_samples,)
        Predicted continuous target values.

    n_buckets : int, default=10
        The number of buckets to divide the sorted `y_score` values into.

    statistic : {'mean', 'median'} or callable, default='mean'
        The statistic to compute for `y_true` values in each bucket.
        - If 'mean', computes the mean of `y_true` in each bucket.
        - If 'median', computes the median of `y_true` in each bucket.
        - If callable, applies the callable function to the `y_true` values in each bucket.

    Returns
    -------
    bucket_positions : ndarray of shape (n_buckets,)
        The positions of the buckets, indexed from 1 to `n_buckets`.

    bucket_values : ndarray of shape (n_buckets,)
        The computed statistic for `y_true` values in each bucket.
    """
    sorted_indices = np.argsort(y_score)
    y_true_sorted = y_true[sorted_indices]

    bucket_edges = np.linspace(0, len(y_true), n_buckets + 1, dtype=int)
    bucket_values = []

    if statistic == 'mean':
        stat_func = np.mean
    elif statistic == 'median':
        stat_func = np.median
    elif callable(statistic):
        stat_func = statistic
    else:
        raise ValueError

    for i in range(n_buckets):
        start, end = bucket_edges[i], bucket_edges[i + 1]
        bin_values = y_true_sorted[start:end]
        if len(bin_values) > 0:
            bucket_stat = stat_func(bin_values)
        else:
            bucket_stat = np.nan
        bucket_values.append(bucket_stat)

    bucket_positions = np.arange(1, n_buckets + 1)
    return bucket_positions, bucket_values

N_BUCKETS = 10
ordering_curve_dict = {}

for score_name, y_score in SCORES.items():
    bins, ordering_curve = ranking_curve(y_true, y_score, N_BUCKETS, 'mean')
    ordering_curve_dict[score_name] = ordering_curve

import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(5, 3), dpi=130)

ax.hlines(y_true.mean(), min(bins), max(bins), label='random ordering', colors='k', alpha=0.5)
for score_name, ordering_curve in ordering_curve_dict.items():
    ax.plot(bins, ordering_curve, '-o', markeredgecolor='k', markeredgewidth=0.5, label=score_name)

ax.set_xticks(bins)
ax.legend()
ax.set_ylabel("Mean of y for each bucket")
ax.set_xlabel("buckets")
plt.tight_layout()

def last_bucket_ordering_curve(y_true, y_score, n_buckets=10, statistic='mean'):
    _, ordering_curve = ranking_curve(y_true, y_score, n_buckets=n_buckets, statistic=statistic)
    *_, last_bucket = ordering_curve
    return last_bucket

def first_bucket_ordering_curve(y_true, y_score, n_buckets=10, statistic='mean'):
    _, ordering_curve = ranking_curve(y_true, y_score, n_buckets=n_buckets, statistic=statistic)
    first_bucket, *_ = ordering_curve
    return first_bucket

def diff_bucket_ordering_curve(y_true, y_score, n_buckets=10, statistic='mean'):
    _, ordering_curve = ranking_curve(y_true, y_score, n_buckets=n_buckets, statistic=statistic)
    first_bucket, *_, last_bucket = ordering_curve
    return last_bucket - first_bucket

from sklearn import linear_model

def linear_regression_coefficient_ordering_curve(y_true, y_score, n_buckets=10, statistic='mean'):
    _, ordering_curve = ranking_curve(y_true, y_score, n_buckets=n_buckets, statistic=statistic)
    x_values = np.arange(n_buckets).reshape(-1, 1)
    y_values = np.array(ordering_curve).reshape(-1, 1)

    model = linear_model.LinearRegression().fit(x_values, y_values)
    return model.coef_[0][0]

for score_name, y_score in SCORES.items():
    last_bucket = last_bucket_ordering_curve(y_true, y_score)
    print(f"Last bucket for {score_name}: {last_bucket:7.5f}")

for score_name, y_score in SCORES.items():
    first_bucket = first_bucket_ordering_curve(y_true, y_score)
    print(f"First bucket for {score_name}: {first_bucket:8.5f}")

for score_name, y_score in SCORES.items():
    diff_bucket = diff_bucket_ordering_curve(y_true, y_score)
    print(f"Diff bucket for {score_name}: {diff_bucket:7.5f}")

for score_name, y_score in SCORES.items():
    lr_bucket = linear_regression_coefficient_ordering_curve(y_true, y_score)
    print(f"Linear regression coefficient for {score_name}: {lr_bucket:7.5f}")

import numpy as np
from sklearn.datasets import fetch_california_housing
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    *fetch_california_housing(return_X_y=True),
    test_size=0.33,
    random_state=42,
)

lr = LinearRegression().fit(X_train, y_train)

def evaluate_model(y_true, y_pred):
    print(f"MSE: {mean_squared_error(y_true, y_pred)}")
    print(f"R^2: {r2_score(y_true, y_pred)}")
    
y_pred_lr =  lr.predict(X_test)
evaluate_model(y_test, y_pred_lr)

y_mean_test = y_test.mean() * np.ones_like(y_test)
evaluate_model(y_test, y_mean_test)

y_pred_terrible_model = np.zeros_like(y_test)
evaluate_model(y_test, y_pred_terrible_model)

y_mean_train = y_train.mean() * np.ones_like(y_test)
evaluate_model(y_test, y_mean_train)

from sklearn.metrics import mean_absolute_error

def r2_score_mae(y_true, y_pred, *args, **kwargs):
    mae_model = mean_absolute_error(y_true=y_true, y_pred=y_pred, *args, **kwargs)
    y_median_true = np.median(y_true) * np.ones_like(y_true)
    mae_median = mean_absolute_error(
        y_true=y_true, y_pred=y_median_true, *args, **kwargs
    )
    return 1 - mae_model / mae_median

def evaluate_model_mae(y_true, y_pred):
    print(f"MAE: {mean_absolute_error(y_true, y_pred)}")
    print(f"R^2_MAE: {r2_score_mae(y_true, y_pred)}")

evaluate_model_mae(y_test, y_pred_lr)

y_median_test = np.median(y_test) * np.ones_like(y_test)
evaluate_model_mae(y_test, y_median_test)

def adapted_sigmoid(x):
    return 0.05 + 0.9 / (1 + np.exp(-x))

def func_0(Z):
    return 10 * np.sin(Z)

def func_1(Z):
    return 10 * np.cos(Z)

def generate_data(size=100, obs=True, random_state=None):
    rs = np.random.RandomState(random_state).randint(
        0, 2**32 - 1, dtype=np.int64, size=4
    )

    Z_obs = np.random.RandomState(rs[0]).uniform(low=-np.pi, high=np.pi, size=size)

    def g_T_noised(Z):
        return (
            np.random.RandomState(rs[1])
            .binomial(n=1, p=adapted_sigmoid(Z))
            .astype(bool)
        )

    T_obs = g_T_noised(Z_obs)

    noise = np.random.RandomState(rs[3]).normal(size=size)

    def g_Y(T, Z, noise):
        return (
            np.select(condlist=[T], choicelist=[func_1(Z)], default=func_0(Z))
            + 0.5 * (1 + T * np.abs(Z)) * noise
        )

    Y_obs = g_Y(T_obs, Z_obs, noise)
    Y_cf = g_Y(~T_obs, Z_obs, noise)

    def generate_df(T, Z, Y):
        return pd.DataFrame(
            np.vstack([T.astype(int), Z, Y]).T,
            columns=["treatment", "confounder", "target"],
        )

    df_obs = generate_df(T_obs, Z_obs, Y_obs)
    df_cf = generate_df(~T_obs, Z_obs, Y_cf)

    return df_obs, df_cf

df_obs, df_cf = generate_data(size=50_000, obs=True, random_state=42)

from sklearn.model_selection import train_test_split

df_train, df_test = train_test_split(
    df_obs.assign(target_cf=df_cf.target),
    test_size=0.2,
    random_state=42,
)

df_train_t0 = df_train.query("treatment == 0")
df_train_t1 = df_train.query("treatment == 1")

def return_TZ_y(df, backdoor_set_list):
    return df.filter(backdoor_set_list), np.array(df.target)

backdoor_set = ["confounder"]

XZ_train_t0, y_train_t0 = return_TZ_y(df_train_t0, backdoor_set)
XZ_train_t1, y_train_t1 = return_TZ_y(df_train_t1, backdoor_set)

XZ_test, y_test = return_TZ_y(df_test, backdoor_set)

from sklearn.metrics import roc_auc_score
from sklearn.linear_model import LogisticRegression

positivity_assumption_check_estimator = LogisticRegression(
    random_state=42,
).fit(df_train.drop(columns=["treatment", "target", "target_cf"]), df_train.treatment)

roc_auc_score(
    df_test.treatment, positivity_assumption_check_estimator.predict_proba(XZ_test)[:, 1]
)

from calibration_stuff import calibration_curve

probs = positivity_assumption_check_estimator.predict_proba(XZ_test)[:, 1]
prob_true, prob_pred, size_bin = calibration_curve(df_test.treatment, probs, n_bins=10)

fig, ax = plt.subplots(ncols=2, figsize=(10, 3))
ax[0].plot([0, 1], "--")
ax[0].scatter(prob_true, prob_pred, s=(0.1 * size_bin).astype(int), edgecolor="k")
ax[0].set_xlabel("True probability of bin")
ax[0].set_ylabel("Mean predicted probability of bin")
ax[1].hist(
    probs, bins=np.linspace(0, 1, 21), weights=np.ones_like(probs) / probs.shape[0]
)
ax[1].set_xlabel("Histogram of predicted probability")
plt.tight_layout()

from functools import partial
from lightgbm import LGBMRegressor
from scipy.stats import loguniform
from sklearn.base import BaseEstimator
from sklearn.metrics import make_scorer
from sklearn.model_selection import RandomizedSearchCV, KFold
from sklearn.utils.validation import check_X_y, check_is_fitted, _check_sample_weight
from statsmodels.stats.weightstats import DescrStatsW

class ConformalizedQuantileRegression(BaseEstimator):
    """
    Conformalized Quantile Regression with LGBMRegressor.

    This estimator provides prediction intervals for one dimension
    regression tasks by using CQR with LightGBM.

    Parameters
    ----------
    alpha : float, default=0.05
        Determines the size of the prediction interval. For example,
        alpha=0.05 results in a 95% coverage prediction interval.

    calibration_size : float, default=0.2
        The proportion of the dataset to be used for the calibration set
        which computes the conformity scores.

    random_state : int, RandomState instance or None, default=None
        Controls the randomness for reproducibility.

    n_iter_cv : int, default=10
        Number of parameter settings that are sampled in RandomizedSearchCV
        for the LightGBM model during fit.
    """

    def __init__(
        self, alpha=0.05, calibration_size=0.2, random_state=None, n_iter_cv=10
    ):
        self.alpha = alpha
        self.calibration_size = calibration_size
        self.random_state = random_state
        self.n_iter_cv = n_iter_cv

    def _quantile_loss(self, y_true, y_pred, quantile=None, sample_weights=None):
        weighted_errors = (y_true - y_pred) * (quantile - (y_true < y_pred))
        if sample_weights is not None:
            weighted_errors *= sample_weights
        return np.mean(weighted_errors)

    def _return_quantile_model(self, quantile):
        quantile_scorer = make_scorer(
            partial(self._quantile_loss, quantile=quantile), greater_is_better=False
        )

        return RandomizedSearchCV(
            estimator=LGBMRegressor(
                random_state=self.random_state, objective="quantile", alpha=quantile
            ),
            cv=KFold(shuffle=True, random_state=self.random_state),
            param_distributions={
                "learning_rate": loguniform.rvs(
                    random_state=self.random_state, a=0.0001, b=1, size=1000
                ),
                "n_estimators": [50, 100, 200],
                "num_leaves": [31, 63, 127],
            },
            scoring=quantile_scorer,
            n_iter=self.n_iter_cv,
            random_state=self.random_state,
            n_jobs=-1,
        )

    def fit(self, X, y, sample_weight=None):
        X, y = check_X_y(X, y)
        sample_weight = _check_sample_weight(sample_weight, X)

        (
            X_train,
            X_cal,
            y_train,
            y_cal,
            sample_weight_train,
            sample_weight_cal,
        ) = train_test_split(
            X,
            y,
            sample_weight,
            test_size=self.calibration_size,
            random_state=self.random_state,
        )

        self.model_lower_ = self._return_quantile_model(quantile=self.alpha / 2).fit(
            X_train, y_train, sample_weight=sample_weight_train
        )
        self.model_upper_ = self._return_quantile_model(
            quantile=1 - self.alpha / 2
        ).fit(X_train, y_train, sample_weight=sample_weight_train)

        self.y_cal_conformity_scores_ = np.maximum(
            self.model_lower_.predict(X_cal) - y_cal,
            y_cal - self.model_upper_.predict(X_cal),
        )
        wq = DescrStatsW(data=self.y_cal_conformity_scores_, weights=sample_weight_cal)
        self.quantile_conformity_scores_ = wq.quantile(
            probs=1 - self.alpha, return_pandas=False
        )[0]

        return self

    def predict(self, X):
        """
        Predicts conformalized quantile regression intervals for X.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            The input samples.

        Returns
        -------
        y_test_interval_pred_cqr : ndarray of shape (n_samples, 2)
            Returns the predicted lower and upper bound for each sample in X.
        """
        check_is_fitted(self)
        y_test_interval_pred_cqr = np.column_stack(
            [
                self.model_lower_.predict(X) - self.quantile_conformity_scores_,
                self.model_upper_.predict(X) + self.quantile_conformity_scores_,
            ]
        )
        return y_test_interval_pred_cqr

model_t0 = ConformalizedQuantileRegression(
    random_state=42, alpha=0.05, n_iter_cv=30
).fit(XZ_train_t0, y_train_t0)
y_test_interval_pred_cqr_t0 = model_t0.predict(XZ_test)

model_t1 = ConformalizedQuantileRegression(
    random_state=42, alpha=0.05, n_iter_cv=30
).fit(XZ_train_t1, y_train_t1)
y_test_interval_pred_cqr_t1 = model_t1.predict(XZ_test)

def return_sample_weight_treatment_i(df_train, df_test):
    df_ood_ti = pd.concat(
        [
            df.assign(train_or_test=j)
            for j, df in enumerate(
                [
                    df_train.drop(columns=["treatment", "target_cf"]),
                    df_test.drop(columns=["treatment", "target_cf"]),
                ]
            )
        ]
    )

    ood_sample_correction_ti = LogisticRegression(
        random_state=42,
    ).fit(df_ood_ti.drop(columns=["train_or_test"]), df_ood_ti.train_or_test)

    roc = roc_auc_score(
        df_ood_ti.train_or_test,
        ood_sample_correction_ti.predict_proba(
            df_ood_ti.drop(columns=["train_or_test"])
        )[:, 1],
    )

    probs = ood_sample_correction_ti.predict_proba(
        df_train.drop(columns=["treatment", "target_cf"])
    )
    # Equivalent to `probs[:, 1]/probs[:, 0]`.
    sample_weights_ti = 1 / probs[:, 0] - 1

    return roc, sample_weights_ti

_, sw_0 = return_sample_weight_treatment_i(df_train=df_train_t0, df_test=df_test)

df_val = (
    df_test.assign(pred_lower_t_0=y_test_interval_pred_cqr_t0[:, 0])
    .assign(pred_upper_t_0=y_test_interval_pred_cqr_t0[:, 1])
    .assign(ic_size_t_0=lambda df_: df_.pred_upper_t_0 - df_.pred_lower_t_0)
    .assign(pred_lower_t_1=y_test_interval_pred_cqr_t1[:, 0])
    .assign(pred_upper_t_1=y_test_interval_pred_cqr_t1[:, 1])
    .assign(ic_size_t_1=lambda df_: df_.pred_upper_t_1 - df_.pred_lower_t_1)
    .assign(
        prob=lambda df_: positivity_assumption_check_estimator.predict_proba(
            df_.filter(backdoor_set)
        )[:, 1]
    )
    .assign(prob_cut=lambda df_: pd.cut(df_.prob, bins=np.linspace(0, 1, 6)))
    .assign(
        coverage=lambda df_: np.select(
            condlist=[df_.treatment == 0],
            choicelist=[
                (df_.target > df_.pred_lower_t_0) & (df_.target < df_.pred_upper_t_0)
            ],
            default=(df_.target > df_.pred_lower_t_1)
            & (df_.target < df_.pred_upper_t_1),
        )
    )
    .assign(
        coverage_cf=lambda df_: np.select(
            condlist=[df_.treatment != 0],
            choicelist=[
                (df_.target_cf > df_.pred_lower_t_0)
                & (df_.target_cf < df_.pred_upper_t_0)
            ],
            default=(df_.target_cf > df_.pred_lower_t_1)
            & (df_.target_cf < df_.pred_upper_t_1),
        )
    )
)

df_val.coverage.mean()

from scipy.stats import bootstrap

def bootstrap_ci(x, ci=0.95):
    boot = bootstrap((x,), np.mean, confidence_level=ci)
    return np.round(boot.confidence_interval, 5)

df_val_cond_aux1 = (
    df_val.groupby(["prob_cut", "treatment"])
    .coverage.apply(bootstrap_ci)
    .to_frame()
    .rename(columns={"coverage": "coverage_confidence_interval"})
)

df_val_cond_aux2 = (
    df_val.groupby(["prob_cut", "treatment"])
    .coverage_cf.apply(bootstrap_ci)
    .to_frame()
    .rename(columns={"coverage_cf": "coverage_cf_confidence_interval"})
)

df_val_cond_aux3 = (
    df_val.groupby(["prob_cut", "treatment"])
    .agg(
        {
            "coverage": np.mean,
            "coverage_cf": np.mean,
            "ic_size_t_0": np.mean,
            "ic_size_t_1": np.mean,
        }
    )
    .rename(columns=lambda col: col + "_mean")
)

pd.concat(
    [df_val_cond_aux1, df_val_cond_aux2, df_val_cond_aux3], axis=1
).reset_index().sort_values(["treatment", "prob_cut"])

probs_coverage = (
    LogisticRegression()
    .fit(df_val.filter(["treatment", "confounder"]), df_val.coverage.astype(int))
    .predict_proba(df_val.filter(["treatment", "confounder"]))[:, 1]
)

roc_auc_score(df_val.coverage.astype(int), probs_coverage)

min(probs_coverage), max(probs_coverage)

df_val_cate = (
    df_val.assign(
        cate_actual=lambda df_: np.select(
            condlist=[(df_.treatment == 0)],
            choicelist=[df_.target_cf - df_.target],
            default=[df_.target - df_.target_cf],
        )[0]
    )
    .assign(cate_ci_lower=lambda df_: df_.pred_lower_t_1 - df_.pred_upper_t_0)
    .assign(cate_ci_upper=lambda df_: df_.pred_upper_t_1 - df_.pred_lower_t_0)
    .assign(cate_ci_size=lambda df_: df_.cate_ci_upper - df_.cate_ci_lower)
    .assign(
        coverage_cate=lambda df_: (df_.cate_actual > df_.cate_ci_lower)
        & (df_.cate_actual < df_.cate_ci_upper)
    )
)

fig, ax = plt.subplots(ncols=3, figsize=(9, 2))
aux_hist = np.hstack([df_val.ic_size_t_0, df_val.ic_size_t_1])
min_hist, max_hist = np.min(aux_hist), np.max(aux_hist)
ax[0].hist(
    df_val.ic_size_t_0,
    bins=np.linspace(min_hist, max_hist, 16),
    weights=np.ones_like(df_val.ic_size_t_0) / df_val.shape[0],
)
ax[1].hist(
    df_val.ic_size_t_1,
    bins=np.linspace(min_hist, max_hist, 16),
    weights=np.ones_like(df_val.ic_size_t_1) / df_val.shape[0],
)
ax[2].hist(
    df_val_cate.cate_ci_size,
    bins=16,
    weights=np.ones_like(df_val_cate.cate_ci_size) / df_val_cate.shape[0],
)
ax[0].set_title(
    "Histogram of interval size for $\mathbb{E}(Y | do(T=0), Z=z)$", fontsize="medium"
)
ax[1].set_title(
    "Histogram of interval size for $\mathbb{E}(Y | do(T=1), Z=z)$", fontsize="medium"
)
ax[2].set_title("Histogram of interval size for CATE(Z=z)", fontsize="medium")
plt.tight_layout()

df_val_cate.coverage_cate.mean()

confounder_plot = np.linspace(XZ_test.confounder.min(), XZ_test.confounder.max(), 10_001)
ci_t1_plot = model_t1.predict(confounder_plot.reshape(-1, 1)).T
ci_t0_plot = model_t0.predict(confounder_plot.reshape(-1, 1)).T
ci_cate_plot = ci_t1_plot - ci_t0_plot[::-1,]

fig, ax = plt.subplots(figsize=(8, 4))

ax.plot(
    confounder_plot,
    func_0(confounder_plot) + 1.96 * 0.5,
    alpha=0.5,
    c="C0",
    label="Real confidence interval for $\mathbb{E}(Y | do(T=0), Z=z)$",
)
ax.plot(confounder_plot, func_0(confounder_plot) - 1.96 * 0.5, alpha=0.5, c="C0")

ax.plot(
    confounder_plot,
    func_1(confounder_plot) + 1.96 * (0.5 * (1 + np.abs(confounder_plot))),
    alpha=0.5,
    c="C1",
    label="Real confidence interval for $\mathbb{E}(Y | do(T=1), Z=z)$",
)
ax.plot(
    confounder_plot,
    func_1(confounder_plot) - 1.96 * (0.5 * (1 + np.abs(confounder_plot))),
    alpha=0.5,
    c="C1",
)

# Variance of CATE(Z=z) is 0.5 * |z| because the term
# related to 1 u_Y is annulled when we do
# \mathbb{E}(g_Y(u_Y, z, 1)) - \mathbb{E}(g_Y(u_Y, z, 0)).
ax.plot(
    confounder_plot,
    func_1(confounder_plot)
    - func_0(confounder_plot)
    + 1.96 * (0.5 * (np.abs(confounder_plot))),
    alpha=0.5,
    c="C2",
    label="Confidence interval for CATE(Z=z)",
)
ax.plot(
    confounder_plot,
    func_1(confounder_plot)
    - func_0(confounder_plot)
    - 1.96 * (0.5 * (np.abs(confounder_plot))),
    alpha=0.5,
    c="C2",
)

ax.fill_between(
    confounder_plot,
    *ci_t0_plot,
    alpha=0.5,
    label="Prediction interval for $\mathbb{E}(Y | do(T=0), Z=z)$",
    color="C0",
)
ax.fill_between(
    confounder_plot,
    *ci_t1_plot,
    alpha=0.5,
    label="Prediction interval for $\mathbb{E}(Y | do(T=1), Z=z)$",
    color="C1",
)
ax.fill_between(
    confounder_plot,
    *ci_cate_plot,
    alpha=0.5,
    label="Prediction interval for CATE(Z=z)",
    color="C2",
)

ax.set_xlabel("z")
ax.legend()
plt.tight_layout()

def mean_function(X):
    return np.sin(2 * np.pi * X)

def deviation_function(X):
    return 0.3 + 0.25 * mean_function(X)

def generate_data_with_normal_noise(
    mean_generator, deviation_generator, size=5_000, random_state=None
):
    def normal_noise_generator(X, deviation_generator, random_state=None):
        noise = np.random.RandomState(random_state).normal(
            loc=0, scale=deviation_generator(X)
        )
        return noise

    rs = np.random.RandomState(random_state).randint(
        0, 2**32 - 1, dtype=np.int64, size=2
    )
    X = np.random.RandomState(rs[0]).uniform(size=size)
    y_pred = mean_generator(X=X)
    noise = normal_noise_generator(
        X=X, deviation_generator=deviation_generator, random_state=rs[1]
    )
    y_pred_noisy = y_pred + noise

    return X, y_pred_noisy

X, y = generate_data_with_normal_noise(
    mean_generator=mean_function,
    deviation_generator=deviation_function,
    random_state=42,
)

x_grid = np.linspace(0, 1, 1000)

fig, ax = plt.subplots(figsize=(8, 3))
ax.plot(x_grid, mean_function(x_grid), color="C0", label="Mean function")
ax.fill_between(
    x_grid,
    mean_function(x_grid) - 1.96 * deviation_function(x_grid),
    mean_function(x_grid) + 1.96 * deviation_function(x_grid),
    color="C0",
    alpha=0.2,
    label="95% confidence interval given x",
)
ax.scatter(X, y, s=2, color="C0", alpha=0.2)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title("Scatter plot of generated data")
ax.legend()
plt.tight_layout()

from scipy.stats import rv_histogram, norm

hist = np.histogram(y, bins=np.linspace(-1.5, 3, 51))
hist_dist = rv_histogram(hist)

def plot_conditional_y_using_near_data(ax, x_value, X, y, c, eps=0.05):
    ax.plot(
        y_grid_refined,
        norm(loc=mean_function(x_value), scale=deviation_function(x_value)).pdf(
            y_grid_refined
        ),
        "--",
        color=c,
        label=f"real $f(y | x = {x_value})$",
    )
    ax.hist(
        y[(X < x_value + eps) & (X > x_value - eps)],
        alpha=0.3,
        bins=y_grid,
        density=True,
        color=c,
        label=f"estimated $f(y | x = {x_value})$ using near data",
    )

min_y, max_y = min(y), max(y)
y_grid = np.linspace(min_y, max_y, 20)
y_grid_refined = np.linspace(min_y, max_y, 1000)

fig, ax = plt.subplots(ncols=2, figsize=(10, 3))
ax[0].bar(y_grid, hist_dist.pdf(y_grid), label="estimated $f(y)$")
ax[0].set_title("Density of y")
ax[0].set_xlabel("y")
ax[0].legend()

plot_conditional_y_using_near_data(ax=ax[1], x_value=0.2, X=X, y=y, c="C1")
plot_conditional_y_using_near_data(ax=ax[1], x_value=0.6, X=X, y=y, c="C2")
ax[1].set_title("Conditional density of y given X=x")
ax[1].set_xlabel("y")
ax[1].legend()
plt.tight_layout()

from sklearn.neighbors import NearestNeighbors, KernelDensity

x_value = 0.2
knn = NearestNeighbors(n_neighbors=100).fit(X.reshape(-1, 1))
_, ind_x_value = knn.kneighbors([[x_value]])

kde = KernelDensity(kernel="gaussian", bandwidth="scott").fit(
    y[ind_x_value].reshape(-1, 1)
)

fig, ax = plt.subplots(figsize=(5, 3))
ax.plot(
    y_grid_refined,
    norm(loc=mean_function(x_value), scale=deviation_function(x_value)).pdf(
        y_grid_refined
    ),
    "--",
    color="C0",
    label=f"real $f(y | x = {x_value})$",
)
ax.hist(
    y[ind_x_value].ravel(),
    alpha=0.3,
    bins=y_grid,
    density=True,
    color="C0",
    label=f"estimated $f(y | x = {x_value})$ using a histogram of nearest neighbors",
)
ax.plot(
    y_grid_refined,
    np.exp(kde.score_samples(y_grid_refined.reshape(-1, 1))),
    color="C0",
    label=f"estimated $f(y | x = {x_value})$ using a kde with nearest neighbors",
)
ax.set_title("Conditional density of y given X=x")
ax.set_xlabel("y")
ax.set_ylim(0, 1.3)
ax.legend()
plt.tight_layout()

from sklearn.base import BaseEstimator, clone

class ConditionalNearestNeighborsKDE(BaseEstimator):
    """Conditional Kernel Density Estimation using nearest neighbors.

    This class implements a Conditional Kernel Density Estimation by applying
    the Kernel Density Estimation algorithm after a nearest neighbors search.

    It allows the use of user-specified nearest neighbor and kernel density
    estimators or, if not provided, defaults will be used.

    Parameters
    ----------
    nn_estimator : NearestNeighbors instance, default=None
        A pre-configured instance of a `~sklearn.neighbors.NearestNeighbors` class
        to use for finding nearest neighbors. If not specified, a
        `~sklearn.neighbors.NearestNeighbors` instance with `n_neighbors=100`
        will be used.

    kde_estimator : KernelDensity instance, default=None
        A pre-configured instance of a `~sklearn.neighbors.KernelDensity` class
        to use for estimating the kernel density. If not specified, a
        `~sklearn.neighbors.KernelDensity` instance with `bandwidth="scott"`
        will be used.
    """

    def __init__(self, nn_estimator=None, kde_estimator=None):
        self.nn_estimator = nn_estimator
        self.kde_estimator = kde_estimator

    def fit(self, X, y=None):
        if self.nn_estimator is None:
            self.nn_estimator_ = NearestNeighbors(n_neighbors=100)
        else:
            self.nn_estimator_ = clone(self.nn_estimator)
        self.nn_estimator_.fit(X, y)
        self.y_train_ = y
        return self

    def predict(self, X):
        """Predict the conditional density estimation of new samples.

        The predicted density of the target for each sample in X is returned.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Vector to be estimated, where `n_samples` is the number of samples
            and `n_features` is the number of features.

        Returns
        -------
        kernel_density_list : list of len n_samples of KernelDensity instances
            Estimated conditional density estimations in the form of
            `~sklearn.neighbors.KernelDensity` instances.
        """
        _, ind_X = self.nn_estimator_.kneighbors(X)
        if self.kde_estimator is None:
            kernel_density_list = [
                KernelDensity(bandwidth="scott").fit(self.y_train_[ind].reshape(-1, 1))
                for ind in ind_X
            ]
        else:
            kernel_density_list = [
                clone(self.kde_estimator).fit(self.y_train_[ind].reshape(-1, 1))
                for ind in ind_X
            ]
        return kernel_density_list

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.33, random_state=42
)

ckde = ConditionalNearestNeighborsKDE().fit(X_train.reshape(-1, 1), y_train)
ckde_preds = ckde.predict(X_test.reshape(-1, 1))

from sklearn.metrics import auc, make_scorer
from joblib import Parallel, delayed

def squared_loss(y_true, cde_preds, y_grid, n_jobs=-1):
    """
    Average squared loss between the true conditional density and predicted one.

    This method can be used to assess the quality of the conditional probability
    density function fit.

    Parameters
    ----------
    y_true : array-like of shape (n_samples,)
        The true values of y for each sample.

    cde_preds : list of len n_samples of KernelDensity instances
        The predicted conditional densitys. Each instance should be a fitted
        KernelDensity instance.

    y_grid : array-like of shape (n_samples,)
        The grid of y values used for computing the area under the curve (AUC)
        for the squared probability density function.

    n_jobs : int, optional
        The number of jobs to run in parallel. '-1' means using all processors.

    Returns
    -------
    average_squared_loss: float
        The average squared loss between the true and predicted conditional
        probability density functions. Note that it is always off by C.
    """

    def _compute_individual_loss(y_, cde_pred):
        # The score_samples and score methods returns stuff on log scale,
        # so we have to exp it.
        squared_auc = auc(
            y_grid, np.exp(cde_pred.score_samples(y_grid.reshape(-1, 1))) ** 2
        )
        expected_value = np.exp(cde_pred.score([[y_]]))
        return squared_auc - 2 * expected_value

    individual_squared_loss = Parallel(n_jobs=n_jobs)(
        delayed(_compute_individual_loss)(y_, cde_pred)
        for y_, cde_pred in zip(y_true, cde_preds)
    )

    average_squared_loss = sum(individual_squared_loss) / len(y_true)
    return average_squared_loss

squared_loss(y_test, ckde_preds, np.linspace(-5, 5, 1000))

squared_loss(
    y_test,
    len(y_test) * [KernelDensity(bandwidth="scott").fit(y_train.reshape(-1, 1))],
    np.linspace(-5, 5, 1000),
)

from functools import partial
from sklearn.model_selection import GridSearchCV

squared_loss_score = make_scorer(
    partial(squared_loss, y_grid=np.linspace(-5, 5, 1000)), greater_is_better=False
)
param_grid = {
    "nn_estimator": [
        NearestNeighbors(n_neighbors=n_neighbors) for n_neighbors in [100, 500, 1000]
    ],
}
gs = GridSearchCV(
    ConditionalNearestNeighborsKDE(), param_grid=param_grid, scoring=squared_loss_score
).fit(X_train.reshape(-1, 1), y_train)

(
    pd.DataFrame(gs.cv_results_)
    .filter(
        ["param_nn_estimator", "mean_score_time", "mean_test_score", "std_test_score"]
    )
    .sort_values(by="mean_test_score", ascending=False)
    .reset_index(drop=True)
)

squared_loss(
    y_test, gs.best_estimator_.predict(X_test.reshape(-1, 1)), np.linspace(-5, 5, 1000)
)

from sklearn.neighbors._base import NeighborsBase
from sklearn.ensemble import RandomForestRegressor

class LeafNeighbors(NeighborsBase):
    """Neighbors search using leaf nodes coincidence in a tree ensemble as a
    similarity measure.

    This class implements a supervised neighbor search using the leaves of an
    ensemble tree estimator as a measure of distance. Examples that occur
    simultaneously in several leaves are naturally close in variables relevant
    to the target.

    Parameters
    ----------
    tree_ensemble_estimator : ForestRegressor instance, default=None
        The ensemble tree estimator to use. If None, a
        `~sklearn.ensemble.RandomForestRegressor` with `max_depth=10` will be
        used.

    n_neighbors : int, default=5
        Number of neighbors to use in the neighbor-based learning method.

    random_state : int, RandomState instance or None, default=None
        Controls the randomness of the ensemble tree estimator. Pass an int
        for reproducible output across multiple function calls.
    """

    def __init__(self, tree_ensemble_estimator=None, n_neighbors=5, random_state=None):
        self.tree_ensemble_estimator = tree_ensemble_estimator
        self.n_neighbors = n_neighbors
        self.random_state = random_state

    def fit(self, X, y=None):
        if self.tree_ensemble_estimator is None:
            self.tree_ensemble_estimator = RandomForestRegressor(
                max_depth=10, random_state=self.random_state
            )
        else:
            self.tree_ensemble_estimator = clone(self.tree_ensemble_estimator)

        self.nn_estimator_ = NearestNeighbors(
            n_neighbors=self.n_neighbors, metric="hamming"
        )

        self.tree_ensemble_estimator.fit(X, y)
        leafs_X = self.tree_ensemble_estimator.apply(X)
        self.nn_estimator_.fit(leafs_X)
        return self

    def kneighbors(self, X):
        leafs_X = self.tree_ensemble_estimator.apply(X)
        return self.nn_estimator_.kneighbors(leafs_X)

    def radius_neighbors(self, X):
        leafs_X = self.tree_ensemble_estimator.apply(X)
        return self.nn_estimator_.radius_neighbors(leafs_X)

crfkde = ConditionalNearestNeighborsKDE(
    nn_estimator=LeafNeighbors(n_neighbors=100)
).fit(X_train.reshape(-1, 1), y_train)
crfkde_preds = crfkde.predict(X_test.reshape(-1, 1))

squared_loss(y_test, crfkde_preds, np.linspace(-5, 5, 1000))

from flexcode.regression_models import RandomForest
from flexcode import FlexCodeModel

flexcode_model = FlexCodeModel(
    RandomForest,
    basis_system="cosine",
    max_basis=31,
    regression_params={"max_depth": 5, "n_estimators": 100},
)
flexcode_model.fit(X_train, y_train)

cdes, y_grid_flexcode = flexcode_model.predict(X_test, n_grid=400)
y_grid_flexcode = y_grid_flexcode.reshape(-1)

fig, ax = plt.subplots(figsize=(8, 3))
for c, sample_index in enumerate(np.random.RandomState(13).choice(len(X_test), size=3)):
    x_value = np.round(X_test[sample_index], 4)
    ax.plot(
        y_grid_refined,
        norm(loc=mean_function(x_value), scale=deviation_function(x_value)).pdf(
            y_grid_refined
        ),
        "--",
        color=f"C{c}",
        label=f"real $f(y | x = {x_value})$",
    )
    ax.plot(
        y_grid_flexcode,
        cdes[sample_index],
        color=f"C{c}",
        label=f"estimated $f(y | x = {x_value})$ using flexcode",
    )
ax.set_title("Conditional density of y given X=x")
ax.set_xlabel("y")
ax.legend()
plt.tight_layout()

class FlexCode_return_to_DensityEstimator:
    def __init__(self, y_grid, pdf_values):
        self.y_grid = y_grid
        self.pdf_values = pdf_values
        self.density = rv_histogram(
            (pdf_values, np.hstack([y_grid, [y_grid[-1] + y_grid[-1] - y_grid[-2]]]))
        )

    def score_samples(self, X):
        return np.log(self.density.pdf(X))

    def score(self, X):
        return np.sum(self.score_samples(X))

density_estimation_preds_flexcode = [
    FlexCode_return_to_DensityEstimator(y_grid=y_grid_flexcode, pdf_values=cde)
    for cde in cdes
]
squared_loss(y_test, density_estimation_preds_flexcode, np.linspace(-5, 5, 1000))

from sklearn.datasets import fetch_california_housing

X_california, y_california = fetch_california_housing(return_X_y=True)
(
    X_california_train,
    X_california_test,
    y_california_train,
    y_california_test,
) = train_test_split(X_california, y_california, test_size=0.33, random_state=42)
print(f"X dimension: {X_california.shape[1]}")

ckde_california = ConditionalNearestNeighborsKDE().fit(
    X_california_train, y_california_train
)
ckde_california_preds = ckde_california.predict(X_california_test)

squared_loss(y_california_test, ckde_california_preds, np.linspace(0, 6, 1000))

crfkde_california = ConditionalNearestNeighborsKDE(
    nn_estimator=LeafNeighbors(n_neighbors=100)
).fit(X_california_train, y_california_train)
crfkde_california_preds = crfkde_california.predict(X_california_test)

squared_loss(y_california_test, crfkde_california_preds, np.linspace(0, 6, 1000))

model_california = FlexCodeModel(
    RandomForest, max_basis=31, regression_params={"max_depth": 10, "n_estimators": 100}
)
model_california.fit(X_california_train, y_california_train)

cdes_california, y_grid_california = model_california.predict(
    X_california_test, n_grid=400
)
y_grid_california = y_grid_california.reshape(-1)
density_estimation_preds_flexcode_california = [
    FlexCode_return_to_DensityEstimator(y_grid=y_grid_california, pdf_values=cde)
    for cde in cdes_california
]

squared_loss(
    y_california_test,
    density_estimation_preds_flexcode_california,
    np.linspace(0, 6, 1000),
)

def predict(self, X):
    ...
    proba = self.predict_proba(X)
    ...
    return self.classes_.take(np.argmax(proba, axis=1), axis=0)

from imblearn.datasets import fetch_datasets

dataset = fetch_datasets()["coil_2000"]
X, y = dataset.data, (dataset.target==1).astype(int)

print(f"Percentage of y=1 is {np.round(y.mean(), 5)*100}%.")
print(f"Number of rows is {X.shape[0]}.")

from sklearn.model_selection import train_test_split

X_train_model, X_test, y_train_model, y_test = train_test_split(X, y, random_state=0, stratify=y)
X_train_model, X_train_threshold, y_train_model, y_train_threshold = \
train_test_split(X_train_model, y_train_model, random_state=0, stratify=y_train_model)

from sklearn.model_selection import StratifiedKFold

params = {
    "max_depth": [2, 4, 10, None],
    "n_estimators": [10, 50, 100],
}

skfold = StratifiedKFold(n_splits=3,
                         shuffle=True,
                         random_state=0)

from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier

def run_experiment(estimator, scoring, X, y, params, cv):
    gscv = (
        GridSearchCV(estimator=estimator,
                     param_grid=params,
                     scoring=scoring,
                     cv=cv)
        .fit(X, y)
    )

    return (
        pd.DataFrame(gscv.cv_results_)
        .pipe(lambda df:
              df[list(map(lambda x: "param_" + x,  params.keys())) + ["mean_test_score", "std_test_score"]])
    )