"""CCD analysis pipeline: POD + force/action/signature CCD.

Usage:
    python ccd/run_ccd.py

Requires resampled data from scripts/resample.py.
"""
from __future__ import annotations

import json
import os
import sys

import numpy as np

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

from CCD_analysis.configs import DATA_DIR
from CCD_analysis.utils.resampling import (
    compute_pod, cumulative_energy, e95_index,
    compute_reduced_ccd, stack_velocity_fields,
)

R_CANDIDATES = [6, 8, 10]
CCD_Q = 12


def load_resampled(name: str):
    p = os.path.join(DATA_DIR, "resampled", name, "resampled.npz")
    if not os.path.isfile(p):
        return None
    return np.load(p)


def main():
    print("=== CCD Pipeline ===\n")

    # Identify which cases have resampled data
    resampled_dir = os.path.join(DATA_DIR, "resampled")
    if not os.path.isdir(resampled_dir):
        print("ERROR: run scripts/resample.py first")
        return 1

    cases = sorted(os.listdir(resampled_dir))
    print(f"Resampled cases: {cases}")

    # --- POD ---
    print("\n--- POD ---")
    snapshots = []
    case_ranges = {}
    idx = 0

    for name in cases:
        d = load_resampled(name)
        if d is None:
            continue
        ux, uy = d.get("ux"), d.get("uy")
        if ux is None:
            print(f"  {name}: no field data, skip POD")
            continue
        n_cyc, n_pt = ux.shape[0], ux.shape[1]
        for c in range(n_cyc):
            for p in range(n_pt):
                q = np.concatenate([ux[c, p].ravel(), uy[c, p].ravel()])
                snapshots.append(q)
        case_ranges[name] = (idx, idx + n_cyc * n_pt)
        idx += n_cyc * n_pt
        print(f"  {name}: {n_cyc}x{n_pt} snapshots")

    if not snapshots:
        print("No field data for POD")
        return 1

    Q = np.column_stack(snapshots)
    mean_field, modes, s, coeffs = compute_pod(Q)
    energy = cumulative_energy(s)
    e95 = e95_index(energy)
    print(f"  POD: {len(s)} modes, E95={e95}")
    for i in range(min(6, len(s))):
        print(f"    mode {i+1}: energy={energy[i]:.4f}")

    # --- CCD for each case ---
    print("\n--- CCD ---")
    all_results = {}
    W_dict = {}

    for r in R_CANDIDATES:
        print(f"\n  POD truncation r={r}")
        for name in cases:
            d = load_resampled(name)
            if d is None:
                continue

            # POD coefficients for this case
            if name in case_ranges:
                start, end = case_ranges[name]
                a_r = coeffs[:r, start:end]
            else:
                # Projection case (not in POD basis)
                ux, uy = d.get("ux"), d.get("uy")
                if ux is None:
                    continue
                proj_snapshots = []
                for c in range(ux.shape[0]):
                    for p in range(ux.shape[1]):
                        q = np.concatenate([ux[c, p].ravel(), uy[c, p].ravel()])
                        proj_snapshots.append(q)
                Q_proj = np.column_stack(proj_snapshots)
                Q_centered = Q_proj - mean_field[:, None]
                a_r = (modes[:, :r].T @ Q_centered)

            N = a_r.shape[1]
            if N < 24:
                print(f"  {name}: too few samples ({N})")
                continue

            # Force CCD
            forces = d.get("forces")
            if forces is not None:
                f = forces.reshape(-1, forces.shape[-1])
                Fx = f[:, 0] + f[:, 2] + f[:, 4]
                Fy = f[:, 1] + f[:, 3] + f[:, 5]
                y_force = np.vstack([Fx, Fy])

                if y_force.shape[1] >= N:
                    y_f = y_force[:, :N]
                else:
                    y_f = y_force

                W, sigma, z = compute_reduced_ccd(a_r[:, :y_f.shape[1]], y_f, Q_delay=CCD_Q)
                ccd_ene = cumulative_energy(sigma)
                m80 = int(np.searchsorted(ccd_ene, 0.80) + 1) if len(ccd_ene) > 0 else 0
                key = f"{name}_force_r{r}"
                W_dict[key] = W
                all_results[key] = {"case": name, "observable": "force", "r": r,
                                    "m80": m80, "sigma_top3": [float(sigma[i]) for i in range(min(3, len(sigma)))]}
                print(f"    {key}: m80={m80}")

            # Action CCD (for controlled cases)
            actions = d.get("actions")
            if actions is not None:
                y_act = actions.reshape(-1, actions.shape[-1]).T
                if y_act.shape[1] >= N:
                    y_a = y_act[:, :N]
                else:
                    y_a = y_act
                W, sigma, z = compute_reduced_ccd(a_r[:, :y_a.shape[1]], y_a, Q_delay=CCD_Q)
                ccd_ene = cumulative_energy(sigma)
                m80 = int(np.searchsorted(ccd_ene, 0.80) + 1) if len(ccd_ene) > 0 else 0
                key = f"{name}_action_r{r}"
                W_dict[key] = W
                all_results[key] = {"case": name, "observable": "action", "r": r,
                                    "m80": m80, "sigma_top3": [float(sigma[i]) for i in range(min(3, len(sigma)))]}
                print(f"    {key}: m80={m80}")

    # --- Modal overlap ---
    print("\n--- Modal Overlap ---")
    force_keys = [k for k in W_dict if "force" in k]
    for i, ka in enumerate(force_keys):
        for kb in force_keys[i+1:]:
            Wa, Wb = W_dict[ka], W_dict[kb]
            n = min(Wa.shape[1], Wb.shape[1], 5)
            ov = []
            for k in range(n):
                ak = Wa[:, k] / (np.linalg.norm(Wa[:, k]) + 1e-12)
                bk = Wb[:, k] / (np.linalg.norm(Wb[:, k]) + 1e-12)
                ov.append(float(abs(ak @ bk)))
            print(f"    O({ka}, {kb}): O1={ov[0]:.4f}, O2={ov[1]:.4f}")

    # Save
    out_dir = os.path.join(DATA_DIR, "ccd")
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "ccd_results.json"), "w") as f:
        json.dump(all_results, f, indent=2)
    print(f"\nSaved to {out_dir}/ccd_results.json")
    return 0


if __name__ == "__main__":
    sys.exit(main())