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

# --------- Load three datasets ----------
df_t1 = pd.read_excel("tester#1.xlsx")
df_t2 = pd.read_excel("tester#2.xlsx")
df_gold = pd.read_excel("Golden_Data.xlsx")

# --------- Get parameter list ----------
df_new = df_gold.head(3).dropna(axis=1, how="all")
df_param = df_new.drop(["T_TIME", "SITE_NUM"], axis=1)
parameters = list(df_param.columns)

# --------- Build master dictionary ----------
data = {
    param: {
        "Tester1": pd.to_numeric(df_t1[param].iloc[3:], errors="coerce").dropna().tolist(),
        "Tester2": pd.to_numeric(df_t2[param].iloc[3:], errors="coerce").dropna().tolist(),
        "Golden":  pd.to_numeric(df_gold[param].iloc[3:], errors="coerce").dropna().tolist(),
        "LimitL": df_new[param].iloc[1],
        "LimitU": df_new[param].iloc[2],
        "Unit": df_new[param].iloc[0]
    }
    for param in parameters
}

# --------- Plot comparison SPC ----------
os.makedirs("./charts_compare", exist_ok=True)

for param in parameters:

    t1 = np.array(data[param]["Tester1"])
    t2 = np.array(data[param]["Tester2"])
    gold = np.array(data[param]["Golden"])

    limit_l = data[param]["LimitL"]
    limit_u = data[param]["LimitU"]
    unit = data[param]["Unit"]

    # Golden reference control band
    mean = gold.mean()
    t1_mean = t1.mean()
    t2_mean = t2.mean()
    std = gold.std(ddof=1)
    UCL = mean + 3 * std
    LCL = mean - 3 * std

    # ---- Fix max length ----
    max_len = max(len(t1), len(t2), len(gold))
    xaxis = np.arange(1, max_len + 1)

    # ---- Correlation ----
    min_len_t1 = min(len(t1), len(gold))
    min_len_t2 = min(len(t2), len(gold))

    corr_t1 = np.corrcoef(t1[:min_len_t1], gold[:min_len_t1])[0,1]
    corr_t2 = np.corrcoef(t2[:min_len_t2], gold[:min_len_t2])[0,1]
    corr_t1_t2 = np.corrcoef(t1[:min_len_t1], t2[:min_len_t2])[0,1]

    # ---- CPK (based on Golden dataset) ----
    cpk = min((mean - limit_l) / (3 * std), (limit_u - mean) / (3 * std))

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

    # --- Plot all datasets ---
    plt.plot(range(1, len(t1)+1), t1, marker='o', label="Tester #1")
    plt.plot(range(1, len(t2)+1), t2, marker='o', label="Tester #2")
    plt.plot(range(1, len(gold)+1), gold, marker='o', label="Gold Reference", linewidth=3)

    # --- SPC lines ---
    plt.axhline(mean, linestyle='--', color='black', label=f"Mean (Gold)")
    plt.axhline(UCL, linestyle='-.', color='red', label="UCL (Mean + 3σ)")
    plt.axhline(LCL, linestyle='-.', color='red', label="LCL (Mean – 3σ)")

    # --- Spec Limits ---
    plt.axhline(limit_l, linestyle=':', color='orange', label="Lower Spec")
    plt.axhline(limit_u, linestyle=':', color='orange', label="Upper Spec")

    # ---- Fix X-axis ----
    # plt.xticks(np.arange(1, max_len + 1, step=1))
    # plt.xlim(1, max_len)

    # ---- Fix X-axis with whole numbers + padding ----
    max_len = max(len(t1), len(t2), len(gold))

    plt.xticks(np.arange(1, max_len + 1, 1))  # whole numbers
    plt.xlim(0, max_len + 1)  # space before 1 and after last point


    # ---- Add CPK + Correlation text box ----
    textstr = (
        f"Gold Mean = {mean:.4f}\n"
        f"Std Dev = {std:.4f}\n"
        f"CPK = {cpk:.4f}\n"
        f"Corr T1–Gold = {corr_t1:.4f}\n"
        f"Corr T2–Gold = {corr_t2:.4f}\n"
        f"Corr T1-T2 = {corr_t1_t2:.4f}\n"
        f"Tester#1 Mean = {t1_mean:.4f}\n"
        f"Tester#2 Mean = {t2_mean:4f}\n"
    )

    plt.gca().text(
        0.02, 0.98, textstr,
        transform=plt.gca().transAxes,
        fontsize=10,
        verticalalignment='top',
        bbox=dict(boxstyle="round,pad=0.4", facecolor="white", alpha=0.8)
    )

    # Labels
    plt.title(f"SPC Comparison - {param} ({unit})")
    plt.xlabel("Sample Index")
    plt.ylabel(f"Value ({unit})")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()

    # Save chart
    plt.savefig(f"./charts_compare/SPC_compare_{param}.png", dpi=300)
    plt.close()