"""LOCO and blocked-split validation for CCD (Round 5).

Reuses shared data loader from resampling.py.
Target-only POD basis. Q_delay=6 for force/action.

0.75L and 1.0L only. 1.5L excluded from validation.

Usage:
    conda run -n pycuda_3_10 python ccd/validate.py
"""
from __future__ import annotations

import json
import os
import sys
import time

import numpy as np

_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from CCD_analysis.configs import DATA_DIR
from CCD_analysis.utils.resampling import (
    compute_reduced_ccd, cumulative_energy,
    load_aligned_fields, make_force_obs,
    build_field_matrix, project_into_basis,
)

R_LIST = [6, 8, 10]
CCD_Q = 6
N_CYCLES = 4
N_PTS = 24
DIAMETERS = [0.75, 1.0]


def r2(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Coefficient of determination."""
    ss_r = np.sum((y_true - y_pred) ** 2)
    ss_t = np.sum((y_true - np.mean(y_true)) ** 2)
    return float(1.0 - ss_r / (ss_t + 1e-12))


def reconstruct_observable(W, sigma, R, a_test, y_train):
    """Reconstruct observable from CCD modes.

    Returns dict with 'mode1' and 'm80' reconstructions.
    """
    am = np.mean(a_test, axis=1, keepdims=True)
    as_ = np.std(a_test, axis=1, keepdims=True) + 1e-12
    a_test_z = (a_test - am) / as_
    z_test = W.T @ a_test_z

    ym = np.mean(y_train, axis=1, keepdims=True)
    ys = np.std(y_train, axis=1, keepdims=True) + 1e-12
    half = CCD_Q // 2
    m_obs = y_train.shape[0]

    en = cumulative_energy(sigma)
    m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 1

    results = {}

    # Mode-1
    if R.shape[1] >= 1:
        pz_1 = R[:, :1] * sigma[:1] @ z_test[:1, :]
        yp_1 = pz_1[half * m_obs:(half + 1) * m_obs, :] * ys + ym
        results["mode1"] = yp_1
    else:
        results["mode1"] = np.zeros_like(y_train[:, :a_test.shape[1]])

    # M80
    n_rm = min(m80, R.shape[1])
    if n_rm >= 1:
        pz_m = R[:, :n_rm] * sigma[:n_rm] @ z_test[:n_rm, :]
        yp_m = pz_m[half * m_obs:(half + 1) * m_obs, :] * ys + ym
        results["m80"] = yp_m
    else:
        results["m80"] = np.zeros_like(y_train[:, :a_test.shape[1]])

    return results


def run_single_diameter(diam: float, scene_data: dict) -> dict:
    """Run validation for one diameter. Returns results dict."""
    tgt_key = f"target_cylinder_{diam}L"
    ill_key = f"illusion_{diam}L"

    tgt_d = scene_data[tgt_key]
    ill_d = scene_data[ill_key]
    unc_d = scene_data["pinball"]

    # Pre-build target-only field matrices
    tgt_f = build_field_matrix(tgt_d["ux"], tgt_d["uy"])
    ill_f = build_field_matrix(ill_d["ux"], ill_d["uy"])
    unc_f = build_field_matrix(unc_d["ux"], unc_d["uy"])

    diam_results = {}
    pod_cache = {}

    # Pre-compute target-only POD for each fold
    for fold in range(N_CYCLES):
        test_cyc = fold
        train_cyc = [c for c in range(N_CYCLES) if c != test_cyc]
        train_idx = sorted([c * N_PTS + p for c in train_cyc for p in range(N_PTS)])
        for r in R_LIST:
            Q_ref = tgt_f[:, train_idx]
            mf = np.mean(Q_ref, axis=1)
            U, _, _ = np.linalg.svd(Q_ref - mf[:, None], full_matrices=False)
            pod_cache[("loco", fold, r)] = (mf, U[:, :r])

    # Blocked split
    train_idx_full = list(range(0, 2 * N_PTS))
    for r in R_LIST:
        Q_ref = tgt_f[:, train_idx_full]
        mf = np.mean(Q_ref, axis=1)
        U, _, _ = np.linalg.svd(Q_ref - mf[:, None], full_matrices=False)
        pod_cache[("blocked", 0, r)] = (mf, U[:, :r])

    # -- LOCO --
    print("\n--- LOCO (4-fold) ---", flush=True)
    loco_results = {}
    for r in R_LIST:
        for obs in ["force_fy", "force_fx", "action"]:
            fold_r2_m1, fold_r2_m80 = [], []
            for fold in range(N_CYCLES):
                test_cyc = fold
                train_cyc = [c for c in range(N_CYCLES) if c != test_cyc]
                train_idx = sorted([c * N_PTS + p for c in train_cyc for p in range(N_PTS)])
                test_idx = sorted([c * N_PTS + p for c in [test_cyc] for p in range(N_PTS)])

                mf, modes_r = pod_cache[("loco", fold, r)]

                for name, d, fld in [
                    (tgt_key, tgt_d, tgt_f), (ill_key, ill_d, ill_f), ("pinball", unc_d, unc_f)
                ]:
                    if obs == "action" and "illusion" not in name:
                        continue
                    if d.get("forces") is None and "force" in obs:
                        continue

                    a_train = modes_r.T @ (fld[:, train_idx] - mf[:, None])
                    a_test = modes_r.T @ (fld[:, test_idx] - mf[:, None])

                    if "force" in obs:
                        f_mode = obs.split("_")[1]  # "fy" or "fx"
                        y_train = make_force_obs(d["forces"][train_idx], name, mode=f_mode)
                        y_test = make_force_obs(d["forces"][test_idx], name, mode=f_mode)
                    else:
                        y_train = d["actions"][train_idx, :].T
                        y_test = d["actions"][test_idx, :].T

                    W, sigma, Rmat, _, _, _ = compute_reduced_ccd(a_train, y_train, Q_delay=CCD_Q)
                    recon = reconstruct_observable(W, sigma, Rmat, a_test, y_train)

                    ch_m1 = [r2(y_test[c], recon["mode1"][c]) for c in range(y_test.shape[0])]
                    ch_m80 = [r2(y_test[c], recon["m80"][c]) for c in range(y_test.shape[0])]
                    fold_r2_m1.append(float(np.mean(ch_m1)))
                    fold_r2_m80.append(float(np.mean(ch_m80)))

            if fold_r2_m1:
                key = f"LOCO_{obs}_r{r}"
                loco_results[key] = {
                    "mode1": {
                        "mean": float(np.mean(fold_r2_m1)),
                        "std": float(np.std(fold_r2_m1)),
                    },
                    "m80": {
                        "mean": float(np.mean(fold_r2_m80)),
                        "std": float(np.std(fold_r2_m80)),
                    },
                }
                print(f"  {key}: R2_m1={loco_results[key]['mode1']['mean']:.4f}+-"
                      f"{loco_results[key]['mode1']['std']:.4f}  "
                      f"R2_m80={loco_results[key]['m80']['mean']:.4f}+-"
                      f"{loco_results[key]['m80']['std']:.4f}", flush=True)

    # -- Blocked split --
    print("\n--- Blocked Split (train=0-47, test=48-95) ---", flush=True)
    blocked = {}
    test_idx_full = list(range(2 * N_PTS, 4 * N_PTS))
    for r in R_LIST:
        for obs in ["force_fy", "force_fx", "action"]:
            mf, modes_r = pod_cache[("blocked", 0, r)]
            per_case_m1, per_case_m80 = {}, {}

            for name, d, fld in [
                (tgt_key, tgt_d, tgt_f), (ill_key, ill_d, ill_f), ("pinball", unc_d, unc_f)
            ]:
                if obs == "action" and "illusion" not in name:
                    continue
                if d.get("forces") is None and "force" in obs:
                    continue

                a_train = modes_r.T @ (fld[:, train_idx_full] - mf[:, None])
                a_test = modes_r.T @ (fld[:, test_idx_full] - mf[:, None])

                if "force" in obs:
                    f_mode = obs.split("_")[1]
                    y_train = make_force_obs(d["forces"][train_idx_full], name, mode=f_mode)
                    y_test = make_force_obs(d["forces"][test_idx_full], name, mode=f_mode)
                else:
                    y_train = d["actions"][train_idx_full, :].T
                    y_test = d["actions"][test_idx_full, :].T

                W, sigma, Rmat, _, _, _ = compute_reduced_ccd(a_train, y_train, Q_delay=CCD_Q)
                recon = reconstruct_observable(W, sigma, Rmat, a_test, y_train)

                ch_m1 = [r2(y_test[c], recon["mode1"][c]) for c in range(y_test.shape[0])]
                ch_m80 = [r2(y_test[c], recon["m80"][c]) for c in range(y_test.shape[0])]
                per_case_m1[name] = float(np.mean(ch_m1))
                per_case_m80[name] = float(np.mean(ch_m80))

            key = f"blocked_{obs}_r{r}"
            blocked[key] = {
                "mode1": {
                    "mean": float(np.mean(list(per_case_m1.values()))),
                    "per_case": per_case_m1,
                },
                "m80": {
                    "mean": float(np.mean(list(per_case_m80.values()))),
                    "per_case": per_case_m80,
                },
            }
            print(f"  {key}: R2_m1={blocked[key]['mode1']['mean']:.4f}  "
                  f"R2_m80={blocked[key]['m80']['mean']:.4f}", flush=True)

    diam_results["LOCO"] = loco_results
    diam_results["blocked_split"] = blocked
    return diam_results


def run():
    print("=" * 60, flush=True)
    print("CCD Validation (Round 5)", flush=True)
    print("=" * 60, flush=True)

    t_start = time.time()
    all_results = {}

    for diam in DIAMETERS:
        tgt_key = f"target_cylinder_{diam}L"
        ill_key = f"illusion_{diam}L"
        print(f"\n{'=' * 60}", flush=True)
        print(f"Diameter {diam}L", flush=True)
        print(f"{'=' * 60}", flush=True)

        t0 = time.time()
        scene_data = {
            tgt_key: load_aligned_fields(tgt_key),
            ill_key: load_aligned_fields(ill_key),
            "pinball": load_aligned_fields("pinball"),
        }
        print(f"  Data loaded in {time.time() - t0:.0f}s", flush=True)

        t1 = time.time()
        all_results[f"{diam}L"] = run_single_diameter(diam, scene_data)
        print(f"  Analysis done in {time.time() - t1:.0f}s", flush=True)

    out_dir = os.path.join(DATA_DIR, "ccd")
    with open(os.path.join(out_dir, "validation_results.json"), "w") as f:
        json.dump(all_results, f, indent=2)

    print(f"\nTotal: {time.time() - t_start:.0f}s", flush=True)
    print(f"Saved to {out_dir}/validation_results.json", flush=True)


if __name__ == "__main__":
    run()
