"""Load legacy fields.npz format for steady scenes (steady_cloak, target_channel).

Converts to the same convention as load_aligned_fields():
    - Transposes fields from (N, NX, NY) -> (N, NY, NX)
    - Loads telemetry from sensors.npz
    - Returns dict with identical key structure
"""
from __future__ import annotations

import json
import os
from typing import Any

import numpy as np

from CCD_analysis.configs import DATA_DIR, NX, NY


def load_legacy_steady(scene_name: str) -> dict:
    """Load steady scene from legacy fields.npz format.

    Parameters
    ----------
    scene_name : str — one of 'steady_cloak' or 'target_channel'

    Returns
    -------
    dict with same keys as load_aligned_fields():
        ux, uy : (N, NY, NX) ndarray
        forces : None (no force telemetry in legacy sensors)
        sensors : (N, 6) ndarray or None
        actions : None (open-loop)
        meta : dict with scene info
        step_indices : list of int
    """
    scene_dir = os.path.join(DATA_DIR, scene_name, scene_name)
    if not os.path.isdir(scene_dir):
        raise FileNotFoundError(f"Scene directory not found: {scene_dir}")

    # -- fields.npz (native simulation order: NX first) --
    fields_path = os.path.join(scene_dir, "fields.npz")
    if not os.path.isfile(fields_path):
        raise FileNotFoundError(f"{fields_path} not found")

    fd = np.load(fields_path)
    ux_raw = fd["ux"]  # (N, NX, NY)
    uy_raw = fd["uy"]
    N = ux_raw.shape[0]
    fd.close()

    # Transpose (N, NX, NY) -> (N, NY, NX) to match load_aligned_fields convention
    ux = np.ascontiguousarray(ux_raw.transpose(0, 2, 1))
    uy = np.ascontiguousarray(uy_raw.transpose(0, 2, 1))

    # -- sensors.npz (telemetry) --
    sensors_path = os.path.join(scene_dir, "sensors.npz")
    sensors = None
    if os.path.isfile(sensors_path):
        sd = np.load(sensors_path)
        if "sensors" in sd:
            sensors = sd["sensors"]  # (N, 6)
            assert sensors.shape[0] == N, (
                f"sensors ({sensors.shape[0]}) != fields ({N})"
            )
        sd.close()

    # -- meta.json --
    meta = {"scene": scene_name, "scene_id": scene_name, "source": "legacy_steady"}
    meta_path = os.path.join(scene_dir, "meta.json")
    if os.path.isfile(meta_path):
        with open(meta_path) as f:
            meta.update(json.load(f))

    result: dict[str, Any] = {
        "ux": ux,
        "uy": uy,
        "forces": None,  # no forces in legacy steady telemetry
        "actions": None,  # open-loop
        "sensors": sensors,
        "meta": meta,
        "step_indices": list(range(N)),
    }

    return result


# ---------------------------------------------------------------------------
# Diagnostic helpers
# ---------------------------------------------------------------------------

def _print_field_stats(label: str, ux: np.ndarray, uy: np.ndarray) -> None:
    """Print mean velocity statistics for a set of fields."""
    ux_mean = ux.mean()
    uy_mean = uy.mean()
    ux_std = ux.std()
    uy_std = uy.std()
    print(f"  {label}:")
    print(f"    shape     = {ux.shape}")
    print(f"    ux_mean   = {ux_mean:.6f}  (expect ~U0={0.01} for channel)")
    print(f"    uy_mean   = {uy_mean:.6f}  (expect near 0)")
    print(f"    ux_rms    = {ux_std:.6f}")
    print(f"    uy_rms    = {uy_std:.6f}")


# ---------------------------------------------------------------------------
# Main (test)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("=" * 60)
    print("Testing load_legacy_steady()")
    print("=" * 60)

    for scene in ["steady_cloak", "target_channel"]:
        print(f"\n--- {scene} ---")
        data = load_legacy_steady(scene)
        ux = data["ux"]
        uy = data["uy"]
        sensors = data["sensors"]
        N = ux.shape[0]

        _print_field_stats(scene, ux, uy)

        print(f"  N_frames   = {N}")
        print(f"  NY x NX    = {ux.shape[1]} x {ux.shape[2]}")
        print(f"  sensors    = {sensors.shape if sensors is not None else None}")
        print(f"  forces     = {data['forces']}")
        print(f"  actions    = {data['actions']}")
        print(f"  step_range = [{data['step_indices'][0]}, {data['step_indices'][-1]}]")

        # Physical reasonableness checks
        ux_max = ux.max()
        uy_max = abs(uy).max()
        print(f"  ux_max     = {ux_max:.4f}  (expect order 0.01)")
        print(f"  |uy|_max   = {uy_max:.4f}  (expect < ux_max)")
        print(f"  metadata   = {list(data['meta'].keys())}")

    # Quick: verify convention matches load_aligned_fields
    print("\n--- Convention check: transpose correctness ---")
    # Load raw from steady_cloak to verify ravel order
    raw = np.load(
        os.path.join(DATA_DIR, "steady_cloak", "steady_cloak", "fields.npz")
    )
    raw_ux = raw["ux"][0]  # (NX, NY)
    loaded = load_legacy_steady("steady_cloak")
    loaded_ux = loaded["ux"][0]  # (NY, NX)

    # raw_ux[NX, NY] should == loaded_ux[NY, NX] after transpose
    match = np.allclose(raw_ux.T, loaded_ux)
    print(f"  Transpose (raw.T == loaded): {match}")
    raw.close()

    print("\nDone.")
