"""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) -----------------------
SR_ANALYSIS_DIR = os.path.abspath(os.path.dirname(__file__))
SRC_DIR = os.path.dirname(SR_ANALYSIS_DIR)
PROJECT_ROOT = os.path.dirname(SRC_DIR)
# Historical private alias retained for callers that imported it.
_PROJ = PROJECT_ROOT
MODEL_DIR = os.path.join(PROJECT_ROOT, "models")
LEGACY_CFG_DIR = os.path.join(PROJECT_ROOT, "configs", "legacy_configs")

# -- 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  # default; per-scene conv_len overrides this (Illusion=36)
CONTROL_TIME_SCALE_STEPS = 2000  # D_CYL / U0 in lattice steps

# Native legacy order follows object insertion in the actual environments.
SENSOR_LAYOUT = ("upper_ux", "upper_uy", "center_ux", "center_uy", "lower_ux", "lower_uy")
BODY_LAYOUT = ("front", "upper", "lower")
FORCE_LAYOUT = ("front_fx", "front_fy", "upper_fx", "upper_fy", "lower_fx", "lower_fy")
ACTION_LAYOUT = BODY_LAYOUT
RAW_LAYOUT = SENSOR_LAYOUT + FORCE_LAYOUT
TARGET_SENSOR_SLICE = (0, 6)
ILLUSION_TARGET_FORCE_SLICE = (0, 2)
ILLUSION_TARGET_SENSOR_SLICE = (2, 8)
DTW_VERSION = "legacy_dtw_v1_abs_n_unclipped"


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


def control_dt(sample_interval: float, control_time_scale_steps: float = CONTROL_TIME_SCALE_STEPS) -> float:
    """Return nondimensional control interval, SI / (D/U0)."""
    return float(sample_interval) / float(control_time_scale_steps)


# -- 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,
    }

# -- Karman Cloak generalization (unseen Re, no PPO model) --------------------
for rc in [25, 70, 150, 300]:
    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": "generalization",
        "model_name": None,
        "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,
    }

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,
    "control_time_scale_steps": CONTROL_TIME_SCALE_STEPS,
    "control_dt": control_dt(800),
    "fifo_len": FIFO_LEN,
    "conv_len": CONV_LEN,
    "sensor_layout": SENSOR_LAYOUT,
    "body_layout": BODY_LAYOUT,
    "action_layout": ACTION_LAYOUT,
    "force_layout": FORCE_LAYOUT,
    "raw_layout": RAW_LAYOUT,
    "raw_sensor_slice": TARGET_SENSOR_SLICE,
    "raw_force_slice": (6, 12),
    "target_sensor_slice": TARGET_SENSOR_SLICE,
    "target_force_slice": None,
    "fifo_init_action": (0.0, 0.0, 0.0),
    "policy_init_action": (0.0, 0.0, 0.0),
    "dtw_version": DTW_VERSION,
    "dtw_lag_channel": 1,
}

# -- Illusion (cylinder imitation, 3 diameters) ----------------------------
# "2U" in model name means S_DIM=14 (2 extra target force channels), NOT 2x velocity.
# u0 is always 0.01, nu is always 0.004 (default) unless Vis suffix present.
def _illusion_key(diam: float) -> str:
    """Generate clean illusion scene key."""
    s = f"{diam:.3f}".rstrip("0").rstrip(".")
    return f"illusion_{s}L"

_ILLUSION_BEST = [
    (0.75, "d1a3o14_250525_imit_075L_2U_400S", 400, 14),
    (1.0, "d1a3o14_250525_imit_1L_2U_600S", 600, 14),
]
for diam, mn, si, sd in _ILLUSION_BEST:
    key = _illusion_key(diam)
    SCENES[key] = {
        "scene_id": "illusion",
        "target_diameter": diam,
        "re_code": 100,
        "mu": 2.0 / 100,
        "nu": 0.004,  # default, no Vis suffix in model name
        "has_disturbance": False,
        "sample_interval": si,
        "conv_len": 36,  # Illusion uses 36, not default 30
        "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": sd,
        "u0": U0,  # always 0.01
    }

# 1.5L Illusion (SAMPLE_INTERVAL=800, default, no S suffix in name)
SCENES[_illusion_key(1.5)] = {
    "scene_id": "illusion",
    "target_diameter": 1.5,
    "re_code": 100,
    "mu": 2.0 / 100,
    "nu": 0.004,  # default, no Vis suffix
    "has_disturbance": False,
    "sample_interval": 800,  # no S suffix in name = default 800
    "conv_len": 36,  # Illusion uses 36
    "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": U0,  # always 0.01
}

# -- Illusion generalization (unseen diameters, no PPO model) ----------------
for diam, mn, si in [
    (0.5, None, 400),
    (0.6, None, 400),   # added 2026-06-25 for generalization test
    (0.8, None, 600),   # added 2026-06-25 for generalization test
    (1.2, None, 600),
    (2.0, None, 800),
]:
    key = _illusion_key(diam)
    SCENES[key] = {
        "scene_id": "illusion",
        "target_diameter": diam,
        "re_code": 100,
        "mu": 2.0 / 100,
        "nu": 0.004,
        "has_disturbance": False,
        "sample_interval": si,
        "conv_len": 36,
        "action_scale": 8.0,
        "action_bias": (0.0, -2.0, 2.0),
        "source": "generalization",
        "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": 14,
        "u0": U0,
    }


# -- Legacy contract completion ---------------------------------------------

def _complete_legacy_contract(cfg: dict) -> None:
    """Add explicit metadata without changing historical fields or values."""
    scene_id = cfg["scene_id"]
    if scene_id not in {"karman", "illusion"}:
        return

    cfg["re_d"] = cfg["re_code"] / 2.0
    cfg["control_time_scale_steps"] = CONTROL_TIME_SCALE_STEPS
    cfg["control_dt"] = control_dt(cfg["sample_interval"])
    cfg["fifo_len"] = FIFO_LEN
    cfg.setdefault("conv_len", CONV_LEN)
    cfg["sensor_layout"] = SENSOR_LAYOUT
    cfg["body_layout"] = BODY_LAYOUT
    cfg["action_layout"] = ACTION_LAYOUT
    cfg["force_layout"] = FORCE_LAYOUT
    cfg["raw_layout"] = RAW_LAYOUT
    cfg["raw_sensor_slice"] = TARGET_SENSOR_SLICE
    cfg["raw_force_slice"] = (6, 12)
    cfg["target_sensor_slice"] = (
        ILLUSION_TARGET_SENSOR_SLICE if scene_id == "illusion" else TARGET_SENSOR_SLICE
    )
    cfg["target_force_slice"] = (
        ILLUSION_TARGET_FORCE_SLICE if scene_id == "illusion" else None
    )
    cfg["fifo_init_action"] = (
        (0.0, -1.0, 1.0) if scene_id == "illusion" else tuple(cfg["action_bias"])
    )
    cfg["policy_init_action"] = cfg["fifo_init_action"]
    cfg["dtw_version"] = DTW_VERSION
    cfg["dtw_lag_channel"] = 1

    if scene_id == "illusion":
        # Legacy add_cylinder takes radius despite the historical diameter name.
        cfg["target_center_x"] = 20.0
        cfg["target_center"] = (20.0, CENTER_Y / L0, 0.0)
        cfg["target_radius"] = cfg["target_diameter"]
        cfg["target_radius_lattice"] = cfg["target_diameter"] * L0


for _scene_cfg in SCENES.values():
    _complete_legacy_contract(_scene_cfg)


# -- 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
