# OID_analysis/analysis/phase1_correction_pod.py
"""
Phase 1: Correction-field POD with rank sensitivity.

For each scene, computes:
  - Delta_q_blk = q_blk - q_in
  - Delta_q_ctl = q_ctl - q_blk
  - POD on Delta_q_ctl (scene-adaptive ROI mask to stay within ~2 GB RAM)
  - Rank sensitivity (r=6,8,10,12,16)

IMPORTANT: Full 1280x512 fields are LOADED but POD is computed on a wake ROI.
Full-field POD with 500 snapshots requires ~50 GB RAM (DOF = 1.31M).
ROI masking reduces DOF to ~180-360K (manageable ~1.5 GB).
Fields are SAVED as full resolution (OID Rule 5); ROI is applied only at analysis stage.

Usage:
    python3 src/OID_analysis/analysis/phase1_correction_pod.py
    python3 src/OID_analysis/analysis/phase1_correction_pod.py --scene karman_re100
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Dict, Tuple, Optional

import numpy as np

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

from OID_analysis.configs import (  # noqa: E402
    get_scene, data_dir_for_scene, DATA_DIR,
)
from OID_analysis.utils.analysis import compute_pod  # noqa: E402


# ---------------------------------------------------------------------------
# ROI mask (scene-adaptive, ~2 GB memory budget)
# ---------------------------------------------------------------------------

def get_scene_roi(scene_key: str) -> Tuple[int, int, int, int]:
    """Return (x_start, x_end, y_start, y_end) in pixel coordinates.

    Uses the empirically verified ROI from original Phase 1 runs:
    x=[400,1000], y=[100,400] for all scenes. This captures the full pinball
    wake region plus sensor zone across all geometry configurations.
    Memory: ~600×300×2 = 360K DOF, snapshot matrix ~1.4 GB for 500 steps.
    """
    return 400, 1000, 100, 400


def mask_field(ux: np.ndarray, uy: np.ndarray,
               x0: int, x1: int, y0: int, y1: int) -> Tuple[np.ndarray, np.ndarray]:
    """Crop field to ROI for POD analysis (memory constraint)."""
    return ux[:, y0:y1, x0:x1], uy[:, y0:y1, x0:x1]


# ---------------------------------------------------------------------------
# Scene groups
# ---------------------------------------------------------------------------

SCENE_GROUPS = {
    "steady_cloak": {
        "q_in_dir": data_dir_for_scene("empty_channel"),
        "q_blk_dir": data_dir_for_scene("pinball_baseline"),
        "q_ctl_dir": data_dir_for_scene("steady_cloak"),
    },
    "karman_re100": {
        "q_in_dir": data_dir_for_scene("disturbance_only"),
        "q_blk_dir": data_dir_for_scene("karman_blk"),
        "q_ctl_dir": data_dir_for_scene("karman_re100"),
    },
    "illusion_0.75L": {
        "q_in_dir": data_dir_for_scene("empty_channel"),
        "q_blk_dir": os.path.join(
            os.path.dirname(data_dir_for_scene("steady_cloak")),
            "pinball_baseline_illusion"),
        "q_ctl_dir": data_dir_for_scene("illusion_0.75L"),
    },
    "illusion_1.0L": {
        "q_in_dir": data_dir_for_scene("empty_channel"),
        "q_blk_dir": os.path.join(
            os.path.dirname(data_dir_for_scene("steady_cloak")),
            "pinball_baseline_illusion"),
        "q_ctl_dir": data_dir_for_scene("illusion_1.0L"),
    },
    "illusion_1.5L": {
        "q_in_dir": data_dir_for_scene("empty_channel"),
        "q_blk_dir": os.path.join(
            os.path.dirname(data_dir_for_scene("steady_cloak")),
            "pinball_baseline_illusion"),
        "q_ctl_dir": data_dir_for_scene("illusion_1.5L"),
    },
}


def load_scene_fields(scene_key: str) -> Optional[Dict]:
    """Load q_in, q_blk, q_ctl. Returns None if missing."""
    groups = SCENE_GROUPS.get(scene_key)
    if groups is None:
        print(f"  Unknown scene group: {scene_key}")
        return None

    result = {}
    for key, dir_path in groups.items():
        fp = os.path.join(dir_path, "fields.npz")
        if not os.path.isfile(fp):
            print(f"  WARNING: {key} fields not found at {fp}")
            return None
        fd = np.load(fp)
        result[key] = (fd["ux"], fd["uy"])
        print(f"  Loaded {key}: {fd['ux'].shape}")

    # Minimum snapshot count
    sizes = [v[0].shape[0] for v in result.values()]
    if len(set(sizes)) > 1:
        print(f"  Mismatched snapshot counts: {sizes}, using min={min(sizes)}")
        min_n = min(sizes)
        for k in result:
            result[k] = (result[k][0][:min_n], result[k][1][:min_n])

    return result


def fields_to_snapshot_matrix(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Convert (N, ny, nx) -> (N, DOF)."""
    N, ny, nx = ux.shape
    DOF = ny * nx * 2
    Q = np.zeros((N, DOF), dtype=np.float64)
    for t in range(N):
        Q[t] = np.concatenate([ux[t].ravel(), uy[t].ravel()])
    return Q


def run_phase1(scene_key: str):
    print(f"\n{'='*60}")
    print(f"Phase 1: Correction-field POD for {scene_key}")
    print(f"{'='*60}")

    fields = load_scene_fields(scene_key)
    if fields is None:
        print(f"  SKIPPED: data not available")
        return

    out_dir = os.path.join(DATA_DIR, "derived", "pod", scene_key)
    os.makedirs(out_dir, exist_ok=True)

    ux_in, uy_in = fields["q_in_dir"]
    ux_blk, uy_blk = fields["q_blk_dir"]
    ux_ctl, uy_ctl = fields["q_ctl_dir"]

    # Delta fields at full resolution
    delta_ux_blk = ux_blk - ux_in
    delta_uy_blk = uy_blk - uy_in
    delta_ux_ctl = ux_ctl - ux_blk
    delta_uy_ctl = uy_ctl - uy_blk

    # Save delta fields full-resolution (OID Rule 5)
    np.savez_compressed(os.path.join(out_dir, "delta_q_blk.npz"),
                        ux=delta_ux_blk, uy=delta_uy_blk)
    np.savez_compressed(os.path.join(out_dir, "delta_q_ctl.npz"),
                        ux=delta_ux_ctl, uy=delta_uy_ctl)

    # ROI analysis: full-field POD would need ~50 GB
    x0, x1, y0, y1 = get_scene_roi(scene_key)
    ux_dm, uy_dm = mask_field(delta_ux_ctl, delta_uy_ctl, x0, x1, y0, y1)
    ux_rm, uy_rm = mask_field(ux_ctl, uy_ctl, x0, x1, y0, y1)
    roi_area = (x1 - x0) * (y1 - y0)
    print(f"  ROI x=[{x0},{x1}) y=[{y0},{y1})  area={roi_area} px  DOF={roi_area*2}")

    # Snapshot matrices
    Q_delta = fields_to_snapshot_matrix(ux_dm, uy_dm)
    Q_raw = fields_to_snapshot_matrix(ux_rm, uy_rm)
    print(f"  Snapshot: {Q_delta.shape}  ~{Q_delta.nbytes/1e9:.1f} GB")

    # POD at ranks 6,8,10,12,16
    ranks = [6, 8, 10, 12, 16]
    results = {}
    prev_modes = None

    for r in ranks:
        if r > Q_delta.shape[0]:
            print(f"  Rank {r} > N, skip")
            continue
        pod = compute_pod(Q_delta, rank=r)
        results[f"r{r}"] = pod

        if prev_modes is not None:
            n_cmp = min(prev_modes.shape[1], pod["modes"].shape[1], 6)
            sims = [abs(np.dot(prev_modes[:, i], pod["modes"][:, i])) for i in range(n_cmp)]
            print(f"  r={r}: cum5={pod['cum_energy'][4]:.4f}  cos_sim={np.mean(sims):.4f}")
        prev_modes = pod["modes"]

    # Save
    for rk, pod in results.items():
        np.savez(os.path.join(out_dir, f"pod_coefs_{rk}.npy"),
                 coefs=pod["coefs"], S=pod["S"],
                 energy=pod["energy"], cum_energy=pod["cum_energy"])
        np.savez_compressed(os.path.join(out_dir, f"pod_modes_{rk}.npz"),
                            modes=pod["modes"], mean=pod["mean"])

    pod_raw = compute_pod(Q_raw, rank=10)
    np.savez(os.path.join(out_dir, "raw_pod_r10.npy"),
             coefs=pod_raw["coefs"], S=pod_raw["S"],
             energy=pod_raw["energy"], cum_energy=pod_raw["cum_energy"])

    summary = {
        "scene": scene_key,
        "n_snapshots": Q_delta.shape[0],
        "dof": Q_delta.shape[1],
        "roi": {"x_start": x0, "x_end": x1, "y_start": y0, "y_end": y1},
    }
    if "r10" in results:
        summary["energy_r10_5modes"] = float(results["r10"]["cum_energy"][4])
    with open(os.path.join(out_dir, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2)

    print(f"  Done.  Saved to {out_dir}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, default=None)
    ap.add_argument("--list", action="store_true")
    args = ap.parse_args()
    scenes = list(SCENE_GROUPS.keys())
    if args.list:
        print("Available:", scenes)
        return
    targets = scenes if (args.scene is None or args.scene == "all") else [args.scene]
    for sn in targets:
        if sn not in scenes:
            print(f"Unknown: {sn}")
            continue
        run_phase1(sn)
    print("\nPhase 1 complete.")


if __name__ == "__main__":
    main()
