"""All scene parameters in one place.

Single source of truth for geometry, action scaling, norm formulas, and DRL settings.
Every scene family needed for reproduction is defined here.

Re convention:
  - "re_code" uses reference length 2*D (=40 lattice units), matching model file naming.
  - nu = U0 * (2*D) / re_code
  - Physical Re_D = re_code / 2 (uses single cylinder diameter D=20)
"""
from __future__ import annotations

from typing import Any, Dict

import numpy as np

# ---------------------------------------------------------------------------
# Physics constants (must match config_lbm_pinball.json)
# ---------------------------------------------------------------------------
U0 = 0.01          # inlet centre velocity (lattice units)
L0 = 20.0          # base length unit = 1 cylinder diameter in lattice
D_CYL = 20.0       # single cylinder diameter
D_REF = 40.0       # reference length for code Re = 2*D
NX = 1280
NY = 512
CENTER_Y = (NY - 1) / 2.0   # 255.5
CFG_PATH = "configs/config_lbm_pinball.json"


def nu_from_re(re_code: float) -> float:
    """Kinematic viscosity from code Reynolds number."""
    return U0 * D_REF / re_code


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

# -- Steady Cloak (clean inflow, no disturbance cylinder) --------------------
# Also serves as the base pinball-only env for Illusion inference and Vortex.
SCENES["steady_cloak_re100"] = {
    "scene_id": "steady_cloak",
    "model": "d1a3o12_re100",
    "model_subdir": "old",
    "re_code": 100,
    "nu": nu_from_re(100),
    "s_dim": 12,
    "a_dim": 3,
    "has_disturbance": False,
    "pinball_front_x": 30.0 * L0,        # 600
    "pinball_rear_x": 31.3 * L0,         # 626
    "pinball_y_span": 0.75 * L0,         # 15
    "sensor_x": 40.0 * L0,               # 800
    "sensor_y_span": 2.0 * L0,           # 40
    "sensor_radius": L0 / 4,             # 5
    "pinball_radius": L0 / 2,            # 10
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": np.array([0.0, -4.0, 4.0], dtype=np.float32),
    "u0": U0,
    "fifo_len": 150,
    "conv_len": 30,
    "max_steps": 500,
    "n_objects": 6,
    "obs_slice": (0, 12),
    "force_norm_formula": "6*max",
    "sens_norm_factor": 5,
    "target_type": "steady",              # mean of clean channel
    "warmup_steps_init": int(4 * NX / U0),
    "warmup_steps_pinball": int(4 * NX / U0),
}

# -- Karman Cloak (upstream disturbance cylinder) ----------------------------
SCENES["karman_cloak_re100"] = {
    "scene_id": "karman_cloak",
    "model": "d1a3o12_re100",
    "model_subdir": "old",
    "re_code": 100,
    "nu": nu_from_re(100),
    "s_dim": 12,
    "a_dim": 3,
    "has_disturbance": True,
    "dist_center_x": 10.0 * L0,          # 200
    "dist_radius": 1.0 * L0,             # 20
    "pinball_front_x": 30.0 * L0,        # 600
    "pinball_rear_x": 31.3 * L0,         # 626
    "pinball_y_span": 0.75 * L0,         # 15
    "sensor_x": 40.0 * L0,               # 800
    "sensor_y_span": 2.0 * L0,           # 40
    "sensor_radius": L0 / 4,             # 5
    "pinball_radius": L0 / 2,            # 10
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": np.array([0.0, -4.0, 4.0], dtype=np.float32),
    "u0": U0,
    "fifo_len": 150,
    "conv_len": 30,
    "max_steps": 500,
    "n_objects": 7,
    "obs_slice": (2, 14),                # skip dist_cyl forces
    "force_norm_formula": "6*max",
    "sens_norm_factor": 5,
    "target_type": "periodic",
    "warmup_steps_dist": int(4 * NX / U0),
    "warmup_steps_pinball": int(4 * NX / U0),
}

# Also define the karman scenes at other Re for completeness
for re_code, name in [(50, "re50"), (200, "re200"), (400, "re400")]:
    key = f"karman_cloak_{name}"
    SCENES[key] = dict(SCENES["karman_cloak_re100"])
    SCENES[key].update({
        "model": f"d1a3o12_{name}",
        "model_subdir": "old",
        "re_code": re_code,
        "nu": nu_from_re(re_code),
        "scene_id": f"karman_cloak_{name}",
    })

# -- Illusion (three target diameters) ---------------------------------------
# NOTE: The positions below (sensors at 40*L0, pinball at 30/31.3*L0) are the
# "unified" inference geometry used by CCD_analysis. The actual training
# geometry (legacy_env_imit.py) used sensors at 30*L0 and pinball at 19/20.3*L0.
# phase3_reproduce.py and legacy_test scripts use the TRAINING positions.
def _illusion_base() -> Dict[str, Any]:
    return {
        "scene_id": "illusion",
        "re_code": 100,
        "nu": nu_from_re(100),
        "s_dim": 14,
        "a_dim": 3,
        "has_disturbance": False,
        "target_center_x": 31.0 * L0,         # 620 (inference position)
        "pinball_front_x": 30.0 * L0,         # 600 (standard inference pinball)
        "pinball_rear_x": 31.3 * L0,          # 626
        "pinball_y_span": 0.75 * L0,
        "sensor_x": 40.0 * L0,                # 800 (inference, not training!)
        "sensor_y_span": 2.0 * L0,
        "sensor_radius": L0 / 4,
        "pinball_radius": L0 / 2,
        "action_scale": 8.0,
        "action_bias": np.array([0.0, -2.0, 2.0], dtype=np.float32),
        "u0": U0,
        "fifo_len": 150,
        "conv_len": 36,                        # illusion uses CONV_LEN=36
        "max_steps": 500,
        "n_objects": 6,
        "obs_slice": (0, 12),
        "force_norm_formula": "6*max",
        "sens_norm_factor": 5,
        "target_type": "harmonics",
        "n_harmonics": 5,
        "warmup_steps_target": int(4 * NX / U0),
        "warmup_steps_pinball": int(4 * NX / U0),
    }

illusion_entries = [
    ("illusion_075L", {
        "model": "d1a3o14_250525_imit_075L_2U_400S",
        "model_subdir": "250525",
        "target_diameter": 0.75 * L0,         # 15
        "sample_interval": 400,
    }),
    ("illusion_1L", {
        "model": "d1a3o14_250525_imit_1L_2U_600S",
        "model_subdir": "250525",
        "target_diameter": 1.0 * L0,          # 20
        "sample_interval": 600,
    }),
    ("illusion_15L", {
        "model": "d1a3o14_250525_imit_15L_2U",
        "model_subdir": "250525",
        "target_diameter": 1.5 * L0,          # 30
        "sample_interval": 800,
    }),
]

for key, overrides in illusion_entries:
    base = _illusion_base()
    base.update(overrides)
    SCENES[key] = base

# -- Vortex (Lamb dipole and Taylor monopole) --------------------------------
def _vortex_base() -> Dict[str, Any]:
    return {
        "scene_id": "vortex",
        "re_code": 100,
        "nu": nu_from_re(100),
        "s_dim": 12,
        "a_dim": 3,
        "has_disturbance": False,
        "pinball_front_x": 30.0 * L0,
        "pinball_rear_x": 31.3 * L0,
        "pinball_y_span": 0.75 * L0,
        "sensor_x": 40.0 * L0,
        "sensor_y_span": 2.0 * L0,
        "sensor_radius": L0 / 4,
        "pinball_radius": L0 / 2,
        "sample_interval": 800,
        "action_scale": 4.0,                   # NOTE: scale=4 not 8
        "action_bias": np.array([0.0, -4.0, 4.0], dtype=np.float32),
        "u0": U0,
        "fifo_len": 150,
        "conv_len": 30,
        "max_steps": 150,                      # transient!
        "n_objects": 6,
        "obs_slice": (0, 12),
        "force_norm_formula": "6*max",
        "sens_norm_factor": 5,
        "target_type": "transient",
        "vortex_center_x": 10.0 * L0,          # target phase
        "vortex_pinball_center_x": 15.0 * L0,  # pinball phase
        "vortex_radius": 2.0 * L0,
    }

vortex_entries = [
    ("vortex_lamb", {
        "model": "vortex_lamb",
        "model_subdir": "old",
        "vortex_type": "lamb",
        "vortex_strength": 0.5 * U0,
    }),
    ("vortex_taylor", {
        "model": "vortex_taylor",
        "model_subdir": "old",
        "vortex_type": "taylor",
        "vortex_strength": 0.03 * U0,
    }),
]

for key, overrides in vortex_entries:
    base = _vortex_base()
    base.update(overrides)
    SCENES[key] = base


def get_scene(name: str) -> Dict[str, Any]:
    """Look up a scene by name. Raises KeyError if not found."""
    if name not in SCENES:
        available = ", ".join(sorted(SCENES.keys()))
        raise KeyError(f"Unknown scene '{name}'. Available: {available}")
    return dict(SCENES[name])  # return a copy