"""Training pipeline for all OncoSim environments."""

from __future__ import annotations

import json
import os

import gymnasium as gym
import numpy as np

import oncosim  # noqa: F401
from oncosim.agents.random_agent import evaluate_random
from oncosim.agents.heuristic_agent import (
    BeamSelectionHeuristic,
    DoseFractionationHeuristic,
    AdaptiveRTHeuristic,
    evaluate_heuristic,
)


ENV_IDS = [
    "oncosim/BeamSelection-v0",
    "oncosim/DoseFractionation-v0",
    "oncosim/AdaptiveRT-v0",
]

HEURISTIC_MAP = {
    "oncosim/BeamSelection-v0": BeamSelectionHeuristic,
    "oncosim/DoseFractionation-v0": DoseFractionationHeuristic,
    "oncosim/AdaptiveRT-v0": AdaptiveRTHeuristic,
}


def train_all(
    timesteps_per_env: int = 90000,
    eval_episodes: int = 100,
    save_dir: str = "results",
) -> dict:
    """Train PPO on all environments, evaluate against baselines."""
    from oncosim.agents.ppo import train_ppo, evaluate_ppo

    os.makedirs(save_dir, exist_ok=True)
    os.makedirs("checkpoints", exist_ok=True)

    all_results = {}

    for env_id in ENV_IDS:
        print(f"\n{'='*60}")
        print(f"Training: {env_id}")
        print(f"{'='*60}")

        # Train PPO
        train_result = train_ppo(env_id, total_timesteps=timesteps_per_env)

        # Evaluate PPO
        ppo_eval = evaluate_ppo(env_id, train_result["model_path"], n_episodes=eval_episodes)

        # Evaluate random baseline
        env = gym.make(env_id)
        random_eval = evaluate_random(env, n_episodes=eval_episodes)

        # Evaluate heuristic baseline
        heuristic_cls = HEURISTIC_MAP[env_id]
        if env_id == "oncosim/BeamSelection-v0":
            heuristic = heuristic_cls(env)
        else:
            heuristic = heuristic_cls()
        heuristic_eval = evaluate_heuristic(env, heuristic, n_episodes=eval_episodes)
        env.close()

        # Compute ratios
        random_mean = random_eval["mean_reward"]
        ppo_mean = ppo_eval["mean_reward"]
        heuristic_mean = heuristic_eval["mean_reward"]

        # Handle negative rewards: use absolute improvement ratio
        if random_mean < 0:
            ratio = ppo_mean / random_mean if random_mean != 0 else float("inf")
            # For negative rewards, PPO should have LESS negative (higher) reward
            # ratio > 1 means PPO is worse (more negative), ratio < 1 means better
            improvement = abs(ppo_mean - random_mean) / (abs(random_mean) + 1e-10)
        else:
            ratio = ppo_mean / (random_mean + 1e-10)
            improvement = ratio

        all_results[env_id] = {
            "ppo": {
                "mean_reward": ppo_mean,
                "std_reward": ppo_eval["std_reward"],
                "episode_rewards": train_result["episode_rewards"],
            },
            "random": {
                "mean_reward": random_mean,
                "std_reward": random_eval["std_reward"],
            },
            "heuristic": {
                "mean_reward": heuristic_mean,
                "std_reward": heuristic_eval["std_reward"],
            },
            "ppo_random_ratio": ratio,
            "ppo_improvement_over_random": improvement,
            "total_timesteps": timesteps_per_env,
        }

        print(f"  PPO mean reward:       {ppo_mean:.3f}")
        print(f"  Random mean reward:    {random_mean:.3f}")
        print(f"  Heuristic mean reward: {heuristic_mean:.3f}")
        print(f"  PPO/Random ratio:      {ratio:.3f}")

    # Save results
    results_path = os.path.join(save_dir, "training_results.json")
    with open(results_path, "w") as f:
        serializable = {}
        for env_id, res in all_results.items():
            serializable[env_id] = {
                k: v for k, v in res.items()
                if k != "ppo" or not isinstance(v, dict)
            }
            serializable[env_id]["ppo"] = {
                "mean_reward": res["ppo"]["mean_reward"],
                "std_reward": res["ppo"]["std_reward"],
                "num_training_episodes": len(res["ppo"]["episode_rewards"]),
            }
            serializable[env_id]["random"] = res["random"]
            serializable[env_id]["heuristic"] = res["heuristic"]
            serializable[env_id]["ppo_random_ratio"] = res["ppo_random_ratio"]
            serializable[env_id]["ppo_improvement_over_random"] = res["ppo_improvement_over_random"]
            serializable[env_id]["total_timesteps"] = res["total_timesteps"]
        json.dump(serializable, f, indent=2)

    print(f"\nResults saved to {results_path}")
    return all_results
