"""Unified scene configuration for SR_analysis.

All scene metadata in one place. Each scene dict contains all parameters
needed for data generation, SINDy fitting, and validation.

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

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

# -- Root paths (resolved when configs.py is imported) -----------------------
_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", "legacy")

# -- Physics constants -------------------------------------------------------
U0 = 0.01   # default inlet velocity (lattice)
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 -------------------------------------------------------

# Each scene dict has fields:
#   scene_id, re_code, mu, nu, has_disturbance, sample_interval,
#   action_scale, action_bias (tuple), source ("PPO_inference"|"open_loop"),
#   model_name (str or None), n_objects_env, obs_slice [start,end],
#   sensor_x, pinball_front_x, pinball_rear_x,
#   target_type ("periodic"|"steady"|"transient"),
#   s_dim (DRL observation dim, 12 or 14)

SCENES: Dict[str, Any] = {}

# -- Karman Cloak (cross-Re) ------------------------------------------------
for rc, mn in [(50, "d1a3o12_re50"), (100, "d1a3o12_re100"),
               (200, "d1a3o12_re200"), (400, "d1a3o12_re400")]:
    key = f"karman_re{rc}"
    SCENES[key] = {
        "scene_id": "karman",
        "re_code": rc,
        "mu": 2.0 / rc,
        "nu": nu_from_re(rc),
        "has_disturbance": True,
        "sample_interval": 800,
        "action_scale": 8.0,
        "action_bias": (0.0, -4.0, 4.0),
        "source": "PPO_inference",
        "model_name": mn,
        "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,
    }

# -- Steady Cloak (open-loop constant rotation) -----------------------------
SCENES["steady"] = {
    "scene_id": "steady",
    "re_code": 100,
    "mu": 2.0 / 100,
    "nu": nu_from_re(100),
    "has_disturbance": False,
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": (0.0, -5.1, 5.1),   # from gen_steady_data.py defaults
    "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,
}

# -- Illusion (cylinder imitation, 3 diameters, 1U=0.01) --------------------
def _illusion_key(diam: float) -> str:
    """Generate clean illusion scene key."""
    s = f"{diam:.3f}".rstrip("0").rstrip(".")
    return f"illusion_{s}L"

_ILLUSION_1U = [
    (0.75, "d1a3o12_250525_imit_075L_1U"),
    (1.0, "d1a3o12_250525_imit_1L_1U"),
]
for diam, mn in _ILLUSION_1U:
    key = _illusion_key(diam)
    SCENES[key] = {
        "scene_id": "illusion",
        "target_diameter": diam,
        "re_code": 100,
        "mu": 2.0 / 100,
        "nu": nu_from_re(100),
        "has_disturbance": False,
        "sample_interval": 600,
        "action_scale": 8.0,
        "action_bias": (0.0, -2.0, 2.0),
        "source": "PPO_inference",
        "model_name": mn,
        "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": 12,
        "u0": U0,
    }

# 1.5L Illusion (2U=0.02 model)
SCENES[_illusion_key(1.5)] = {
    "scene_id": "illusion",
    "target_diameter": 1.5,
    "re_code": 100,
    "mu": 2.0 / 100,
    "nu": nu_from_re(100, u0=0.02),
    "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_15L_2U",
    "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,
}

# -- Vortex Cloak (Lamb dipole + Taylor monopole) --------------------------
_SCENES_VORTEX = [
    ("lamb", "vortex_lamb", 0.5),
    ("taylor", "vortex_taylor", 0.03),
]
for vtype, mn, strength in _SCENES_VORTEX:
    key = f"vortex_{vtype}"
    SCENES[key] = {
        "scene_id": "vortex",
        "vortex_type": vtype,
        "vortex_strength": strength,
        "re_code": 100,
        "mu": 2.0 / 100,
        "nu": nu_from_re(100),
        "has_disturbance": False,
        "sample_interval": 800,
        "max_steps": 150,
        "action_scale": 4.0,
        "action_bias": (0.0, -4.0, 4.0),
        "source": "PPO_inference",
        "model_name": mn,
        "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",
        "s_dim": 12,
        "u0": U0,
    }


# -- 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
    # Check model directories in priority order
    candidate_dirs = [
        os.path.join(MODEL_DIR, "old"),
        os.path.join(MODEL_DIR, "250525"),
        os.path.join(MODEL_DIR, "250729"),
        os.path.join(MODEL_DIR, "250326"),
    ]
    for d in candidate_dirs:
        p = os.path.join(d, f"{mn}.zip")
        if os.path.isfile(p):
            return p
    return None
