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

from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, roc_curve
from scipy import stats

# =========================================================
# GLOBAL FIGURE SETTINGS
# =========================================================
plt.style.use("seaborn-v0_8-white")

plt.rcParams["font.weight"] = "bold"
plt.rcParams["axes.labelweight"] = "bold"
plt.rcParams["axes.titleweight"] = "bold"
plt.rcParams["figure.dpi"] = 300
plt.rcParams["savefig.dpi"] = 300

sns.set_context("talk")

# =========================================================
# LOAD DATA FUNCTION
# =========================================================
def load_dataset(llp_csv, slp_csv):

    llp = pd.read_csv(llp_csv)
    slp = pd.read_csv(slp_csv)

    llp["label"] = 1
    slp["label"] = 0

    df = pd.concat([llp, slp], axis=0).reset_index(drop=True)

    X = df.drop(columns=["Uniprot_ID","label"]).values
    y = df["label"].values

    return X, y


# =========================================================
# 5-FOLD CROSS VALIDATION
# =========================================================
def run_cv_auc(X, y, n_splits=5):

    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)

    aucs = []
    tprs = []
    mean_fpr = np.linspace(0,1,100)

    for train_idx, test_idx in skf.split(X,y):

        X_train, X_test = X[train_idx], X[test_idx]
        y_train, y_test = y[train_idx], y[test_idx]

        # Standardization
        scaler = StandardScaler()
        X_train = scaler.fit_transform(X_train)
        X_test = scaler.transform(X_test)

        # PCA
        pca = PCA(n_components=30)
        X_train = pca.fit_transform(X_train)
        X_test = pca.transform(X_test)

        # Logistic regression
        clf = LogisticRegression(max_iter=2000)
        clf.fit(X_train, y_train)

        probs = clf.predict_proba(X_test)[:,1]

        auc = roc_auc_score(y_test, probs)
        aucs.append(auc)

        fpr, tpr, _ = roc_curve(y_test, probs)

        interp_tpr = np.interp(mean_fpr, fpr, tpr)
        interp_tpr[0] = 0
        tprs.append(interp_tpr)

    return np.array(aucs), np.array(tprs), mean_fpr


# =========================================================
# LOAD RAW DATA
# =========================================================
X_raw, y_raw = load_dataset(
    "LLP_650M_embeddings.csv",
    "SLP_650M_embeddings.csv"
)

# =========================================================
# LOAD LENGTH-MATCHED DATA
# =========================================================
X_matched, y_matched = load_dataset(
    "LLP_cleaned_650M_embeddings.csv",
    "SLP_cleaned_650M_embeddings.csv"
)

# =========================================================
# RUN CROSS VALIDATION
# =========================================================
raw_aucs, raw_tprs, mean_fpr = run_cv_auc(X_raw, y_raw)
matched_aucs, matched_tprs, _ = run_cv_auc(X_matched, y_matched)

print("RAW AUCs:", raw_aucs)
print("MATCHED AUCs:", matched_aucs)

print("\nMean RAW AUC:", raw_aucs.mean())
print("Mean MATCHED AUC:", matched_aucs.mean())

print("RAW SD:", raw_aucs.std())
print("MATCHED SD:", matched_aucs.std())

# =========================================================
# STATISTICAL TEST
# =========================================================
t_stat, p_value = stats.ttest_rel(raw_aucs, matched_aucs)

print("\nPaired t-test:")
print("t-statistic:", t_stat)
print("p-value:", p_value)

# =========================================================
# EFFECT SIZE
# =========================================================
diff = raw_aucs - matched_aucs
cohens_d = diff.mean() / diff.std()

print("\nCohen's d effect size:", cohens_d)

# =========================================================
# ROC CURVE PLOT
# =========================================================
plt.figure(figsize=(6,6))

mean_tpr_raw = raw_tprs.mean(axis=0)
mean_tpr_matched = matched_tprs.mean(axis=0)

plt.plot(
    mean_fpr,
    mean_tpr_raw,
    linewidth=2,
    label=f"Raw (AUC = {raw_aucs.mean():.3f})"
)

plt.plot(
    mean_fpr,
    mean_tpr_matched,
    linewidth=2,
    label=f"Matched (AUC = {matched_aucs.mean():.3f})"
)

plt.plot([0,1],[0,1], linestyle="--", color="gray", linewidth=1)

plt.xlabel("False Positive Rate", fontweight="bold")
plt.ylabel("True Positive Rate", fontweight="bold")

legend = plt.legend(frameon=False)

for text in legend.get_texts():
    text.set_fontweight("bold")

plt.tight_layout()
plt.savefig("ROC_Curve_Comparison.png", bbox_inches="tight")
plt.show()


# =========================================================
# AUC DISTRIBUTION PLOT
# =========================================================
plt.figure(figsize=(6,6))

data = pd.DataFrame({
    "AUC": np.concatenate([raw_aucs, matched_aucs]),
    "Dataset": ["Raw"]*len(raw_aucs) + ["Matched"]*len(matched_aucs)
})

sns.boxplot(
    x="Dataset",
    y="AUC",
    data=data,
    showfliers=False
)

sns.stripplot(
    x="Dataset",
    y="AUC",
    data=data,
    color="black",
    size=7
)

plt.ylabel("ROC-AUC", fontweight="bold")
plt.xlabel("")

for label in plt.gca().get_xticklabels():
    label.set_fontweight("bold")

for label in plt.gca().get_yticklabels():
    label.set_fontweight("bold")

plt.tight_layout()
plt.savefig("AUC_Distribution_Comparison.png", bbox_inches="tight")
plt.show()