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 scipy.stats import mannwhitneyu
from matplotlib import rcParams

# -----------------------------
# FIGURE SETTINGS
# -----------------------------
rcParams['font.weight'] = 'bold'
rcParams['axes.labelweight'] = 'bold'
rcParams['axes.titleweight'] = 'bold'
rcParams['figure.dpi'] = 300

# -----------------------------
# LOAD LENGTH-MATCHED DATA
# -----------------------------
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

# -----------------------------
# CROSS-VALIDATED SSS
# -----------------------------
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

sss_out = np.zeros(len(y))

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

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

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

    # PCA (fit only on train)
    pca = PCA(n_components=30)
    X_train = pca.fit_transform(X_train)
    X_test = pca.transform(X_test)

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

    # Stability axis
    stability_axis = clf.coef_[0]

    # Compute SSS for test fold only
    sss_fold = X_test @ stability_axis + clf.intercept_[0]

    sss_out[test_idx] = sss_fold

# -----------------------------
# SEPARATE GROUPS
# -----------------------------
sss_llp = sss_out[y == 1]
sss_slp = sss_out[y == 0]

# -----------------------------
# STATISTICS
# -----------------------------
mean_llp = sss_llp.mean()
mean_slp = sss_slp.mean()

print("Cross-Validated Mean SSS LLP:", mean_llp)
print("Cross-Validated Mean SSS SLP:", mean_slp)

u_stat, p_value = mannwhitneyu(sss_llp, sss_slp)
print("Mann–Whitney p-value:", p_value)

def cohens_d(a, b):
    pooled_std = np.sqrt((a.std()**2 + b.std()**2) / 2)
    return (a.mean() - b.mean()) / pooled_std

d_value = cohens_d(sss_llp, sss_slp)
print("Cross-Validated Cohen's d:", d_value)

# -----------------------------
# PLOT
# -----------------------------
plt.figure(figsize=(6,6))

sns.kdeplot(sss_llp, fill=True, label="LLP", linewidth=2)
sns.kdeplot(sss_slp, fill=True, label="SLP", linewidth=2)

plt.xlabel("Cross-Validated Stability Score (SSS)", fontsize=12, fontweight='bold')
plt.ylabel("Density", fontsize=12, fontweight='bold')
plt.title("Cross-Validated Stability Axis Separation", fontsize=14, fontweight='bold')
plt.legend(prop={'weight':'bold'})

plt.tight_layout()
plt.savefig("Cross_Validated_Stability_Axis.png", dpi=300)
plt.show()