#!/usr/bin/env python3
# =============================================================================
# FrESH - Dynamic residual analysis during ramp experiments
# Final paper-facing version
# =============================================================================

from __future__ import annotations

from pathlib import Path
import json
import re

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

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
OUTPUT_DIR = BASE_DIR / "output"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

RAMP_WINDOWS_FILE = DATA_DIR / "ramp_windows.json"
MODEL_SUMMARY_FILE = OUTPUT_DIR / "table_transfer.csv"
MODEL_ROWS_FILE = OUTPUT_DIR / "all_characterization_models.csv"

CHILLERS = ["RE1050", "RP1845"]
PROBES = ["probe1", "probe2", "probe3", "probe4", "probe5"]

PAIR_RE = re.compile(
    r"(?P<chiller>RE1050|RP1845)_RAMP_(?P<run>.+)_(?P<dtype>pt100|chiller)\.csv$",
    re.IGNORECASE,
)

CHILLER_COLORS = {
    "RE1050": "#2E6E4D",
    "RP1845": "#D2874C",
}

SAVE_TIMELINES = False
SAVE_SUPPORTING_CSV = True

TARGET_STEPS = np.array([0, -5, -10, -15, -20, -25, -30, -35], dtype=float)
HALF_WIDTH = 2.5


# ---------------------------------------------------------------------------
# Metadata
# ---------------------------------------------------------------------------
def load_ramp_windows() -> dict:
    if not RAMP_WINDOWS_FILE.exists():
        raise FileNotFoundError(f"Missing ramp windows file: {RAMP_WINDOWS_FILE}")
    with open(RAMP_WINDOWS_FILE, "r", encoding="utf-8") as f:
        return json.load(f)


def load_static_model_summary() -> pd.DataFrame:
    if MODEL_SUMMARY_FILE.exists():
        df = pd.read_csv(MODEL_SUMMARY_FILE)
        required = {"Chiller", "a_mean", "b_mean"}
        if not required.issubset(df.columns):
            raise ValueError(
                f"{MODEL_SUMMARY_FILE} exists but does not contain the required columns: {required}"
            )
        out = (
            df[["Chiller", "a_mean", "b_mean"]]
            .rename(columns={"Chiller": "chiller", "a_mean": "a", "b_mean": "b"})
            .copy()
        )
        out["chiller"] = out["chiller"].astype(str).str.strip()
        return out.sort_values("chiller").reset_index(drop=True)

    if MODEL_ROWS_FILE.exists():
        df = pd.read_csv(MODEL_ROWS_FILE)
        required = {"chiller", "a_slope", "b_intercept"}
        if not required.issubset(df.columns):
            raise ValueError(
                f"{MODEL_ROWS_FILE} exists but does not contain the required columns: {required}"
            )
        out = (
            df.groupby("chiller")[["a_slope", "b_intercept"]]
            .mean()
            .reset_index()
            .rename(columns={"a_slope": "a", "b_intercept": "b"})
        )
        out["chiller"] = out["chiller"].astype(str).str.strip()
        return out.sort_values("chiller").reset_index(drop=True)

    raise FileNotFoundError(
        "No static model summary found. Run analyze_step.py first so that "
        "'output/table_transfer.csv' or 'output/all_characterization_models.csv' exists."
    )


# ---------------------------------------------------------------------------
# Time parsing
# ---------------------------------------------------------------------------
def parse_datetime_flexible(series: pd.Series) -> pd.Series:
    s = series.astype(str).str.strip()

    dt_time = pd.to_datetime(
        "2000-01-01 " + s,
        format="%Y-%m-%d %H:%M:%S",
        errors="coerce",
    )
    if dt_time.notna().mean() > 0.8:
        return dt_time

    dt_full = pd.to_datetime(s, errors="coerce")
    if dt_full.notna().mean() > 0.8:
        return dt_full

    td = pd.to_timedelta(s, errors="coerce")
    if td.notna().mean() > 0.8:
        return pd.Timestamp("2000-01-01") + td

    raise ValueError("Could not interpret time column as HH:MM:SS, datetime, or timedelta.")


def hhmmss_to_timestamp(time_str: str, reference_date: pd.Timestamp) -> pd.Timestamp:
    h, m, s = map(int, str(time_str).split(":"))
    return reference_date.normalize() + pd.Timedelta(hours=h, minutes=m, seconds=s)


# ---------------------------------------------------------------------------
# Pair discovery
# ---------------------------------------------------------------------------
def find_ramp_pairs(data_dir: Path) -> dict:
    grouped = {}
    for path in sorted(data_dir.rglob("*.csv")):
        m = PAIR_RE.search(path.name)
        if not m:
            continue
        key = (m.group("chiller").upper(), m.group("run"))
        grouped.setdefault(key, {})
        grouped[key][m.group("dtype").lower()] = path
    return grouped


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------
def load_chiller_data(filepath: Path) -> pd.DataFrame:
    df = pd.read_csv(filepath, index_col=0)
    df.columns = df.columns.astype(str).str.strip()

    idx = parse_datetime_flexible(pd.Series(df.index.astype(str)))
    df.index = pd.DatetimeIndex(idx)
    df.index.name = "datetime"

    df = df[df.index.notna()].copy()

    for col in ["SP", "BT"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")

    df = df.sort_index()
    df = df[~df.index.duplicated(keep="first")]

    return df


def load_pt100_data(filepath: Path) -> pd.DataFrame:
    df = pd.read_csv(
        filepath,
        names=["datetime", "probe1", "probe2", "probe3", "probe4", "probe5", "avg_probes"],
        header=None,
        on_bad_lines="skip",
    )

    df["datetime"] = parse_datetime_flexible(df["datetime"])

    for col in PROBES:
        df[col] = pd.to_numeric(df[col], errors="coerce")

    df = (
        df.dropna(subset=["datetime"])
        .sort_values("datetime")
        .drop_duplicates(subset=["datetime"], keep="first")
        .set_index("datetime")
    )

    return df


# ---------------------------------------------------------------------------
# Windowing
# ---------------------------------------------------------------------------
def extract_ramp_window(df: pd.DataFrame, ramp_windows: dict, chiller: str, run_key: str):
    if chiller not in ramp_windows or run_key not in ramp_windows[chiller]:
        raise KeyError(f"Missing ramp window for {chiller} {run_key} in {RAMP_WINDOWS_FILE}")

    start_str, end_str = ramp_windows[chiller][run_key]
    ref_date = pd.Timestamp(df["datetime"].dropna().iloc[0])

    start_dt = hhmmss_to_timestamp(start_str, ref_date)
    end_dt = hhmmss_to_timestamp(end_str, ref_date)

    df_ramp = df[(df["datetime"] >= start_dt) & (df["datetime"] <= end_dt)].copy()
    return df_ramp, start_dt, end_dt


# ---------------------------------------------------------------------------
# Plotting
# ---------------------------------------------------------------------------
def plot_ramp_timeline(
    df: pd.DataFrame,
    chiller: str,
    run_key: str,
    start_dt: pd.Timestamp,
    end_dt: pd.Timestamp,
) -> None:
    fig, ax = plt.subplots(figsize=(14, 6))

    if "SP" in df.columns:
        ax.plot(df["datetime"], df["SP"], color="#7f7f7f", linestyle="--", linewidth=1.6, label="SP")
    if "BT" in df.columns:
        ax.plot(df["datetime"], df["BT"], color="k", linewidth=2.2, label="BT")

    for probe in PROBES:
        if probe in df.columns:
            ax.plot(df["datetime"], df[probe], ".", markersize=2, alpha=0.45, label=probe)

    ax.axvspan(start_dt, end_dt, color=CHILLER_COLORS.get(chiller, "#cccccc"), alpha=0.16, label="Ramp window")
    ax.axvline(start_dt, color=CHILLER_COLORS.get(chiller, "#666666"), linewidth=1.2, alpha=0.85)
    ax.axvline(end_dt, color=CHILLER_COLORS.get(chiller, "#666666"), linewidth=1.2, alpha=0.85)

    ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M"))
    ax.xaxis.set_major_locator(mdates.MinuteLocator(interval=5))
    ax.set_xlabel("Time")
    ax.set_ylabel("Temperature (°C)")
    ax.set_title(f"{chiller} {run_key}")
    ax.grid(True, alpha=0.25)

    handles, labels = ax.get_legend_handles_labels()
    unique = dict(zip(labels, handles))
    ax.legend(unique.values(), unique.keys(), loc="best", fontsize=8, ncol=2)

    plt.xticks(rotation=45)
    plt.tight_layout()

    out = OUTPUT_DIR / f"timeline_ramp_{chiller}_{run_key}.png"
    plt.savefig(out, dpi=150, bbox_inches="tight")

    out_pdf = OUTPUT_DIR / "timeline_ramp_{chiller}_{run_key}.pdf"
    plt.savefig(out_pdf, bbox_inches="tight")

    plt.close()


# ---------------------------------------------------------------------------
# Core computation
# ---------------------------------------------------------------------------
def compute_dynamic_residuals(model_df: pd.DataFrame, ramp_windows: dict) -> tuple[pd.DataFrame, pd.DataFrame]:
    models = model_df.set_index("chiller")[["a", "b"]].copy()
    rows = []
    run_rows = []

    pairs = find_ramp_pairs(DATA_DIR)

    for (chiller, run_key), pair in sorted(pairs.items()):
        if chiller not in models.index:
            continue
        if "chiller" not in pair or "pt100" not in pair:
            continue

        a = float(models.loc[chiller, "a"])
        b = float(models.loc[chiller, "b"])

        try:
            df_ch = load_chiller_data(pair["chiller"])
            df_pt = load_pt100_data(pair["pt100"])

            ch_flat = df_ch[["SP", "BT"]].reset_index().sort_values("datetime").reset_index(drop=True)
            pt_flat = df_pt.reset_index().sort_values("datetime").reset_index(drop=True)

            df = pd.merge_asof(
                pt_flat,
                ch_flat[["datetime", "SP", "BT"]],
                on="datetime",
                direction="nearest",
                tolerance=pd.Timedelta(seconds=2),
            )

            df["SP"] = pd.to_numeric(df["SP"], errors="coerce")
            df["BT"] = pd.to_numeric(df["BT"], errors="coerce")
            df = df.dropna(subset=["SP", "BT"]).reset_index(drop=True)

            if len(df) < 50:
                continue

            df_ramp, start_dt, end_dt = extract_ramp_window(df, ramp_windows, chiller, run_key)

            if len(df_ramp) < 50:
                continue

            if SAVE_TIMELINES:
                plot_ramp_timeline(df, chiller, run_key, start_dt, end_dt)

            df_ramp = df_ramp.copy()
            df_ramp["T_corr"] = a * df_ramp["BT"] + b

            valid_points = 0
            valid_probes = 0

            for probe in [p for p in PROBES if p in df_ramp.columns]:
                sub = df_ramp[["datetime", "BT", "T_corr", probe]].copy()
                sub[probe] = pd.to_numeric(sub[probe], errors="coerce")
                sub = sub.dropna()
                if sub.empty:
                    continue

                sub = sub.rename(columns={probe: "T_probe"})
                sub["delta_dyn"] = sub["T_probe"] - sub["T_corr"]
                sub["chiller"] = chiller
                sub["run_key"] = run_key
                sub["probe"] = probe

                valid_points += len(sub)
                valid_probes += 1

                rows.append(
                    sub[["chiller", "run_key", "probe", "datetime", "BT", "T_probe", "T_corr", "delta_dyn"]]
                )

            run_rows.append(
                {
                    "chiller": chiller,
                    "run_key": run_key,
                    "a_used": a,
                    "b_used": b,
                    "matched_rows": int(len(df)),
                    "ramp_rows": int(len(df_ramp)),
                    "n_valid_probes": int(valid_probes),
                    "n_valid_points": int(valid_points),
                    "BT_min": float(df_ramp["BT"].min()),
                    "BT_max": float(df_ramp["BT"].max()),
                    "start_dt": start_dt,
                    "end_dt": end_dt,
                }
            )

        except Exception as exc:
            run_rows.append(
                {
                    "chiller": chiller,
                    "run_key": run_key,
                    "a_used": a,
                    "b_used": b,
                    "matched_rows": 0,
                    "ramp_rows": 0,
                    "n_valid_probes": 0,
                    "n_valid_points": 0,
                    "BT_min": np.nan,
                    "BT_max": np.nan,
                    "start_dt": pd.NaT,
                    "end_dt": pd.NaT,
                    "error": str(exc),
                }
            )

    if rows:
        residuals_df = pd.concat(rows, ignore_index=True)
    else:
        residuals_df = pd.DataFrame(
            columns=["chiller", "run_key", "probe", "datetime", "BT", "T_probe", "T_corr", "delta_dyn"]
        )

    run_summary_df = pd.DataFrame(run_rows)

    residuals_df.to_csv(OUTPUT_DIR / "dynamic_residuals_ramps.csv", index=False)
    if SAVE_SUPPORTING_CSV:
        run_summary_df.to_csv(OUTPUT_DIR / "dynamic_run_summary.csv", index=False)

    return residuals_df, run_summary_df


# ---------------------------------------------------------------------------
# Figure 4 summary
# ---------------------------------------------------------------------------
def build_binned_summary(residuals_df: pd.DataFrame) -> pd.DataFrame:
    rows = []

    for chiller in CHILLERS:
        sub = residuals_df[residuals_df["chiller"] == chiller].copy()
        sub = sub[np.isfinite(sub["BT"]) & np.isfinite(sub["delta_dyn"])]
        if sub.empty:
            continue

        for step in TARGET_STEPS:
            mask = (sub["BT"] > step - HALF_WIDTH) & (sub["BT"] <= step + HALF_WIDTH)
            vals = sub.loc[mask, "delta_dyn"].dropna()

            if len(vals) == 0:
                continue

            rows.append(
                {
                    "chiller": chiller,
                    "BT_bin_center": float(step),
                    "delta_dyn_mean": float(vals.mean()),
                    "delta_dyn_std": float(vals.std()),
                    "n": int(len(vals)),
                }
            )

    return pd.DataFrame(rows)


def plot_dynamic_residuals_amt(residuals_df: pd.DataFrame) -> pd.DataFrame:
    if residuals_df.empty:
        print("No residuals to plot.")
        return pd.DataFrame()

    binned_df = build_binned_summary(residuals_df)

    fig, ax = plt.subplots(figsize=(5.2, 4.2))

    for chiller in CHILLERS:
        g = binned_df[binned_df["chiller"] == chiller].sort_values("BT_bin_center")
        if g.empty:
            continue

        ax.errorbar(
            g["BT_bin_center"],
            g["delta_dyn_mean"],
            yerr=g["delta_dyn_std"],
            fmt="o-",
            color=CHILLER_COLORS[chiller],
            markersize=4.5,
            linewidth=1.4,
            capsize=3,
            elinewidth=1.0,
            label=chiller,
        )

    ax.axhline(0, color="k", linestyle="--", linewidth=1.0, alpha=0.6)
    ax.set_xlabel(r"Bath temperature $T_\mathrm{BT}$ ($^\circ$C)", fontsize=10)
    ax.set_ylabel(
        r"Dynamic residual $T_\mathrm{probe} - T_\mathrm{well,corr}$ ($^\circ$C)",
        fontsize=10,
    )
    ax.set_xlim(-36, 1)
    ax.set_xticks([-35, -30, -25, -20, -15, -10, -5, 0])
    ax.grid(True, alpha=0.25, linewidth=0.6)
    ax.tick_params(axis="both", labelsize=9)
    ax.legend(fontsize=8, frameon=False, loc="upper right")

    plt.tight_layout()
    out = OUTPUT_DIR / "fig4_dynamic_residuals_amt.png"
    plt.savefig(out, dpi=300, bbox_inches="tight")

    out_pdf = OUTPUT_DIR / "fig4_dynamic_residuals_amt.pdf"
    plt.savefig(out_pdf, bbox_inches="tight")

    plt.close()

    return binned_df


# ---------------------------------------------------------------------------
# Summaries and reporting
# ---------------------------------------------------------------------------
def build_probe_summary(residuals_df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    if residuals_df.empty:
        empty_probe = pd.DataFrame(columns=["chiller", "probe", "mean_delta_dyn", "std_delta_dyn", "n"])
        empty_dyn = pd.DataFrame(columns=["chiller", "u_dynamic", "mean_delta_dyn", "overall_std_delta_dyn", "n"])
        return empty_probe, empty_dyn

    probe_summary = (
        residuals_df.groupby(["chiller", "probe"])["delta_dyn"]
        .agg(["mean", "std", "count"])
        .reset_index()
        .rename(columns={"mean": "mean_delta_dyn", "std": "std_delta_dyn", "count": "n"})
        .sort_values(["chiller", "probe"])
    )

    overall_summary = (
        residuals_df.groupby("chiller")["delta_dyn"]
        .agg(["mean", "std", "count"])
        .reset_index()
        .rename(columns={"mean": "mean_delta_dyn", "std": "overall_std_delta_dyn", "count": "n"})
    )

    dynamic_term = (
        probe_summary.groupby("chiller")["std_delta_dyn"]
        .mean()
        .reset_index()
        .rename(columns={"std_delta_dyn": "u_dynamic"})
    )

    dynamic_summary = overall_summary.merge(dynamic_term, on="chiller", how="left")
    dynamic_summary = dynamic_summary[["chiller", "u_dynamic", "mean_delta_dyn", "overall_std_delta_dyn", "n"]]
    dynamic_summary = dynamic_summary.sort_values("chiller").reset_index(drop=True)

    return probe_summary, dynamic_summary


def build_metrics_dict(
    residuals_df: pd.DataFrame,
    run_summary_df: pd.DataFrame,
    probe_summary_df: pd.DataFrame,
    dynamic_summary_df: pd.DataFrame,
    binned_df: pd.DataFrame,
) -> dict:
    metrics = {
        "n_residual_rows": int(len(residuals_df)),
        "n_runs_processed": int(run_summary_df["run_key"].nunique()) if not run_summary_df.empty else 0,
        "chillers": {},
    }

    for chiller in CHILLERS:
        dyn = dynamic_summary_df[dynamic_summary_df["chiller"] == chiller]
        probe = probe_summary_df[probe_summary_df["chiller"] == chiller]
        bins = binned_df[binned_df["chiller"] == chiller]
        runs = run_summary_df[run_summary_df["chiller"] == chiller]

        if dyn.empty:
            continue

        metrics["chillers"][chiller] = {
            "u_dynamic": float(dyn["u_dynamic"].iloc[0]),
            "mean_delta_dyn": float(dyn["mean_delta_dyn"].iloc[0]),
            "overall_std_delta_dyn": float(dyn["overall_std_delta_dyn"].iloc[0]),
            "n_points": int(dyn["n"].iloc[0]),
            "n_runs": int(runs["run_key"].nunique()) if not runs.empty else 0,
            "probe_std_values": {
                row["probe"]: float(row["std_delta_dyn"])
                for _, row in probe.iterrows()
            },
            "binned_means": {
                str(int(row["BT_bin_center"])): {
                    "mean": float(row["delta_dyn_mean"]),
                    "std": float(row["delta_dyn_std"]),
                    "n": int(row["n"]),
                }
                for _, row in bins.iterrows()
            },
        }

    return metrics


def write_report(
    run_summary_df: pd.DataFrame,
    probe_summary_df: pd.DataFrame,
    dynamic_summary_df: pd.DataFrame,
    binned_df: pd.DataFrame,
    outpath: Path,
) -> None:
    lines = []
    lines.append("FrESH dynamic-ramp analysis report")
    lines.append("=" * 80)
    lines.append("")

    if not run_summary_df.empty:
        valid_runs = run_summary_df[run_summary_df["n_valid_points"] > 0].copy()
        lines.append(f"Runs with valid residual output: {len(valid_runs)}")
        lines.append("")

    lines.append("Dynamic uncertainty summary")
    lines.append("-" * 80)
    for _, row in dynamic_summary_df.iterrows():
        lines.append(
            f"{row['chiller']}: "
            f"u_dynamic = {row['u_dynamic']:.3f} °C, "
            f"mean(delta_dyn) = {row['mean_delta_dyn']:.3f} °C, "
            f"overall std(delta_dyn) = {row['overall_std_delta_dyn']:.3f} °C, "
            f"n = {int(row['n'])}"
        )
    lines.append("")

    lines.append("Per-probe residual spread")
    lines.append("-" * 80)
    for _, row in probe_summary_df.iterrows():
        lines.append(
            f"{row['chiller']} {row['probe']}: "
            f"mean = {row['mean_delta_dyn']:.3f} °C, "
            f"std = {row['std_delta_dyn']:.3f} °C, "
            f"n = {int(row['n'])}"
        )
    lines.append("")

    lines.append("Figure-4 binned values")
    lines.append("-" * 80)
    for chiller in CHILLERS:
        lines.append(chiller)
        sub = binned_df[binned_df["chiller"] == chiller].sort_values("BT_bin_center")
        for _, row in sub.iterrows():
            lines.append(
                f"  BT {int(row['BT_bin_center']):>3} °C: "
                f"mean = {row['delta_dyn_mean']:.3f} °C, "
                f"std = {row['delta_dyn_std']:.3f} °C, "
                f"n = {int(row['n'])}"
            )
        lines.append("")

    outpath.write_text("\n".join(lines) + "\n", encoding="utf-8")


def print_paper_summary(dynamic_summary_df: pd.DataFrame) -> None:
    print("\n" + "=" * 80)
    print("RAMP ANALYSIS COMPLETE")
    print("=" * 80)

    print("\nPaper-facing outputs written to output/:")
    print("  - fig4_dynamic_residuals_amt.png")
    print("  - dynamic_residuals_ramps.csv")
    print("  - dynamic_probe_summary.csv")
    print("  - dynamic_uncertainty_summary.csv")
    print("  - dynamic_binned_summary.csv")
    print("  - dynamic_report.txt")
    print("  - dynamic_text_metrics.json")

    if SAVE_TIMELINES:
        print("  - timeline_ramp_<chiller>_<run>.png")

    print("\nTable-2 / Results values:")
    for _, row in dynamic_summary_df.iterrows():
        print(
            f"  {row['chiller']}: "
            f"mean delta_dyn = {row['mean_delta_dyn']:.3f} °C, "
            f"overall std = {row['overall_std_delta_dyn']:.3f} °C, "
            f"u_dynamic (mean per-probe std) = {row['u_dynamic']:.3f} °C, "
            f"n = {int(row['n'])}"
        )


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def main() -> None:
    ramp_windows = load_ramp_windows()
    model_df = load_static_model_summary()

    residuals_df, run_summary_df = compute_dynamic_residuals(model_df, ramp_windows)
    binned_df = plot_dynamic_residuals_amt(residuals_df)
    probe_summary_df, dynamic_summary_df = build_probe_summary(residuals_df)

    probe_summary_df.to_csv(OUTPUT_DIR / "dynamic_probe_summary.csv", index=False)
    dynamic_summary_df.to_csv(OUTPUT_DIR / "dynamic_uncertainty_summary.csv", index=False)
    binned_df.to_csv(OUTPUT_DIR / "dynamic_binned_summary.csv", index=False)

    metrics = build_metrics_dict(
        residuals_df=residuals_df,
        run_summary_df=run_summary_df,
        probe_summary_df=probe_summary_df,
        dynamic_summary_df=dynamic_summary_df,
        binned_df=binned_df,
    )

    with open(OUTPUT_DIR / "dynamic_text_metrics.json", "w", encoding="utf-8") as f:
        json.dump(metrics, f, indent=2, default=str)

    write_report(
        run_summary_df=run_summary_df,
        probe_summary_df=probe_summary_df,
        dynamic_summary_df=dynamic_summary_df,
        binned_df=binned_df,
        outpath=OUTPUT_DIR / "dynamic_report.txt",
    )

    print_paper_summary(dynamic_summary_df)


if __name__ == "__main__":
    main()