"""Unified scene configuration for CCD_analysis.

All scene metadata in one place. Each scene dict contains all parameters
needed for data collection, resampling, POD, and CCD.

CRITICAL: Illusion models use ONLY 2U series (d1a3o14_*).
1U series (d1a3o12_*) are NOT to be used.
All models use u0=0.01, nu=0.004.

Re convention:
  - "re_code" uses reference length 2*D (matching model file naming).
  - Re_D = re_code / 2 is the true physical Reynolds number.
"""
from __future__ import annotations

import os
from typing import Any, Dict, List, Optional

# -- Root paths ---------------------------------------------------------------
_PROJ = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
MODEL_DIR = os.path.join(_PROJ, "..", "models")
LEGACY_CFG_DIR = os.path.join(os.path.dirname(__file__), "configs")
DATA_DIR = os.path.join(os.path.dirname(__file__), "data")

# -- Physics constants -------------------------------------------------------
U0 = 0.01      # standard inlet center velocity (all models use this)
# NOTE: "2U" in model name means S_DIM=14 (2 extra target force obs), NOT u0 scaling
D_CYL = 20.0
D_REF = 40.0
L0 = 20.0
NX = 1280
NY = 512
CENTER_Y = (NY - 1) / 2.0
FIFO_LEN = 150
# CONV_LEN is per-scene. Illusion=36, Karman/Steady=30.
# Set locally in collection scripts, not as a global here.


def nu_from_re(re_code: float, u0: float = U0) -> float:
    return u0 * D_REF / re_code


# -- Scene definitions -------------------------------------------------------
SCENES: Dict[str, Any] = {}

# -- Pure Pinball (uncontrolled baseline) ------------------------------------
SCENES["pinball"] = {
    "scene_id": "pinball",
    "re_code": 100,
    "has_disturbance": False,
    "sample_interval": 800,
    "source": "open_loop",
    "model_name": None,
    "n_objects_env": 6,
    "obs_slice": (0, 12),
    "sensor_x": 40.0,
    "pinball_front_x": 30.0,
    "pinball_rear_x": 31.3,
    "target_type": "periodic",
    "s_dim": 12,
    "u0": U0,
    "nu": nu_from_re(100),
}

# -- Steady Cloak (open-loop constant rotation) ------------------------------
SCENES["steady_cloak"] = {
    "scene_id": "steady_cloak",
    "re_code": 100,
    "has_disturbance": False,
    "sample_interval": 800,
    "source": "open_loop",
    "model_name": None,
    "n_objects_env": 6,
    "obs_slice": (0, 12),
    "sensor_x": 40.0,
    "pinball_front_x": 30.0,
    "pinball_rear_x": 31.3,
    "target_type": "steady",
    "s_dim": 12,
    "u0": U0,
    "nu": nu_from_re(100),
    "omega_front": 0.0,
    "omega_rear_scale": 5.1,
}

# -- Karman Cloak re100 (PPO, cloak validation only) -------------------------
SCENES["karman_re100"] = {
    "scene_id": "karman",
    "re_code": 100,
    "has_disturbance": True,
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": (0.0, -4.0, 4.0),
    "source": "PPO_inference",
    "model_name": "d1a3o12_re100",
    "model_subdir": "old",
    "n_objects_env": 7,
    "obs_slice": (2, 14),
    "sensor_x": 40.0,
    "pinball_front_x": 30.0,
    "pinball_rear_x": 31.3,
    "target_type": "periodic",
    "s_dim": 12,
    "u0": U0,
    "nu": nu_from_re(100),
}

# -- Karman q_in (target/incoming vortex street, no pinball) ------------------
SCENES["karman_q_in"] = {
    "scene_id": "karman_target",
    "re_code": 100,
    "has_disturbance": True,
    "sample_interval": 800,
    "source": "open_loop",
    "n_objects_env": 4,
    "obs_slice": (0, 8),
    "sensor_x": 40.0,
    "pinball_front_x": None,
    "pinball_rear_x": None,
    "target_type": "periodic",
    "s_dim": None,
    "u0": U0,
    "nu": 0.004,
}

# -- Karman q_blk (pinball in vortex street, zero control) -------------------
SCENES["karman_q_blk"] = {
    "scene_id": "karman_blocked",
    "re_code": 100,
    "has_disturbance": True,
    "sample_interval": 800,
    "source": "open_loop",
    "n_objects_env": 7,
    "obs_slice": (2, 14),
    "sensor_x": 40.0,
    "pinball_front_x": 30.0,
    "pinball_rear_x": 31.3,
    "target_type": "periodic",
    "s_dim": None,
    "u0": U0,
    "nu": 0.004,
}

# -- Illusion scenes (S_DIM=14) -----------------------------------------------
# All use u0=0.01, SAMPLE_INTERVAL per diameter, nu=0.004 confirmed
# Sweep results: 0.004=0.962, 0.008=0.957, 0.002=0.882
# "2U" in model name = S_DIM=14 (2 extra target force dimensions), NOT 2x velocity
_ILLUSION_2U = [
    ("illusion_0.75L", "d1a3o14_250525_imit_075L_2U_400S", 0.75, 400),
    ("illusion_1.0L",  "d1a3o14_250525_imit_1L_2U_600S",  1.0,  600),
    ("illusion_1.5L",  "d1a3o14_250525_imit_15L_2U",      1.5,  800),
]
for key, mn, diam, si in _ILLUSION_2U:
    SCENES[key] = {
        "scene_id": "illusion",
        "target_diameter": diam,
        "re_code": 100,  # u0=0.01, nu=0.004
        "has_disturbance": False,
        "sample_interval": si,
        "conv_len": 36,  # Illusion uses 36 (see legacy_env_imit.py)
        "action_scale": 8.0,
        "action_bias": (0.0, -2.0, 2.0),
        "source": "PPO_inference",
        "model_name": mn,
        "model_subdir": "250525",
        "n_objects_env": 6,
        "obs_slice": (0, 12),
        "sensor_x": 40.0,                    # UNIFIED: was 30.0
        "pinball_front_x": 30.0,              # UNIFIED: was 19.0
        "pinball_rear_x": 31.3,               # UNIFIED: was 20.3
        "target_type": "periodic",
        "s_dim": 14,               # CRITICAL: all 2U are 14-dim
        "u0": U0,                  # 0.01 (NOT 0.02)
        "nu": 0.004,               # confirmed correct via sweep: 0.004=0.962, 0.008=0.957, 0.002=0.882
    }

# -- Target cylinders (per-diameter, for signature-CCD reference) ------------
# Each illusion diameter needs its own target cylinder data
# SAMPLE_INTERVAL per diameter, matching the corresponding illusion scene
for diam, si in [(0.75, 400), (1.0, 600), (1.5, 800)]:
    key = f"target_cylinder_{diam}L"
    SCENES[key] = {
        "scene_id": "target_cylinder",
        "target_diameter": diam,
        "re_code": 100,  # u0=0.01, nu=0.004
        "has_disturbance": False,
        "sample_interval": si,  # per-diameter, matching corresponding illusion
        "conv_len": 36,  # matching corresponding illusion
        "source": "open_loop",
        "model_name": None,
        "n_objects_env": 4,        # 1 cylinder + 3 sensors
        "obs_slice": (0, 8),       # cylinder force(2) + sensor(6)
        "sensor_x": 40.0,                    # UNIFIED: was 30.0
        "cylinder_x": 30.65,                  # UNIFIED: was 20.0 (pinball center)
        "target_type": "periodic",
        "s_dim": None,
        "u0": U0,                  # 0.01, matching illusion
        "nu": 0.004,               # confirmed correct via sweep: 0.004=0.962, 0.008=0.957, 0.002=0.882
    }

# -- Target Channel (empty channel, for steady metrics) ----------------------
SCENES["target_channel"] = {
    "scene_id": "target_channel",
    "re_code": 100,
    "has_disturbance": False,
    "sample_interval": 800,
    "source": "open_loop",
    "model_name": None,
    "n_objects_env": 3,
    "obs_slice": (0, 6),
    "sensor_x": 40.0,
    "pinball_front_x": None,
    "pinball_rear_x": None,
    "target_type": "steady",
    "s_dim": 6,
    "u0": U0,
    "nu": nu_from_re(100),
}

# -- Vortex cloak scenes (transient, Taylor monopole & Lamb dipole) ---------
# Taylor: monopole vortex, strength=0.03*U0, r=2*L0=40
# Lamb: dipole vortex, strength=0.5*U0, r=2*L0=40
# Both: MAX_STEPS=150, action_scale=4, action_bias=(0,-4,4), s_dim=12
# Geometry: pinball at (30, 31.3)xL0, sensors at 40*xL0
# Vortex at 10*xL0 (target) or 15*xL0 (pinball phase)
_VORTEX_SCENES = [
    ("vortex_lamb",          "vortex_lamb",     "lamb",  0.50),
    ("vortex_taylor",        "vortex_taylor",   "taylor", 0.03),
    ("vortex_uncontrolled_lamb",   None,         "lamb",  0.50),
    ("vortex_uncontrolled_taylor", None,         "taylor", 0.03),
    ("vortex_target_lamb",         None,         "lamb",  0.50),
    ("vortex_target_taylor",       None,         "taylor", 0.03),
]
for key, mn, vtype, vstrength in _VORTEX_SCENES:
    is_controlled = mn is not None
    SCENES[key] = {
        "scene_id": key,
        "re_code": 100,
        "has_disturbance": False,
        "sample_interval": 800,
        "vortex_type": vtype,
        "vortex_strength": vstrength,
        "conv_len": 30,
        "action_scale": 4.0 if is_controlled else None,
        "action_bias": (0.0, -4.0, 4.0) if is_controlled else None,
        "source": "PPO_inference" if is_controlled else "open_loop",
        "model_name": mn,
        "model_subdir": "old",
        "n_objects_env": 6,
        "obs_slice": (0, 12),
        "sensor_x": 40.0,
        "pinball_front_x": 30.0,
        "pinball_rear_x": 31.3,
        "target_type": "transient",
        "max_steps": 150,
        "s_dim": 12 if is_controlled else None,
        "u0": U0,
        "nu": nu_from_re(100),
    }


# -- Utility helpers ---------------------------------------------------------

def get_scene(name: str) -> dict:
    if name not in SCENES:
        raise KeyError(f"Unknown scene: {name}. Available: {list(SCENES.keys())}")
    return dict(SCENES[name])


def get_scene_list(scene_id: Optional[str] = None) -> List[str]:
    if scene_id is None:
        return list(SCENES.keys())
    return [k for k, v in SCENES.items() if v["scene_id"] == scene_id]


def model_path_for_scene(scene_name: str) -> Optional[str]:
    cfg = get_scene(scene_name)
    mn = cfg.get("model_name")
    if mn is None:
        return None
    subdir = cfg.get("model_subdir", "old")
    p = os.path.join(MODEL_DIR, subdir, f"{mn}.zip")
    return p if os.path.isfile(p) else None


def data_dir_for_scene(scene_name: str) -> str:
    cfg = get_scene(scene_name)
    scene_id = cfg["scene_id"]
    d = os.path.join(DATA_DIR, scene_id, scene_name)
    os.makedirs(d, exist_ok=True)
    return d
