"""
=================================================
Calculating p-value with multiview data (CoMIGHT)
=================================================
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

from treeple.datasets import make_trunk_classification
from treeple.ensemble import HonestForestClassifier
from treeple.stats import PermutationHonestForestClassifier, build_coleman_forest
from treeple.tree import MultiViewDecisionTreeClassifier

sns.set(color_codes=True, style="white", context="talk", font_scale=1.5)
PALETTE = sns.color_palette("Set1")
sns.set_palette(PALETTE[1:5] + PALETTE[6:], n_colors=9)
sns.set_style("white", {"axes.edgecolor": "#dddddd"})
# %%
# Independence Testing
# --------------------
#
# Given samples from ``X``, ``Z``, and ``Y``, the independent hypothesis
# and its alternative are stated as:
#
# .. math:: H_0 : F_{X,Y|Z} = F_{X|Z} F_{Y|Z}
#
# .. math:: H_A : F_{X,Y|Z} \neq F_{X|Z} F_{Y|Z}
#
# By computing the p-value using ``treeple``, we can test if :math:`H_0`
# would be rejected, which confirms that ``X|Z`` and ``Y|Z`` are not independent.
# The p-value is generated by comparing the observed
# statistic difference with permuted differences, using conditional mutual
# information as a test statistic in this example.

# %%
# CMI
# ---
#
# Conditional mutual information (*CMI*) measures the dependence of *Y* on
# *X* conditioned on *Z*. It can be calculated by the difference between
# the joint *MI* (``I([X, Z]; Y)``) and the *MI* of Y on Z (``I(Y; Z)``):
#
# .. math:: I(Y; X \mid Z) = I(Y; [X, Z]) - I(Y; Z)
#
# Under the null hypothesis :math:`H_0`, the joint *MI* ``I(Y; [X, Z])``
# is equal to the *MI* of Y on Z ``I(Y; Z)``, so the *CMI* becomes zero. Thus, if
# the *CMI* is significantly larger than zero, we can reject the null hypothesis
# :math:`H_0`.
#
# With a multiview binary class simulation as an example, this tutorial
# will show how to use ``treeple`` to calculate the statistic and test the
# hypothesis with multiview data.

# %%
# Create a simulation with two gaussians
# --------------------------------------


# create a binary class simulation with two gaussians
# 500 samples for each class, class zero is standard
# gaussian, and class one has a mean at one for Z
Z, y = make_trunk_classification(
    n_samples=1000,
    n_dim=1,
    mu_0=0,
    mu_1=1,
    n_informative=1,
    seed=1,
)

# class one has a mean at two for X
X, y = make_trunk_classification(
    n_samples=1000,
    n_dim=1,
    mu_0=0,
    mu_1=2,
    n_informative=1,
    seed=2,
)

Z_X = np.hstack((Z, X))


Z_X_y = np.hstack((Z_X, y.reshape(-1, 1)))
Z_X_y = pd.DataFrame(Z_X_y, columns=["Z", "X", "y"])
Z_X_y = Z_X_y.replace({"y": 0.0}, "Class Zero")
Z_X_y = Z_X_y.replace({"y": 1.0}, "Class One")

fig, ax = plt.subplots(figsize=(6, 6))
fig.tight_layout()
ax.tick_params(labelsize=15)
sns.scatterplot(data=Z_X_y, x="Z", y="X", hue="y", palette=PALETTE[:2][::-1], alpha=0.2)
sns.kdeplot(data=Z_X_y, x="Z", y="X", hue="y", palette=PALETTE[:2][::-1], alpha=0.6)
ax.set_ylabel("Variable Two", fontsize=15)
ax.set_xlabel("Variable One", fontsize=15)
plt.legend(frameon=False, fontsize=15)

# %%
# Modify the simulation with permuted variable two
# ------------------------------------------------


# shuffle the labels
X_null = np.copy(X)
np.random.shuffle(X_null)

Z_X_y.insert(2, "Permuted X", X_null)

fig, ax = plt.subplots(figsize=(6, 6))
fig.tight_layout()
ax.tick_params(labelsize=15)
sns.scatterplot(data=Z_X_y, x="Z", y="Permuted X", hue="y", palette=PALETTE[:2][::-1], alpha=0.2)
sns.kdeplot(data=Z_X_y, x="Z", y="Permuted X", hue="y", palette=PALETTE[:2][::-1], alpha=0.6)
ax.set_ylabel("Permuted Variable Two", fontsize=15)
ax.set_xlabel("Variable One", fontsize=15)
plt.legend(frameon=False, fontsize=15)


# %%
# Fit the models and calculate the p-value
# ----------------------------------------
#
# Construct two forests, one trained with original data,
# and the other trained with permuted data. The test randomly
# permutes the two forests for a specified number of times.
#
# Each pair of forest outputs a set of mutual information statistics,
# and the statistic differences are used to calculate the p-vale.


# initialize the forest with 100 trees
est = HonestForestClassifier(
    n_estimators=100,
    max_samples=1.6,
    max_features=0.3,
    bootstrap=True,
    stratify=True,
    tree_estimator=MultiViewDecisionTreeClassifier(),
    random_state=1,
)

# initialize another forest with 100 trees
est_null = PermutationHonestForestClassifier(
    n_estimators=100,
    max_samples=1.6,
    max_features=0.3,
    bootstrap=True,
    stratify=True,
    tree_estimator=MultiViewDecisionTreeClassifier(),
    random_state=1,
)

PERMUTE = 10000

# conduct the hypothesis test with conditional mutual information
observed_diff, _, _, pvalue, mix_diff = build_coleman_forest(
    est,
    est_null,
    Z_X,
    y,
    metric="mi",
    n_repeats=PERMUTE,
    covariate_index=[1],
    return_posteriors=False,
    seed=1,
)

fig, ax = plt.subplots(figsize=(6, 6))
fig.tight_layout()
ax.tick_params(labelsize=15)

# histogram plot the statistic differences
ax.hist(mix_diff, bins=50, alpha=0.6, color=PALETTE[1], label="null")
ax.axvline(x=observed_diff, color=PALETTE[0], linestyle="--", label="observed")
ax.set_xlabel("Conditional Mutual Information Diff", fontsize=15)
ax.set_ylabel("# of Samples", fontsize=15)
plt.legend(frameon=False, fontsize=15)
plt.show()

print("p-value is:", round(pvalue, 2))
if pvalue < 0.05:
    print("The null hypothesis is rejected.")
else:
    print("The null hypothesis is not rejected.")
# sphinx_gallery_thumbnail_number = -1
