"""Load vortex scene fields.npz format for transient CCD analysis.

Vortex scenes (vortex_lamb, vortex_taylor, vortex_target_*, vortex_uncontrolled_*)
have field snapshots saved as raw fields.npz with shape (N, NX, NY) and no phase
plan (transient, not periodic).

Converts to the same convention as load_aligned_fields():
    - Transposes fields from (N, NX, NY) -> (N, NY, NX)
    - Loads telemetry from controlled.npz or 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, SCENES


def load_vortex_fields(scene_name: str) -> dict:
    """Load vortex scene field data from raw fields.npz.

    Parameters
    ----------
    scene_name : str — one of the vortex scene names registered in configs.

    Returns
    -------
    dict with same keys as load_aligned_fields():
        ux, uy : (N, NY, NX) ndarray — field snapshots (transposed)
        forces : (N, 6) ndarray or None
        sensors : (N, 6) ndarray or None
        actions : (N, 3) ndarray or None
        meta : dict with scene info
        step_indices : list of int (sequential 0..N-1 for transient)
    """
    if scene_name not in SCENES:
        raise KeyError(f"Unknown scene: {scene_name}")

    cfg = SCENES[scene_name]
    scene_id = cfg["scene_id"]
    data_dir = os.path.join(DATA_DIR, scene_id, scene_name)

    if not os.path.isdir(data_dir):
        raise FileNotFoundError(f"Vortex scene directory not found: {data_dir}")

    # -- fields.npz (native simulation order: NX first) --
    fields_path = os.path.join(data_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))

    # -- Telemetry (controlled.npz or sensors.npz) --
    tele_path = None
    for p in [
        os.path.join(data_dir, "controlled.npz"),
        os.path.join(data_dir, "sensors.npz"),
    ]:
        if os.path.isfile(p):
            tele_path = p
            break

    sensors, forces, actions = None, None, None
    if tele_path is not None:
        td = np.load(tele_path)
        if "sensors" in td:
            sensors = td["sensors"]  # (N, 6)
        if "forces" in td:
            forces = td["forces"]    # (N, 6)
        if "actions" in td:
            actions = td["actions"]  # (N, 3)
        td.close()

    # Verify N matches
    if sensors is not None and sensors.shape[0] != N:
        raise ValueError(
            f"sensors ({sensors.shape[0]}) != fields ({N})"
        )
    if forces is not None and forces.shape[0] != N:
        raise ValueError(
            f"forces ({forces.shape[0]}) != fields ({N})"
        )

    meta = {
        "scene": scene_name,
        "scene_id": scene_id,
        "source": "vortex_transient",
        "target_type": "transient",
        "n_frames": N,
    }

    result: dict[str, Any] = {
        "ux": ux,
        "uy": uy,
        "forces": forces,
        "actions": actions,
        "sensors": sensors,
        "meta": meta,
        "step_indices": list(range(N)),  # sequential — no phase plan
    }

    return result
