"""Build q_in, q_blk, q_ctl, q_tar field references and compute Delta corrections.

Each scene type maps to different data sources:

| Scene type      | q_in              | q_blk   | q_ctl               | q_tar                    |
|-----------------|-------------------|---------|---------------------|--------------------------|
| illusion_0.75L  | target_channel*   | pinball | illusion_0.75L      | target_cylinder_0.75L    |
| illusion_1.0L   | target_channel*   | pinball | illusion_1.0L       | target_cylinder_1.0L     |
| illusion_1.5L   | target_channel*   | pinball | illusion_1.5L       | target_cylinder_1.5L     |
| steady_cloak    | target_channel*   | pinball | steady_cloak*       | None (target=q_in)       |
| karman_re100    | karman_q_in       | karman_q_blk | karman_re100    | karman_q_in              |

(*) loaded via load_legacy_steady(). Vortex scenes use load_vortex_fields().
    All scenes now use UNIFIED geometry (pinball center at 613px, sensors at 800px).
"""
from __future__ import annotations

import os
import sys
import warnings
from typing import Any, Optional

import numpy as np

from CCD_analysis.configs import DATA_DIR, NX, NY
from CCD_analysis.utils.resampling import (
    load_aligned_fields,
    build_field_matrix as _build_field_matrix,
)

from CCD_analysis.correction_analysis.process_legacy_steady import load_legacy_steady
from CCD_analysis.utils.load_vortex_fields import load_vortex_fields


# ---------------------------------------------------------------------------
# Source mapping
# ---------------------------------------------------------------------------

def _resolve_source(name: str) -> Optional[dict]:
    """Load a named data source, dispatching to the correct loader.

    Parameters
    ----------
    name : str
        Scene or special name:
        - 'target_channel', 'steady_cloak' -> load_legacy_steady
        - 'vortex_*'                       -> load_vortex_fields (transient)
        - 'karman_re100'                   -> _load_karman_re100 (handles 72 vs 96 mismatch)
        - all others                       -> load_aligned_fields

    Returns
    -------
    dict or None
    """
    _LEGACY_SCENES = {"target_channel", "steady_cloak"}
    _VORTEX_SCENES = {
        "vortex_lamb", "vortex_taylor",
        "vortex_uncontrolled_lamb", "vortex_uncontrolled_taylor",
        "vortex_target_lamb", "vortex_target_taylor",
    }

    if name in _LEGACY_SCENES:
        return load_legacy_steady(name)
    elif name in _VORTEX_SCENES:
        return load_vortex_fields(name)
    elif name == "karman_re100":
        return _load_karman_re100()
    else:
        return load_aligned_fields(name)


# ---------------------------------------------------------------------------
# Karman re100 special loader (handles 72 vs 96 frame mismatch)
# ---------------------------------------------------------------------------

def _load_karman_re100() -> dict:
    """Load karman_re100 aligned fields, handling the 72 vs 96 frame mismatch.

    karman_re100's fields_aligned.npz has 72 snapshots (3 cycles x 24 pts)
    but the phase_plan.json lists 96 step_indices (4 cycles x 24 pts).
    This loader truncates step_indices to match.
    """
    import json as _json

    scene_name = "karman_re100"
    from CCD_analysis.configs import SCENES as _SCENES

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

    # Load fields_aligned.npz
    fa_path = os.path.join(data_dir, "fields_aligned.npz")
    fd = np.load(fa_path)
    ux_raw = fd["ux"]
    uy_raw = fd["uy"]
    N = ux_raw.shape[0]  # 72
    fd.close()

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

    # Load phase_plan, truncate to N
    plan_path = os.path.join(DATA_DIR, "resampled", scene_name, "phase_plan.json")
    with open(plan_path) as f:
        plan = _json.load(f)
    step_indices = list(plan["step_indices"][:N])

    # Load telemetry
    tele_path = os.path.join(data_dir, "controlled.npz")
    td = np.load(tele_path)

    result = {
        "ux": ux,
        "uy": uy,
        "forces": td["forces"][step_indices] if "forces" in td else None,
        "actions": td["actions"][step_indices] if "actions" in td else None,
        "sensors": td["sensors"][step_indices] if "sensors" in td else None,
        "meta": {
            "scene": scene_name,
            "scene_id": scene_id,
            "gate": plan.get("gate", "unknown"),
            "CV_T": plan.get("CV_T"),
            "f_dom": plan.get("f_dom"),
            "N_raw_per_cycle": plan.get("N_raw_per_cycle"),
            "rho_interp": plan.get("rho_interp"),
            "sample_interval": cfg.get("sample_interval"),
            "note": "step_indices truncated from 96 to 72 (3 cycles, not 4)",
        },
        "step_indices": step_indices,
    }
    td.close()
    return result


# ---------------------------------------------------------------------------
# Correction computation
# ---------------------------------------------------------------------------

# Scene map: (scene_type -> (q_in_source, q_blk_source, q_ctl_source, q_tar_source))
_SCENE_MAP = {
    "illusion_0.75L": ("target_channel", "pinball", "illusion_0.75L", "target_cylinder_0.75L"),
    "illusion_1.0L":  ("target_channel", "pinball", "illusion_1.0L",  "target_cylinder_1.0L"),
    "illusion_1.5L":  ("target_channel", "pinball", "illusion_1.5L",  "target_cylinder_1.5L"),
    "steady_cloak":   ("target_channel", "pinball", "steady_cloak",   None),
    "karman_re100":   ("karman_q_in",    "karman_q_blk", "karman_re100", "karman_q_in"),
    "vortex_lamb":    ("target_channel", "vortex_uncontrolled_lamb",  "vortex_lamb",  "vortex_target_lamb"),
    "vortex_taylor":  ("target_channel", "vortex_uncontrolled_taylor", "vortex_taylor", "vortex_target_taylor"),
}


def get_diameter(scene_type: str) -> Optional[float]:
    """Extract target diameter from scene type string (e.g. 'illusion_1.0L' -> 1.0)."""
    if "illusion" in scene_type or "target_cylinder" in scene_type:
        try:
            return float(scene_type.split("_")[-1].replace("L", ""))
        except (ValueError, IndexError):
            return None
    return None


def subtract_fields(q_a: dict, q_b: dict) -> Optional[dict]:
    """Compute q_a - q_b field difference.

    Both must have the same N. Returns dict with:
        ux, uy : (N, NY, NX) -- field difference
        forces : from q_a (reference)
        sensors : from q_a (reference)
        actions : from q_a (reference)
        meta : combined
        step_indices : from q_a

    Parameters
    ----------
    q_a : dict -- reference (minuend)
    q_b : dict -- subtrahend

    Returns
    -------
    dict or None if either input is None
    """
    if q_a is None or q_b is None:
        return None

    N_a = q_a["ux"].shape[0]
    N_b = q_b["ux"].shape[0]
    if N_a != N_b:
        raise ValueError(
            f"Frame count mismatch: q_a has {N_a} frames, q_b has {N_b}. "
            "Use common_length() to align."
        )

    return {
        "ux": q_a["ux"] - q_b["ux"],
        "uy": q_a["uy"] - q_b["uy"],
        "forces": q_a.get("forces"),
        "sensors": q_a.get("sensors"),
        "actions": q_a.get("actions"),
        "meta": {**q_a.get("meta", {}), "delta_from": q_b.get("meta", {}).get("scene", "unknown")},
        "step_indices": q_a.get("step_indices"),
    }


def dict_to_field_matrix(q: dict) -> np.ndarray:
    """Wrapper: build_field_matrix(q['ux'], q['uy']) with error checking.

    Parameters
    ----------
    q : dict -- must contain 'ux' and 'uy' with shape (N, NY, NX).

    Returns
    -------
    Q : (2 * NX * NY, N) ndarray -- snapshot matrix for POD.
    """
    if q is None:
        raise ValueError("Cannot build field matrix from None")
    ux = q["ux"]
    uy = q["uy"]
    if ux.ndim != 3 or ux.shape[-2:] != (NY, NX):
        raise ValueError(
            f"Expected field shape (N, {NY}, {NX}), got {ux.shape}"
        )
    return _build_field_matrix(ux, uy)


def compute_correction(scene_type: str) -> dict:
    """Load q_in, q_blk, q_ctl, q_tar and compute Delta fields.

    Parameters
    ----------
    scene_type : str -- one of the keys in _SCENE_MAP.

    Returns
    -------
    dict with:
        scene_type : str
        diam : float or None
        q_in, q_blk, q_ctl, q_tar : dict or None -- raw loaded data
        dq_blk, dq_ctl, dq_tar : dict or None -- field differences
        dq_tar_minus_blk : dict or None
        N : int -- aligned frame count (min across all loaded sources)
        meta : combined metadata dict
    """
    if scene_type not in _SCENE_MAP:
        raise KeyError(
            f"Unknown scene_type: {scene_type}. "
            f"Available: {list(_SCENE_MAP.keys())}"
        )

    q_in_name, q_blk_name, q_ctl_name, q_tar_name = _SCENE_MAP[scene_type]
    diam = get_diameter(scene_type)

    print(f"\n{'=' * 60}")
    print(f"Computing correction fields for: {scene_type}")
    print(f"  q_in = {q_in_name}, q_blk = {q_blk_name}, "
          f"q_ctl = {q_ctl_name}, q_tar = {q_tar_name}")
    print(f"{'=' * 60}")

    # -- Load all sources --
    q_in = _resolve_source(q_in_name) if q_in_name else None
    q_blk = _resolve_source(q_blk_name) if q_blk_name else None
    q_ctl = _resolve_source(q_ctl_name) if q_ctl_name else None
    q_tar = _resolve_source(q_tar_name) if q_tar_name else None

    # -- Determine aligned N --
    all_N = []
    for label, q in [("q_in", q_in), ("q_blk", q_blk), ("q_ctl", q_ctl), ("q_tar", q_tar)]:
        if q is not None:
            n = q["ux"].shape[0]
            all_N.append(n)
            print(f"  {label}: {n} frames, shape={q['ux'].shape}")
        else:
            print(f"  {label}: None")

    N = min(all_N) if all_N else 0

    # -- Compute Delta fields --
    # dq_blk  = q_ctl - q_blk  (controller adds beyond pinball)
    # dq_tar  = q_tar - q_blk  (target cylinder wake beyond pinball)
    # dq_ctl  = q_ctl - q_in   (ctl perturbation from inflow)
    # dq_tar_in = q_tar - q_in (target perturbation from inflow)

    dq_blk = _safe_subtract(q_blk, q_in, "dq_blk = q_blk - q_in (pinball blockage)")
    dq_ctl = _safe_subtract(q_ctl, q_blk, "dq_ctl = q_ctl - q_blk (control correction)")
    dq_tar = _safe_subtract(q_tar, q_blk, "dq_tar = q_tar - q_blk (target correction)")

    result: dict[str, Any] = {
        "scene_type": scene_type,
        "diam": diam,
        "q_in": q_in,
        "q_blk": q_blk,
        "q_ctl": q_ctl,
        "q_tar": q_tar,
        "dq_blk": dq_blk,            # q_blk - q_in
        "dq_ctl": dq_ctl,            # q_ctl - q_blk
        "dq_tar": dq_tar,            # q_tar - q_blk
        "N": N,
        "meta": {
            "scene_type": scene_type,
            "q_in": q_in_name,
            "q_blk": q_blk_name,
            "q_ctl": q_ctl_name,
            "q_tar": q_tar_name,
            "N_aligned": N,
        },
    }

    return result


def _safe_subtract(q_a: Optional[dict], q_b: Optional[dict],
                   label: str) -> Optional[dict]:
    """Subtract fields with optional trimming and None safety."""
    if q_a is None or q_b is None:
        print(f"  {label}: skipped (None input)")
        return None

    N_a = q_a["ux"].shape[0]
    N_b = q_b["ux"].shape[0]

    if N_a != N_b:
        N_min = min(N_a, N_b)
        print(f"  {label}: N mismatch ({N_a} vs {N_b}), "
              f"trimming to min N={N_min}")
        q_a_trim = _trim_to(q_a, N_min)
        q_b_trim = _trim_to(q_b, N_min)
    else:
        q_a_trim = q_a
        q_b_trim = q_b

    dq = subtract_fields(q_a_trim, q_b_trim)
    if dq is not None:
        _print_field_summary(f"  {label}", dq["ux"], dq["uy"])
    return dq


def _trim_to(q: dict, N: int) -> dict:
    """Trim first N frames from field dict."""
    return {
        "ux": q["ux"][:N],
        "uy": q["uy"][:N],
        "forces": q.get("forces")[:N] if q.get("forces") is not None else None,
        "sensors": q.get("sensors")[:N] if q.get("sensors") is not None else None,
        "actions": q.get("actions")[:N] if q.get("actions") is not None else None,
        "step_indices": q.get("step_indices")[:N] if q.get("step_indices") is not None else None,
        "meta": q.get("meta", {}),
    }


def _print_field_summary(label: str, ux: np.ndarray, uy: np.ndarray) -> None:
    """Print one-line field statistics."""
    ux_mean = ux.mean()
    uy_mean = uy.mean()
    ux_rms = ux.std()
    uy_rms = uy.std()
    mag_mean = np.sqrt(ux_mean**2 + uy_mean**2)
    print(f"  {label}:")
    print(f"    shape    = {ux.shape}")
    print(f"    ux_mean  = {ux_mean:.6f}   uy_mean  = {uy_mean:.6f}")
    print(f"    ux_rms   = {ux_rms:.6f}   uy_rms   = {uy_rms:.6f}")
    print(f"    |q|_mean = {mag_mean:.6f}")


# ---------------------------------------------------------------------------
# Main (test / verification)
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("=" * 60)
    print("Testing compute_correction_fields.py")
    print("=" * 60)

    # -- 1. Load target_channel + pinball, compute dq_blk --
    print("\n--- 1. Loading target_channel (legacy) + pinball (aligned) ---")
    target_channel = load_legacy_steady("target_channel")
    pinball = load_aligned_fields("pinball")

    _print_field_summary("target_channel (mean)", target_channel["ux"], target_channel["uy"])
    _print_field_summary("pinball (mean)", pinball["ux"], pinball["uy"])

    # Verify shapes
    print(f"\n  target_channel: N={target_channel['ux'].shape[0]}, "
          f"shape={target_channel['ux'].shape}")
    print(f"  pinball: N={pinball['ux'].shape[0]}, "
          f"shape={pinball['ux'].shape}")
    print(f"  sensors (target): {target_channel['sensors'].shape}")
    print(f"  sensors (pinball): {pinball['sensors'].shape}")
    print(f"  forces (pinball): {pinball['forces'].shape}")

    # -- 2. Compute dq_blk = pinball - target_channel --
    # Trim to smaller N
    dq_blk_test = _safe_subtract(pinball, target_channel, "pinball - target_channel")
    if dq_blk_test is not None:
        print(f"\n  dq_blk shape: {dq_blk_test['ux'].shape}")
        print(f"  dq_blk ux_mean (mean of difference): {dq_blk_test['ux'].mean():.6f}")

    # -- 3. Test full compute_correction for illusion_1.0L --
    print("\n--- 3. Full correction field pipeline: illusion_1.0L ---")
    result = compute_correction("illusion_1.0L")

    print(f"\n--- Result summary for {result['scene_type']} ---")
    print(f"  diam     = {result['diam']}")
    print(f"  N        = {result['N']}")
    print(f"  q_in     = {result['q_in']['ux'].shape if result['q_in'] else None}")
    print(f"  q_blk    = {result['q_blk']['ux'].shape if result['q_blk'] else None}")
    print(f"  q_ctl    = {result['q_ctl']['ux'].shape if result['q_ctl'] else None}")
    print(f"  q_tar    = {result['q_tar']['ux'].shape if result['q_tar'] else None}")

    for key in ["dq_blk", "dq_ctl", "dq_tar", "dq_tar_minus_blk"]:
        dq = result.get(key)
        if dq is not None:
            print(f"  {key}: mean(ux)={dq['ux'].mean():.6f}, "
                  f"mean(uy)={dq['uy'].mean():.6f}")
        else:
            print(f"  {key}: None")

    # -- 4. Verify dict_to_field_matrix --
    print("\n--- 4. Testing dict_to_field_matrix ---")
    Q = dict_to_field_matrix(pinball)
    print(f"  pinball snapshot matrix: {Q.shape} "
          f"(expect ({2 * NX * NY}, 96))")
    assert Q.shape == (2 * NX * NY, 96), f"Unexpected shape: {Q.shape}"
    print(f"  Q range: [{Q.min():.6f}, {Q.max():.6f}]")

    # -- 5. Test steady_cloak --
    print("\n--- 5. Testing steady_cloak correction ---")
    result_sc = compute_correction("steady_cloak")
    print(f"  steady_cloak N = {result_sc['N']}")
    if result_sc.get("dq_blk") is not None:
        print(f"  dq_blk (cloak-pinball) ux_mean: "
              f"{result_sc['dq_blk']['ux'].mean():.6f}")

    print("\nAll tests passed.")
