# OID_analysis/analysis/phase2_build_observables.py
"""
Phase 2: Scene-specific observable construction.

For each scene, builds standardized observable matrices Y from sensor/force/action data.
All outputs saved to data/derived/observables/<scene>/

Usage:
    python3 src/OID_analysis/analysis/phase2_build_observables.py
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Dict, 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, SCENES, DATA_DIR, CONV_LEN_DEFAULT, CONV_LEN_ILLUSION,
)
from OID_analysis.utils.analysis import standardize  # noqa: E402


def build_steady_observables(scene_key: str) -> Optional[Dict]:
    """Build observables for steady cloak."""
    dd = data_dir_for_scene(scene_key)

    # Load forces
    fp = os.path.join(dd, "forces.npz")
    if not os.path.isfile(fp):
        print(f"  WARNING: forces not found at {fp}")
        return None
    fd = np.load(fp)
    forces = fd["forces"]  # (N, 6): Fx,Fy per cylinder
    N = forces.shape[0]

    # Load sensors
    sp = os.path.join(dd, "sensors.npz")
    sensors = np.load(sp)["sensors"] if os.path.isfile(sp) else np.zeros((N, 6))

    # Total force
    Fx_total = np.sum(forces[:, 0::2], axis=1, keepdims=True)  # (N, 1)
    Fy_total = np.sum(forces[:, 1::2], axis=1, keepdims=True)  # (N, 1)

    # Control power (requires torque -- not available from legacy env directly)
    # Use force magnitude as proxy
    force_mag = np.sqrt(forces[:, 0::2]**2 + forces[:, 1::2]**2)  # (N, 3)
    total_force_mag = np.sum(force_mag, axis=1, keepdims=True)

    # Fluctuation observable: windowed RMS of sensor uy
    conv_len = CONV_LEN_DEFAULT
    rms_uy = np.zeros((N, 1))
    for t in range(conv_len, N):
        rms_uy[t] = np.std(sensors[t-conv_len:t, 1::2])  # all uy channels
    # First conv_len steps: pad with first valid value
    rms_uy[:conv_len] = rms_uy[conv_len]

    # Mean wake deviation: difference from downstream sensor mean to uniform
    # Use sensor ux as proxy for wake recovery
    ux_mean = sensors[:, 0::2]  # (N, 3)
    ux_deviation = np.std(ux_mean, axis=1, keepdims=True)  # cross-sensor std

    Y = {
        "force_total": np.hstack([Fx_total, Fy_total]),
        "force_mag": total_force_mag,
        "rms_uy": rms_uy,
        "ux_deviation": ux_deviation,
    }
    return Y


def build_karman_observables(scene_key: str) -> Optional[Dict]:
    """Build observables for Karman cloak."""
    dd = data_dir_for_scene(scene_key)
    fp = os.path.join(dd, "controlled.npz")
    if not os.path.isfile(fp):
        print(f"  WARNING: controlled.npz not found at {fp}")
        return None

    data = np.load(fp)
    sensors = data["sensors"]  # (N, 6)
    forces = data["forces"]    # (N, 6)
    actions = data["actions"]  # (N, 3)
    N = len(sensors)

    # Load target
    tp = os.path.join(dd, "target.npz")
    target_states = np.load(tp)["target_states"] if os.path.isfile(tp) else None

    # Sensor error
    if target_states is not None:
        # Align to target length, repeat if needed
        tlen = target_states.shape[0]
        if N > tlen:
            e_s = sensors[:tlen] - target_states  # aligned portion
            # Pad with last values for remaining
            pad = np.tile(e_s[-1:], (N - tlen, 1))
            e_s = np.vstack([e_s, pad])
        else:
            e_s = sensors - target_states[:N]
    else:
        e_s = sensors * 0.0  # no target

    # Total force
    Fx_total = np.sum(forces[:, 0::2], axis=1, keepdims=True)
    Fy_total = np.sum(forces[:, 1::2], axis=1, keepdims=True)

    # Delay estimate: pinball~sensor distance ~ 10*L0, Uconv ~ U0
    # tau_c steps = 10*L0 / (U0 * SI) = 200 / (0.01 * 800) = 25
    tau_c = 25
    # Create delayed sensor error
    e_s_delayed = np.zeros_like(e_s)
    if N > tau_c:
        e_s_delayed[:N-tau_c] = e_s[tau_c:]
        e_s_delayed[N-tau_c:] = e_s[-1]

    # Full delay stack (3 time points)
    tau_quarter = 50 // (800 // 200)  # ~12-13 steps for ~1/4 period at Re=100
    tau_q = max(8, tau_quarter)
    p_sig_stack = np.zeros((N, 18))
    if N > tau_c + 2 * tau_q:
        p_sig_stack[:N-tau_c-2*tau_q, 0:6] = e_s[tau_c: N-2*tau_q]
        p_sig_stack[:N-tau_c-2*tau_q, 6:12] = e_s[tau_c+tau_q: N-tau_q]
        p_sig_stack[:N-tau_c-2*tau_q, 12:18] = e_s[tau_c+2*tau_q: N]
        # Pad tail
        for i in range(1, 3):
            p_sig_stack[N-i:, :] = p_sig_stack[N-i-1, :]

    Y = {
        "force_total": np.hstack([Fx_total, Fy_total]),
        "sensor_error": e_s,
        "sensor_error_delayed": e_s_delayed,
        "p_sig_stack": p_sig_stack,
        "actions": actions,
    }
    return Y


def build_illusion_observables(scene_key: str) -> Optional[Dict]:
    """Build observables for illusion scenes."""
    dd = data_dir_for_scene(scene_key)
    fp = os.path.join(dd, "controlled.npz")
    if not os.path.isfile(fp):
        print(f"  WARNING: controlled.npz not found at {fp}")
        return None

    data = np.load(fp)
    sensors = data["sensors"]
    forces = data["forces"]
    actions = data["actions"]
    N = len(sensors)

    # Target sensors from target cylinder (stored in target_cylinder dir)
    diam = scene_key.split("_")[1]  # e.g. "0.75L"
    tgt_key = f"target_cylinder_{diam}"
    tgt_dd = data_dir_for_scene(tgt_key)
    tp = os.path.join(tgt_dd, "target.npz")
    if os.path.isfile(tp):
        tgt = np.load(tp)
        target_states = tgt["target_states"]
        # target_states shape: (FIFO_LEN, 8) - [cylinder_fx,fy, s0_ux,uy, s1_ux,uy, s2_ux,uy]
        target_sensors = target_states[:, 2:8]  # (150, 6)
    else:
        target_sensors = sensors * 0.0

    # Sensor error
    tlen = target_sensors.shape[0]
    if N > tlen:
        e_s = sensors[:tlen] - target_sensors
        pad = np.tile(e_s[-1:], (N - tlen, 1))
        e_s = np.vstack([e_s, pad])
    else:
        e_s = sensors - target_sensors[:N]

    # Force observable
    Fx_total = np.sum(forces[:, 0::2], axis=1, keepdims=True)
    Fy_total = np.sum(forces[:, 1::2], axis=1, keepdims=True)

    # Delay (illusion: sensor_x=30, pinball_rear_x=20.3, distance ~10*L0)
    # tau_c depends on SI: 200/(0.01*SI)
    cfg = get_scene(scene_key)
    si = cfg["sample_interval"]
    tau_c = max(10, int(200 / (0.01 * si)))
    tau_q = max(5, int(50 / (0.01 * si) / 4))

    e_s_delayed = np.zeros_like(e_s)
    if N > tau_c:
        e_s_delayed[:N-tau_c] = e_s[tau_c:]
        e_s_delayed[N-tau_c:] = e_s[-1]

    # Full delay stack
    p_sig_stack = np.zeros((N, 18))
    if N > tau_c + 2 * tau_q:
        p_sig_stack[:N-tau_c-2*tau_q, 0:6] = e_s[tau_c: N-2*tau_q]
        p_sig_stack[:N-tau_c-2*tau_q, 6:12] = e_s[tau_c+tau_q: N-tau_q]
        p_sig_stack[:N-tau_c-2*tau_q, 12:18] = e_s[tau_c+2*tau_q: N]
        for i in range(1, 3):
            p_sig_stack[N-i:, :] = p_sig_stack[N-i-1, :]

    Y = {
        "force_total": np.hstack([Fx_total, Fy_total]),
        "sensor_error": e_s,
        "sensor_error_delayed": e_s_delayed,
        "p_sig_stack": p_sig_stack,
        "actions": actions,
    }
    return Y


def run_phase2(scene_key: str):
    print(f"\n--- Phase 2: Observables for {scene_key} ---")

    if scene_key.startswith("steady") or scene_key == "steady_cloak":
        Y = build_steady_observables(scene_key)
    elif scene_key.startswith("karman"):
        Y = build_karman_observables(scene_key)
    elif scene_key.startswith("illusion"):
        Y = build_illusion_observables(scene_key)
    else:
        print(f"  Unknown scene type: {scene_key}")
        return

    if Y is None:
        print(f"  SKIPPED")
        return

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

    # Standardize each observable and save
    for name, arr in Y.items():
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        arr_std, mean, std = standardize(arr)
        np.savez(os.path.join(out_dir, f"{name}.npz"),
                 raw=arr, standardized=arr_std, mean=mean, std=std)
        print(f"  {name}: {arr.shape}")

    # Save metadata
    meta = {"scene": scene_key, "n_steps": len(Y[list(Y.keys())[0]]),
            "observables": list(Y.keys())}
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)

    print(f"  Saved to {out_dir}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, default=None)
    args = ap.parse_args()

    scenes = ["steady_cloak", "karman_re100",
              "illusion_0.75L", "illusion_1.0L", "illusion_1.5L"]

    if args.scene:
        targets = [args.scene] if args.scene in scenes else scenes
    else:
        targets = scenes

    for sn in targets:
        run_phase2(sn)

    print("Phase 2 complete.")


if __name__ == "__main__":
    main()
