# OID_analysis/analysis/phase7_whitebox.py
"""
Phase 7: White-box control chain comparison.

Compares how well different state representations predict the action:
  Model A: obs (raw sensor) -> act
  Model B: POD coord -> act
  Model C: OID coord -> act
  Model D: OID coord + force -> act

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

import argparse
import json
import os
import 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  # noqa: E402
from OID_analysis.utils.analysis import standardize  # noqa: E402

try:
    from sklearn.linear_model import LinearRegression
    from sklearn.metrics import r2_score
    HAS_SKLEARN = True
except ImportError:
    HAS_SKLEARN = False

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


def run_whitebox(scene_key: str):
    print(f"\n--- Phase 7: White-box for {scene_key} ---")

    # Check for controlled.npz (PPO scenes) or forces (open-loop scenes)
    data_dir_base = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                                "data")

    if scene_key == "steady_cloak":
        dd = os.path.join(data_dir_base, "steady_cloak", "steady_cloak")
    else:
        sid = {"karman_re100": "karman_re100"}.get(scene_key, scene_key)
        dd = os.path.join(data_dir_base, scene_key.replace("steady_", ""),
                          scene_key) if "illusion" in scene_key else \
             os.path.join(data_dir_base, "karman_cloak", scene_key)

    controlled_fp = os.path.join(dd, "controlled.npz")
    forces_fp = os.path.join(dd, "forces.npz")

    if os.path.isfile(controlled_fp):
        data = np.load(controlled_fp)
        actions = data["actions"]
        sensors = data["sensors"]
    elif os.path.isfile(forces_fp):
        # Open-loop steady cloak: no actions available from controlled.npz
        # But we know the steady cloak action: [0, -5.1*U0, 5.1*U0]
        from OID_analysis.configs import get_scene
        cfg = get_scene(scene_key)
        u0 = cfg["u0"]
        sensors_n = np.load(os.path.join(dd, "sensors.npz"))["sensors"]
        N = len(sensors_n)
        sensors = sensors_n
        omega_rear = cfg.get("omega_rear_scale", 5.1)
        actions = np.tile([0.0, -omega_rear * u0 / 0.01, omega_rear * u0 / 0.01], (N, 1))
        # Actually these should be in normalized [-1,1] range
        # rear = 5.1 -> normalized = (5.1 - bias)/scale where bias=5.1, scale=8
        # Actually for steady: bias=[0,-5.1,5.1], scale=8
        # So action = (omega/u0 - bias)/scale
        actions = np.tile([0.0, 0.0, 0.0], (N, 1))  # zero action = bias actions
    else:
        print(f"  SKIPPED: no action data found")
        return

    N = min(len(sensors), len(actions))
    sensors = sensors[:N]
    actions = actions[:N]

    # Normalize actions per channel
    actions_std, act_mean, act_std = standardize(actions)

    # POD coefs
    pod_dir = os.path.join(DATA_DIR, "derived", "pod", scene_key)
    coefs = None
    pod_fp = os.path.join(pod_dir, "pod_coefs_r10.npy.npz")
    if os.path.isfile(pod_fp):
        pod_npz = np.load(pod_fp, allow_pickle=True)
        pod_n = pod_npz["coefs"].shape[0]
        N = min(N, pod_n)
    
    # Re-apply truncation based on final N
    sensors = sensors[:N]
    actions = actions[:N]
    actions_std, act_mean, act_std = standardize(actions)

    # OID coords
    oid_dir = os.path.join(DATA_DIR, "derived", "oid")
    oid_coords = None
    oid_fp = os.path.join(oid_dir, "force", scene_key, "force_oid.npz")
    if os.path.isfile(oid_fp):
        oid_coords = np.load(oid_fp)["z"][:N]

    # Force observable
    obs_dir = os.path.join(DATA_DIR, "derived", "observables", scene_key)
    force_obs = None
    force_fp = os.path.join(obs_dir, "force_total.npz")
    if os.path.isfile(force_fp):
        force_obs = np.load(force_fp)["standardized"][:N]

    if not HAS_SKLEARN:
        print("  sklearn not available")
        return

    split = int(N * 0.7)

    # Model A: raw sensor -> act
    X_A = sensors[:split]
    Y_train = actions_std[:split]
    # Test on last segment
    X_A_test = sensors[split:N]

    results = {}

    # Model A
    if X_A.shape[1] > 0:
        reg = LinearRegression().fit(X_A, Y_train)
        r2_a = r2_score(Y_train, reg.predict(X_A))
        results["obs_act_train"] = float(r2_a)

    # Model B: POD coord -> act
    if coefs is not None:
        for m in [3, 5]:
            X = standardize(coefs[:split, :m])[0]
            reg = LinearRegression().fit(X, Y_train)
            r2 = r2_score(Y_train, reg.predict(X))
            results[f"pod_m{m}_act_train"] = float(r2)

    # Model C: OID coord -> act
    if oid_coords is not None:
        for m in [3, 5]:
            X = oid_coords[:split, :m]
            reg = LinearRegression().fit(X, Y_train)
            r2 = r2_score(Y_train, reg.predict(X))
            results[f"oid_m{m}_act_train"] = float(r2)

    # Model D: OID + force -> act
    if oid_coords is not None and force_obs is not None:
        X = np.hstack([oid_coords[:split, :3], force_obs[:split, :2]])
        reg = LinearRegression().fit(X, Y_train)
        r2 = r2_score(Y_train, reg.predict(X))
        results["oid_force_act_train"] = float(r2)

    print(f"  Results: {json.dumps(results, indent=2)}")

    # Save
    out_dir = os.path.join(DATA_DIR, "derived", "whitebox")
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, f"{scene_key}.json"), "w") as f:
        json.dump(results, f, indent=2)


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

    targets = [args.scene] if args.scene and args.scene in SCENES else SCENES

    for sn in targets:
        run_whitebox(sn)

    print("\nPhase 7 complete.")


if __name__ == "__main__":
    main()