import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
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 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

# -----------------------------
# STANDARDIZE + PCA
# -----------------------------
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

pca = PCA(n_components=30)
X_pca = pca.fit_transform(X_scaled)

# -----------------------------
# TRAIN LOGISTIC REGRESSION
# -----------------------------
clf = LogisticRegression(max_iter=5000)
clf.fit(X_pca, y)

# Stability axis = weight vector
stability_axis = clf.coef_[0]

# -----------------------------
# COMPUTE STABILITY SCORE (SSS)
# -----------------------------
sss = X_pca @ stability_axis + clf.intercept_[0]

df["SSS"] = sss

# Separate groups
sss_llp = df[df["label"] == 1]["SSS"].values
sss_slp = df[df["label"] == 0]["SSS"].values

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

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

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

# Cohen's d
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("Cohen's d:", d_value)

# -----------------------------
# PLOT DISTRIBUTIONS
# -----------------------------
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("Stability Score (SSS)", fontsize=12, fontweight='bold')
plt.ylabel("Density", fontsize=12, fontweight='bold')
plt.title("Stability Axis Separation", fontsize=14, fontweight='bold')
plt.legend(prop={'weight':'bold'})

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