"""Model inventory — metadata and loading utilities for all PPO models.

Provides:
- ModelInventory: registry of all pre-trained PPO models
- DummyEnv: minimal SB3-compatible env for PPO.load()
"""
from __future__ import annotations

import os
from typing import Any, Dict, Optional

import numpy as np
import gymnasium as gym
from gymnasium import spaces
from stable_baselines3 import PPO
from torch.nn import Module as TorchModule

_PROJECT_ROOT = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")
)
_MODELS_DIR = os.path.join(_PROJECT_ROOT, "models")


# ---------------------------------------------------------------------------
# DummyEnv for model loading
# ---------------------------------------------------------------------------
class DummyEnv(gym.Env):
    """Minimal SB3-compatible environment for loading PPO models.

    PPO.load() requires an env (or env's observation_space / action_space).
    This dummy env provides the correct spaces without any CFD backend.
    """

    def __init__(self, s_dim: int = 12, a_dim: int = 3):
        super().__init__()
        self.observation_space = spaces.Box(
            low=-1.0, high=1.0, shape=(s_dim,), dtype=np.float32,
        )
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, shape=(a_dim,), dtype=np.float32,
        )

    def reset(self, *, seed=None, options=None):
        return np.zeros(self.observation_space.shape, dtype=np.float32), {}

    def step(self, action):
        obs = np.zeros(self.observation_space.shape, dtype=np.float32)
        return obs, 0.0, False, False, {}


# ---------------------------------------------------------------------------
# Custom Sin activation (must match training)
# ---------------------------------------------------------------------------
class Sin(TorchModule):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        import torch
        return torch.sin(x)


# ---------------------------------------------------------------------------
# Model metadata
# ---------------------------------------------------------------------------
MODEL_META: Dict[str, Dict[str, Any]] = {
    # Old models: Karman cloak, cross-Re
    "d1a3o12_re50":       {"scene": "karman_cloak_re50",  "s_dim": 12, "subdir": "old"},
    "d1a3o12_re100":      {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "old"},
    "d1a3o12_re200":      {"scene": "karman_cloak_re200", "s_dim": 12, "subdir": "old"},
    "d1a3o12_re400":      {"scene": "karman_cloak_re400", "s_dim": 12, "subdir": "old"},
    # Vortex (transfer-learned from re100)
    "vortex_lamb":        {"scene": "vortex_lamb",        "s_dim": 12, "subdir": "old"},
    "vortex_taylor":      {"scene": "vortex_taylor",      "s_dim": 12, "subdir": "old"},
    # Re-trained cloak 250326
    "d1a3o12_250326":     {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "250326"},
    # No-offset cloak
    "d0a3o12_250329_nooffset": {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "250329"},
    # Reduced obs
    "d1a3o12_250421_forces02":            {"scene": "karman_cloak_re100", "s_dim": 3,  "subdir": "250421"},
    "d1a3o12_250421_torque+forces02":     {"scene": "karman_cloak_re100", "s_dim": 5,  "subdir": "250421"},
    "d1a3o12_250421_torque+forces02+sens24": {"scene": "karman_cloak_re100", "s_dim": 9,  "subdir": "250421"},
    "d1a3o12_250421_torque+forces04+sens04": {"scene": "karman_cloak_re100", "s_dim": 5,  "subdir": "250421"},
    "d1a3o12_250421_torque+total_force":  {"scene": "karman_cloak_re100", "s_dim": 5,  "subdir": "250421"},
    "d1a3o12_250421_total_force":         {"scene": "karman_cloak_re100", "s_dim": 3,  "subdir": "250421"},
    # Illusion (250525)
    "d1a3o14_250525_imit_075L_2U_400S":   {"scene": "illusion_075L", "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_1L_2U_600S":     {"scene": "illusion_1L",   "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_15L_2U":         {"scene": "illusion_15L",  "s_dim": 14, "subdir": "250525"},
    # Additional illusion variants
    "d1a3o14_250525_imit_075L_2U":        {"scene": "illusion_075L", "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_1L_2U":          {"scene": "illusion_1L",   "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_1L_2U_1":        {"scene": "illusion_1L",   "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_1L_2U_1000S_08Vis": {"scene": "illusion_1L", "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_1L_2U_800S_08Vis":  {"scene": "illusion_1L", "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_1L_2U_400S_02Vis":  {"scene": "illusion_1L", "s_dim": 14, "subdir": "250525"},
    "d1a3o14_250525_imit_075L_2U_1":      {"scene": "illusion_075L", "s_dim": 14, "subdir": "250525"},
    # Early illusion models (S_DIM=12, 1U variants)
    "d1a3o12_250525_imit_075L_1U":        {"scene": "illusion_075L", "s_dim": 12, "subdir": "250525"},
    "d1a3o12_250525_imit_1L_1U":          {"scene": "illusion_1L",   "s_dim": 12, "subdir": "250525"},
    "d1a3o12_250525_imit_1L_1U_trans":    {"scene": "illusion_1L",   "s_dim": 12, "subdir": "250525"},
    # Erase models (for reference, not primary focus)
    "d1a3o12_250729_250326_erase":        {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "250729"},
    "d1a3o12_250729_250326_erase_250804_20D_retrain2": {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "250729"},
    "d1a3o12_250729_250326_erase_250804_20D_retrain3": {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "250729"},
    "d1a3o12_250729_250326_cloak_800S_02Vis": {"scene": "karman_cloak_re100", "s_dim": 12, "subdir": "250729"},
}


class ModelInventory:
    """Registry for loading pre-trained PPO models."""

    def __init__(self):
        self.models_dir = _MODELS_DIR

    def get_model_path(self, name: str) -> str:
        if name not in MODEL_META:
            raise KeyError(f"Unknown model '{name}'")
        meta = MODEL_META[name]
        return os.path.join(self.models_dir, meta["subdir"], f"{name}.zip")

    def load(self, name: str, device: str = "cuda:0") -> PPO:
        """Load a PPO model with correct observation space and Sin activation."""
        if name not in MODEL_META:
            raise KeyError(f"Unknown model '{name}'")
        meta = MODEL_META[name]
        dummy = DummyEnv(s_dim=meta["s_dim"])
        path = self.get_model_path(name)
        model = PPO.load(path, env=dummy, device=device)
        return model

    def list_models(self, scene: Optional[str] = None) -> list:
        """List model names, optionally filtered by scene name."""
        if scene is None:
            return sorted(MODEL_META.keys())
        return sorted(k for k, v in MODEL_META.items() if v["scene"] == scene)
