# OID_analysis/analysis/phase3a_oid_mode_viz.py
"""
Phase 3a: OID Force/Signature Mode Spatial Visualization & Zone Partition

Loads OID spatial modes from derived/oid/, reshapes to velocity/vorticity fields,
projects onto CCD's three-zone masks, and outputs:
  - Partition energy table JSON
  - Multi-panel figure: (force+sig) x 3 modes x (ux, uy, vorticity)

Pure CPU. No GPU. Derived data only (small).

Usage:
  PYTHONPATH="src:$PYTHONPATH" python3 src/OID_analysis/analysis/phase3a_oid_mode_viz.py
"""
import json, os, sys
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 DATA_DIR, get_scene  # noqa: E402

# ROI coords used in Phase 1 POD (hardcoded consistently)
ROI_X0, ROI_X1 = 400, 1000
ROI_Y0, ROI_Y1 = 100, 400
NY_ROI = ROI_Y1 - ROI_Y0  # 300
NX_ROI = ROI_X1 - ROI_X0  # 600

SCENES = ["steady_cloak", "karman_re100",
          "illusion_0.75L", "illusion_1.0L", "illusion_1.5L"]
DERIVED = os.path.join(DATA_DIR, "derived")

# ---------------------------------------------------------------------------
# Zone mask definitions (in ROI pixel coordinates)
# ---------------------------------------------------------------------------

def build_zone_masks(scene_key: str):
    """Build near_body, body_wake, sensor_zone bool masks in ROI frame."""
    cfg = get_scene(scene_key)
    L0 = 20.0
    front_x = cfg["pinball_front_x"] * L0  # lattice coords
    rear_x = cfg.get("pinball_rear_x", front_x + 1.3) * L0
    sensor_x = cfg["sensor_x"] * L0

    # Convert global -> ROI-local coords
    def gx(x): return x - ROI_X0
    def gy(y): return y - ROI_Y0

    ny, nx = NY_ROI, NX_ROI
    center_y = 255.5  # global CENTER_Y

    masks = {}
    # near_body: pinball ± 2D
    m = np.zeros((ny, nx), dtype=bool)
    x0 = max(0, int(gx(front_x - 2*L0)))
    x1 = min(nx, int(gx(rear_x + 2*L0)))
    y0 = max(0, int(gy(center_y - 2*L0)))
    y1 = min(ny, int(gy(center_y + 2*L0)))
    m[y0:y1, x0:x1] = True
    masks["near_body"] = m

    # body_wake: from right of near_body to sensors
    m = np.zeros((ny, nx), dtype=bool)
    x0 = int(gx(rear_x + 2*L0))
    x1 = min(nx, int(gx(sensor_x - L0)))
    y0 = max(0, int(gy(center_y - 3*L0)))
    y1 = min(ny, int(gy(center_y + 3*L0)))
    if x1 > x0:
        m[y0:y1, x0:x1] = True
    masks["body_wake"] = m

    # sensor_zone: ±2D around sensor x
    m = np.zeros((ny, nx), dtype=bool)
    x0 = max(0, int(gx(sensor_x - 2*L0)))
    x1 = min(nx, int(gx(sensor_x + 2*L0)))
    y0 = max(0, int(gy(center_y - 3*L0)))
    y1 = min(ny, int(gy(center_y + 3*L0)))
    m[y0:y1, x0:x1] = True
    masks["sensor_zone"] = m

    return masks


# ---------------------------------------------------------------------------
# OID mode loading & reshaping
# ---------------------------------------------------------------------------

def load_oid_mode_fields(oid_type: str, scene_key: str, k: int):
    """Load OID mode k as (ux, uy, vorticity) arrays of shape (NY_ROI, NX_ROI)."""
    fpath = os.path.join(DERIVED, "oid", oid_type, scene_key,
                         f"{oid_type}_oid_modes.npz")
    if not os.path.isfile(fpath):
        raise FileNotFoundError(f"Missing: {fpath}")
    data = np.load(fpath)
    # 'modes' key: shape (DOF_ROI, r)
    modes_all = data["modes"]  # (DOF, r)
    dof_half = NY_ROI * NX_ROI
    ux = modes_all[:dof_half, k].reshape(NY_ROI, NX_ROI)
    uy = modes_all[dof_half:, k].reshape(NY_ROI, NX_ROI)
    # Vorticity via central difference
    vort = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
    return ux, uy, vort


# ---------------------------------------------------------------------------
# Zone energy partitioning
# ---------------------------------------------------------------------------

def zone_energy_fraction(mode_field: np.ndarray, masks: dict) -> dict:
    """Fraction of L2 energy of mode_field in each zone."""
    total = np.sum(mode_field ** 2)
    if total < 1e-30:
        return {k: 0.0 for k in masks}
    return {k: float(np.sum(mode_field[masks[k]] ** 2) / total) for k in masks}


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def compute_all():
    results = {}
    for scene in SCENES:
        masks = build_zone_masks(scene)
        scene_result = {}
        for otype in ["force", "signature"]:
            fpath = os.path.join(DERIVED, "oid", otype, scene,
                                 f"{otype}_oid_modes.npz")
            if not os.path.isfile(fpath):
                print(f"  SKIP {scene}/{otype}: no modes file")
                continue
            for km in range(3):  # modes 0,1,2
                try:
                    ux, uy, vort = load_oid_mode_fields(otype, scene, km)
                except Exception as e:
                    print(f"  SKIP {scene}/{otype}/mode{km}: {e}")
                    continue
                key = f"{otype}_m{km}"
                scene_result[key] = {
                    "ux_energy": zone_energy_fraction(ux, masks),
                    "uy_energy": zone_energy_fraction(uy, masks),
                    "vort_energy": zone_energy_fraction(vort, masks),
                }
        results[scene] = scene_result
        n = len(scene_result)
        print(f"  {scene}: {n} mode-fields computed")

    # Save
    out_dir = os.path.join(DERIVED, "oid_viz")
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, "zone_energy.json"), "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved to {out_dir}/zone_energy.json")

    # Print key table
    print("\nForce-OID mode 1 ux energy partition:")
    print(f"{'Scene':20s} {'near_body':>10s} {'body_wake':>10s} {'sensor_zone':>12s}")
    for scene in SCENES:
        if "force_m0" in results.get(scene, {}):
            r = results[scene]["force_m0"]["ux_energy"]
            print(f"{scene:20s} {r['near_body']:10.4f} {r['body_wake']:10.4f} "
                  f"{r['sensor_zone']:12.4f}")

    print("\nSignature-OID mode 1 ux energy partition:")
    print(f"{'Scene':20s} {'near_body':>10s} {'body_wake':>10s} {'sensor_zone':>12s}")
    for scene in SCENES:
        if "signature_m0" in results.get(scene, {}):
            r = results[scene]["signature_m0"]["ux_energy"]
            print(f"{scene:20s} {r['near_body']:10.4f} {r['body_wake']:10.4f} "
                  f"{r['sensor_zone']:12.4f}")

    return results


if __name__ == "__main__":
    compute_all()
