"""Simulate a trained policy in the MuJoCo viewer.

Uses Hydra to load the same BrittleStarConfig that was used during training.
Override settings via CLI, e.g.:
    python scripts/simulate.py morphology=3_arms
"""

from __future__ import annotations

import hydra
from omegaconf import DictConfig, OmegaConf
from pathlib import Path

from brittle_star_project import (
    BrittleStarEnv,
    BrittleStarEnvFactory,
    SimulationConfig,
    simulate_policy,
)
from brittle_star_project.configs.main_config import BrittleStarConfig
from brittle_star_project.configs.register_configs import register_configs
from brittle_star_project.rl import RLModel
from brittle_star_project.rl.base import get_rl_model_registry

MODEL_BY_NAME = get_rl_model_registry()
MODEL_OPTIONS = sorted(MODEL_BY_NAME)


@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3")
def main(dict_cfg: DictConfig) -> None:
    # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly.
    config: BrittleStarConfig = OmegaConf.to_object(
        OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)
    )

    # Use the configurable settings from the simulation group
    backend = config.simulation.backend
    model_type = config.simulation.model_type

    # Hydra chdir changes CWD; we map CLI relative paths relative to invocation originally.
    model_path = config.simulation.model_path
    if model_path is not None:
        model_path = hydra.utils.to_absolute_path(model_path)

    seed = config.experiment.seed

    # ======= ENVIRONMENT SETUP =======
    factory = BrittleStarEnvFactory()
    raw_env = factory.create_environment(
        backend, config.morphology, config.arena, config.environment
    )
    env = BrittleStarEnv(raw_env, backend=backend, config=config.environment)

    state = env.reset(seed=seed)

    # ======= MODEL SETUP =======
    nu = int(state.mj_model.nu)

    if model_path is not None:
        # TODO: Refactoring Notice - The .flax checkpoint payload no longer encapsulates the config
        # and no longer wraps parameters into a hardcoded list.
        # The file now natively contains solely the pure raw Jax 'agent_state.params' FrozenDict mapping.
        # The entire BrittleStarConfig is safely exported alongside it down at '..._metadata.yaml'.
        #
        # Example parsed layout from flax.serialization.from_bytes():
        # {
        #    'sensor_params': FrozenDict({...}),
        #    'actor_params': FrozenDict({...}),
        #    'critic_params': FrozenDict({...}),
        #    ...
        # }
        # Update the RLModel.load function or subsequent destructuring to support this raw dictionary natively.
        policy = RLModel.load(Path(model_path))
        if hasattr(policy, "nu"):
            policy.nu = nu
    else:
        model_cls = MODEL_BY_NAME[model_type]
        policy = model_cls(seed=seed)
        if hasattr(policy, "nu"):
            policy.nu = nu

    default_seed = int(getattr(policy, "seed", seed))

    # ======= SIMULATION =======
    rollout_cfg = SimulationConfig(
        realtime=True,
        seed=default_seed,
    )

    simulate_policy(policy, rollout_cfg, state)
    env.close()


if __name__ == "__main__":
    register_configs()
    main()