"""Unified scene configuration for OID_analysis.

Mirrors SR_analysis/configs.py and CCD_analysis/configs.py pattern.
All scene metadata in one place.

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, Tuple

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

# -- Physics constants -------------------------------------------------------
U0 = 0.01       # standard inlet center velocity (all models use this)
D_CYL = 20.0    # single cylinder diameter (lattice units)
D_REF = 40.0    # reference length = 2*D for code Re
L0 = 20.0       # base length unit
NX = 1280
NY = 512
CENTER_Y = (NY - 1) / 2.0
FIFO_LEN = 150
CONV_LEN_DEFAULT = 30   # Karman/Steady
CONV_LEN_ILLUSION = 36  # Illusion


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] = {}

# -- Empty Channel (reference) -----------------------------------------------
SCENES["empty_channel"] = {
    "scene_id": "empty_channel",
    "re_code": 100,
    "nu": 0.004,
    "has_disturbance": False,
    "sample_interval": 800,
    "source": "open_loop",
    "n_objects_env": 3,
    "obs_slice": (0, 6),
    "sensor_x": 40.0,
    "target_type": "steady",
    "u0": U0,
}

# -- Pure Pinball (uncontrolled baseline, steady/illusion positions) ---------
SCENES["pinball_baseline"] = {
    "scene_id": "pinball_baseline",
    "re_code": 100,
    "nu": 0.004,
    "has_disturbance": False,
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": (0.0, -4.0, 4.0),
    "source": "open_loop",
    "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,
}

# -- Pure Pinball (illusion positions) ---------------------------------------
SCENES["pinball_baseline_illusion"] = {
    "scene_id": "pinball_baseline_illusion",
    "re_code": 100,
    "nu": 0.004,
    "has_disturbance": False,
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": (0.0, -2.0, 2.0),
    "source": "open_loop",
    "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,
}

# -- Disturbance Only (Karman inflow) ----------------------------------------
SCENES["disturbance_only"] = {
    "scene_id": "disturbance_only",
    "re_code": 100,
    "nu": 0.004,
    "has_disturbance": True,
    "sample_interval": 800,
    "source": "open_loop",
    "n_objects_env": 4,
    "obs_slice": (2, 8),
    "sensor_x": 40.0,
    "disturbance_x": 10.0,
    "disturbance_radius": 1.0,  # L0 units
    "target_type": "periodic",
    "u0": U0,
}

# -- Disturbance + Pinball (Karman q_blk) ------------------------------------
SCENES["karman_blk"] = {
    "scene_id": "karman_blk",
    "re_code": 100,
    "nu": 0.004,
    "has_disturbance": True,
    "sample_interval": 800,
    "action_scale": 8.0,
    "action_bias": (0.0, -4.0, 4.0),
    "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": 12,
    "u0": U0,
}

# -- Steady Cloak (open-loop constant rotation) ------------------------------
SCENES["steady_cloak"] = {
    "scene_id": "steady_cloak",
    "re_code": 100,
    "nu": 0.004,
    "has_disturbance": False,
    "sample_interval": 800,
    "source": "open_loop",
    "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,
    "omega_front": 0.0,
    "omega_rear_scale": 5.1,
}

# -- Karman Cloak re100 (PPO) ------------------------------------------------
SCENES["karman_re100"] = {
    "scene_id": "karman",
    "re_code": 100,
    "nu": 0.004,
    "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,
}

# -- Illusion scenes (S_DIM=14, 3 diameters) ---------------------------------
_ILLUSION_SCENES = [
    ("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_SCENES:
    SCENES[key] = {
        "scene_id": "illusion",
        "target_diameter": diam,
        "re_code": 100,
        "nu": 0.004,
        "has_disturbance": False,
        "sample_interval": si,
        "conv_len": CONV_LEN_ILLUSION,
        "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": 30.0,
        "pinball_front_x": 19.0,
        "pinball_rear_x": 20.3,
        "target_type": "periodic",
        "s_dim": 14,
        "u0": U0,
    }

# -- Target cylinders (for illusion comparison) ------------------------------
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,
        "nu": 0.004,
        "has_disturbance": False,
        "sample_interval": si,
        "conv_len": CONV_LEN_ILLUSION,
        "source": "open_loop",
        "n_objects_env": 4,
        "obs_slice": (0, 8),
        "sensor_x": 30.0,
        "cylinder_x": 20.0,
        "target_type": "periodic",
        "u0": U0,
    }


# -- 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]:
    """Return full path to PPO model zip, or None if open-loop."""
    cfg = get_scene(scene_name)
    if cfg.get("source") != "PPO_inference":
        return None
    mn = cfg.get("model_name")
    if mn is None:
        return None
    subdir = cfg.get("model_subdir", "old")
    return os.path.join(MODEL_DIR, subdir, f"{mn}.zip")


def data_dir_for_scene(scene_name: str) -> str:
    """Return data directory for a given scene name.
    
    Each scene config gets its OWN directory based on scene_name
    to prevent filename collisions (different fields all save fields.npz).
    """
    cfg = get_scene(scene_name)
    sid = cfg["scene_id"]
    
    # Steady group: q_in=q_in, q_blk=q_blk, q_ctl=q_ctl
    if sid == "empty_channel":
        return os.path.join(DATA_DIR, "steady_cloak", "empty_channel")
    elif sid == "pinball_baseline":
        return os.path.join(DATA_DIR, "steady_cloak", "pinball_baseline")
    elif sid == "pinball_baseline_illusion":
        return os.path.join(DATA_DIR, "steady_cloak", "pinball_baseline_illusion")
    elif sid == "steady_cloak":
        return os.path.join(DATA_DIR, "steady_cloak", "steady_cloak")
    # Karman group: separate dirs for q_in, q_blk, q_ctl
    elif sid == "disturbance_only":
        return os.path.join(DATA_DIR, "karman_cloak", "disturbance_only")
    elif sid == "karman_blk":
        return os.path.join(DATA_DIR, "karman_cloak", "karman_blk")
    elif sid == "karman":
        return os.path.join(DATA_DIR, "karman_cloak", "karman_re100")
    # Target cylinders get their own dirs under target_cylinders
    elif sid == "target_cylinder":
        diam = cfg["target_diameter"]
        # Preserve one decimal: 0.75->"0.75", 1.0->"1.0", 1.5->"1.5"
        ds = f"{diam:.1f}" if diam == int(diam) else f"{diam:.2f}".rstrip("0")
        return os.path.join(DATA_DIR, "target_cylinder", f"target_cylinder_{ds}L")
    # Illusion scenes
    elif sid == "illusion":
        return os.path.join(DATA_DIR, "illusion", scene_name)
    else:
        return os.path.join(DATA_DIR, scene_name)
