"""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.

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
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 = 30


def nu_from_re(re_code: float, u0: float = U0) -> float:
    """Viscosity from code Reynolds number (reference length = 2*D)."""
    return u0 * D_REF / re_code


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

# -- Pure Pinball (uncontrolled baseline, 6 objects, no disturbance) ---------
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, 6 objects) --------------------
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 (PPO, 7 objects, disturbance cylinder) ---------------------
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),
}

# -- 1L Illusion (PPO, 6 objects, 2U=0.02) ----------------------------------
SCENES["illusion_1L"] = {
    "scene_id": "illusion",
    "re_code": 100,
    "target_diameter": 1.0,
    "has_disturbance": False,
    "sample_interval": 600,
    "action_scale": 8.0,
    "action_bias": (0.0, -2.0, 2.0),
    "source": "PPO_inference",
    "model_name": "d1a3o14_250525_imit_1L_2U_600S",
    "model_subdir": "250525",
    "n_objects_env": 6,
    "obs_slice": (0, 12),
    "sensor_x": 30.0,
    "pinball_front_x": 19.0,
    "pinball_rear_x": 20.3,
    "target_type": "periodic",
    "s_dim": 14,
    "u0": 0.02,
    "nu": nu_from_re(100, u0=0.02),
}


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

def get_scene(name: str) -> dict:
    """Return scene config dict by name. Raises KeyError if not found."""
    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]:
    """Return list of scene names, optionally filtered by scene_id."""
    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]:
    """Return absolute path to PPO model .zip file, or None."""
    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:
    """Return the data directory for a scene, creating it if needed."""
    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