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

# ===============================
# FILE PATHS
# ===============================

MOTOR_FILE = r"C:\Users\DURGA\motor_health.csv"
SURVIVAL_FILE = r"C:\Users\DURGA\survival.csv"

# ===============================
# LOAD DATA
# ===============================

motor_df = pd.read_csv(MOTOR_FILE)
survival_df = pd.read_csv(SURVIVAL_FILE)

# ===============================
# CONVERT ALIVE → SURVIVAL FRACTION
# ===============================

survival_df["initial"] = survival_df.groupby("group")["alive"].transform("max")
survival_df["survival_fraction"] = survival_df["alive"] / survival_df["initial"]

# ===============================
# CALCULATE SURVIVAL AUC
# ===============================

auc_list = []

for grp in survival_df["group"].unique():
    df_grp = survival_df[survival_df["group"] == grp]
    x = df_grp["day"].values
    y = df_grp["survival_fraction"].values
    
    auc = np.trapezoid(y, x)
    
    auc_list.append({"group": grp, "survival_AUC": auc})

auc_df = pd.DataFrame(auc_list)

# ===============================
# NORMALIZE MOTOR HEALTH (Day 7)
# ===============================

control_day7 = motor_df[motor_df["group"] == "CS"]["day7_slope"].values[0]
motor_df["norm_day7_motor"] = motor_df["day7_slope"] / control_day7

# ===============================
# MERGE DATA
# ===============================

final_df = pd.merge(motor_df, auc_df, on="group")

print("\nIntegrated Dataset:")
print(final_df)

# ===============================
# CORRELATION
# ===============================

corr = final_df["norm_day7_motor"].corr(final_df["survival_AUC"])
print("\nCorrelation:", corr)

# ===============================
# PLOT
# ===============================

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

plt.scatter(final_df["norm_day7_motor"], final_df["survival_AUC"])

for i, row in final_df.iterrows():
    plt.text(row["norm_day7_motor"], row["survival_AUC"], row["group"])

plt.xlabel("Normalized Day 7 Motor Health")
plt.ylabel("Survival AUC")
plt.title("Early Motor Health vs Survival")
plt.tight_layout()
plt.show()