# =========================================================
# IMPORTS
# =========================================================
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 sklearn.metrics import roc_auc_score, accuracy_score, roc_curve

# =========================================================
# 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 TRAIN DATA (U2OS)
# =========================================================
llp_u2os = pd.read_csv("LLP (Cell 1).csv")
slp_u2os = pd.read_csv("SLP (Cell 2).csv")

llp_u2os["label"] = 1
slp_u2os["label"] = 0

train_df = pd.concat([llp_u2os, slp_u2os]).reset_index(drop=True)

X_train = train_df.drop(columns=["Uniprot_ID","label"]).values
y_train = train_df["label"].values

# =========================================================
# LOAD TEST DATA (HCT116)
# =========================================================
llp_test = pd.read_csv("LLL (Cell X).csv")
slp_test = pd.read_csv("SLP (Cell X).csv")

llp_test["label"] = 1
slp_test["label"] = 0

test_df = pd.concat([llp_test, slp_test]).reset_index(drop=True)

X_test = test_df.drop(columns=["Uniprot_ID","label"]).values
y_test = test_df["label"].values

# =========================================================
# TRAIN MODEL (U2OS)
# =========================================================
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

pca = PCA(n_components=30)
X_train_pca = pca.fit_transform(X_train_scaled)
X_test_pca = pca.transform(X_test_scaled)

clf = LogisticRegression(max_iter=5000)
clf.fit(X_train_pca, y_train)

# =========================================================
# TEST MODEL (HCT116)
# =========================================================
probs = clf.predict_proba(X_test_pca)[:,1]
pred_labels = clf.predict(X_test_pca)

auc = roc_auc_score(y_test, probs)
acc = accuracy_score(y_test, pred_labels)

print("External Validation (Train U2OS → Test HCT116)")
print("ROC-AUC:", auc)
print("Accuracy:", acc)

# =========================================================
# FIGURE 1 — ROC CURVE
# =========================================================
fpr, tpr, _ = roc_curve(y_test, probs)

plt.figure(figsize=(6,6))

plt.plot(
    fpr,
    tpr,
    linewidth=2,
    label=f"External Validation (AUC = {auc:.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("External_ROC.png", bbox_inches="tight")
plt.show()

# =========================================================
# FIGURE 2 — SCORE DISTRIBUTION
# =========================================================
test_df["score"] = probs

scores_llp = test_df[test_df["label"]==1]["score"]
scores_slp = test_df[test_df["label"]==0]["score"]

plt.figure(figsize=(6,6))

sns.kdeplot(
    scores_llp,
    fill=True,
    linewidth=2,
    label="LLP"
)

sns.kdeplot(
    scores_slp,
    fill=True,
    linewidth=2,
    label="SLP"
)

plt.xlabel("Predicted Stability Score", fontweight="bold")
plt.ylabel("Density", fontweight="bold")

legend = plt.legend(frameon=False)

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

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