"""Train all agents and generate results for AnestheSim.

Trains PPO agents on all environments, evaluates against baselines,
generates training_results.json and publication figures.
"""

import json
import os
import random
import sys
import time

import gymnasium as gym
import numpy as np

import anesthesim  # noqa: F401
from anesthesim.training.configs import ENV_CONFIGS
from anesthesim.agents.random_agent import evaluate_random
from anesthesim.agents.heuristic import evaluate_heuristic
from anesthesim.envs.wrappers import NormalizeObservation, ClipAction

SEED = 42
EVAL_EPISODES = 50
RESULTS_DIR = "results"
FIGURES_DIR = "paper/figures"


def train_and_evaluate():
    """Train PPO on all envs and evaluate against baselines."""
    random.seed(SEED)
    np.random.seed(SEED)
    print(f"Training with seed={SEED}")

    os.makedirs(RESULTS_DIR, exist_ok=True)
    os.makedirs(f"{RESULTS_DIR}/models", exist_ok=True)
    os.makedirs(f"{RESULTS_DIR}/training_logs", exist_ok=True)
    os.makedirs(FIGURES_DIR, exist_ok=True)

    try:
        from stable_baselines3 import PPO
        from stable_baselines3.common.callbacks import BaseCallback
    except ImportError:
        print("ERROR: stable-baselines3 required. Install: pip install anesthesim[train]")
        sys.exit(1)

    class RewardLoggerCallback(BaseCallback):
        """Log episode rewards during training."""
        def __init__(self):
            super().__init__()
            self.episode_rewards = []
            self.episode_lengths = []
            self.current_rewards = None

        def _on_training_start(self):
            self.current_rewards = np.zeros(self.training_env.num_envs)

        def _on_step(self):
            rewards = self.locals.get("rewards", [])
            dones = self.locals.get("dones", [])
            if self.current_rewards is not None:
                self.current_rewards += rewards
                for i, done in enumerate(dones):
                    if done:
                        self.episode_rewards.append(float(self.current_rewards[i]))
                        self.current_rewards[i] = 0.0
            return True

    results = {}

    for env_id, config in ENV_CONFIGS.items():
        env_name = env_id.split("/")[-1]
        print(f"\n{'='*60}")
        print(f"Training {env_name}")
        print(f"{'='*60}")

        # Evaluate baselines first
        print("Evaluating random baseline...")
        random_result = evaluate_random(env_id, n_episodes=EVAL_EPISODES, seed=SEED)
        print(f"  Random: {random_result['mean_reward']:.2f} +/- {random_result['std_reward']:.2f}")

        print("Evaluating heuristic baseline...")
        heuristic_result = evaluate_heuristic(env_id, n_episodes=EVAL_EPISODES, seed=SEED)
        print(f"  Heuristic: {heuristic_result['mean_reward']:.2f} +/- {heuristic_result['std_reward']:.2f}")

        # Train PPO
        print(f"Training PPO for {config['total_timesteps']:,} steps...")
        env = ClipAction(NormalizeObservation(gym.make(env_id)))

        callback = RewardLoggerCallback()
        model = PPO(
            "MlpPolicy",
            env,
            learning_rate=config["learning_rate"],
            n_steps=config["n_steps"],
            batch_size=config["batch_size"],
            gamma=config["gamma"],
            gae_lambda=config["gae_lambda"],
            clip_range=config["clip_range"],
            ent_coef=config["ent_coef"],
            seed=config["seed"],
            verbose=0,
        )

        start_time = time.time()
        model.learn(total_timesteps=config["total_timesteps"], callback=callback)
        train_time = time.time() - start_time
        print(f"  Training took {train_time:.1f}s")

        # Save model
        model_path = f"{RESULTS_DIR}/models/{env_name}_ppo"
        model.save(model_path)
        env.close()

        # Evaluate PPO (use same wrapper chain as training)
        print("Evaluating PPO agent...")
        ppo_rewards = []
        rng = np.random.default_rng(SEED)
        for ep in range(EVAL_EPISODES):
            eval_env = ClipAction(NormalizeObservation(gym.make(env_id)))
            obs, _ = eval_env.reset(seed=int(rng.integers(0, 2**31)))
            total_reward = 0.0
            done = False
            while not done:
                action, _ = model.predict(obs, deterministic=True)
                obs, reward, terminated, truncated, info = eval_env.step(action)
                total_reward += float(reward)
                done = terminated or truncated
            ppo_rewards.append(total_reward)
            eval_env.close()

        ppo_mean = float(np.mean(ppo_rewards))
        ppo_std = float(np.std(ppo_rewards))
        print(f"  PPO: {ppo_mean:.2f} +/- {ppo_std:.2f}")

        # Compute ratio
        random_mean = random_result["mean_reward"]
        if abs(random_mean) < 0.01:
            ppo_vs_random_ratio = ppo_mean - random_mean
            ratio_note = "Raw difference used (random baseline near zero)"
        elif random_mean < 0:
            ppo_vs_random_ratio = abs(ppo_mean - random_mean) / abs(random_mean)
            ratio_note = "Absolute difference ratio (negative baseline)"
        else:
            ppo_vs_random_ratio = ppo_mean / random_mean
            ratio_note = "Direct ratio"

        # Training reward curve (windowed averages for plotting)
        training_rewards = callback.episode_rewards
        window = max(1, len(training_rewards) // 50)
        if len(training_rewards) > 0:
            smoothed = []
            for i in range(0, len(training_rewards), window):
                chunk = training_rewards[i:i+window]
                smoothed.append(float(np.mean(chunk)))
            reward_curve = smoothed
        else:
            reward_curve = []

        # Compute sample efficiency: steps to reach 50% of peak
        if len(reward_curve) > 0:
            peak = max(reward_curve)
            threshold = 0.5 * peak if peak > 0 else peak * 1.5
            steps_to_50pct = len(reward_curve) * window
            for i, r in enumerate(reward_curve):
                if r >= threshold:
                    steps_to_50pct = i * window
                    break
        else:
            steps_to_50pct = config["total_timesteps"]

        # Stability: std of final 20% of training
        if len(reward_curve) > 5:
            final_20pct = reward_curve[int(len(reward_curve) * 0.8):]
            stability_std = float(np.std(final_20pct))
        else:
            stability_std = 0.0

        results[env_name] = {
            "env_id": env_id,
            "total_timesteps": config["total_timesteps"],
            "training_time_seconds": round(train_time, 1),
            "random": {
                "mean_reward": round(random_result["mean_reward"], 2),
                "std_reward": round(random_result["std_reward"], 2),
            },
            "heuristic": {
                "mean_reward": round(heuristic_result["mean_reward"], 2),
                "std_reward": round(heuristic_result["std_reward"], 2),
            },
            "ppo": {
                "mean_reward": round(ppo_mean, 2),
                "std_reward": round(ppo_std, 2),
            },
            "ppo_vs_random_ratio": round(ppo_vs_random_ratio, 2),
            "ppo_vs_random_ratio_note": ratio_note,
            "sample_efficiency_steps_to_50pct": steps_to_50pct,
            "stability_std_final_20pct": round(stability_std, 2),
            "reward_curve": reward_curve,
            "seed": config["seed"],
        }

        # Save training log
        log_path = f"{RESULTS_DIR}/training_logs/{env_name}.log"
        with open(log_path, "w") as f:
            f.write(f"env_id={env_id}\n")
            f.write(f"seed={config['seed']}\n")
            f.write(f"total_timesteps={config['total_timesteps']}\n")
            f.write(f"training_time={train_time:.1f}s\n")
            f.write(f"ppo_mean_reward={ppo_mean:.4f}\n")
            f.write(f"random_mean_reward={random_mean:.4f}\n")
            f.write(f"heuristic_mean_reward={heuristic_result['mean_reward']:.4f}\n")
            f.write(f"ppo_vs_random_ratio={ppo_vs_random_ratio:.4f}\n")
            f.write(f"n_training_episodes={len(training_rewards)}\n")

    # Save results
    results_path = f"{RESULTS_DIR}/training_results.json"
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {results_path}")

    return results


def generate_figures(results: dict):
    """Generate publication-quality figures from training results."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams["font.family"] = "serif"
    plt.rcParams["font.size"] = 11
    plt.rcParams["axes.labelsize"] = 12
    plt.rcParams["axes.titlesize"] = 13

    os.makedirs(FIGURES_DIR, exist_ok=True)

    # Figure 1: Training curves with variance bands
    fig, axes = plt.subplots(1, 3, figsize=(14, 4))
    colors = ["#2196F3", "#4CAF50", "#FF5722"]

    for idx, (env_name, data) in enumerate(results.items()):
        ax = axes[idx]
        curve = data["reward_curve"]
        if len(curve) == 0:
            continue

        x = np.arange(len(curve))
        y = np.array(curve)

        # Compute rolling std for variance bands
        window = max(1, len(y) // 10)
        y_smooth = np.convolve(y, np.ones(window)/window, mode="valid")
        x_smooth = x[:len(y_smooth)]

        # Variance band: use rolling std
        y_std = []
        for i in range(len(y_smooth)):
            start = max(0, i - window)
            end = min(len(y), i + window)
            y_std.append(np.std(y[start:end]))
        y_std = np.array(y_std)

        ax.plot(x_smooth, y_smooth, color=colors[idx], linewidth=1.5)
        ax.fill_between(x_smooth, y_smooth - y_std, y_smooth + y_std,
                        color=colors[idx], alpha=0.2)
        ax.set_title(env_name.replace("-v0", ""))
        ax.set_xlabel("Training Episode Window")
        ax.set_ylabel("Mean Episode Reward")
        ax.grid(True, alpha=0.3)

        # Add random baseline line
        random_mean = data["random"]["mean_reward"]
        ax.axhline(y=random_mean, color="gray", linestyle="--", alpha=0.5, label="Random")
        ax.legend(fontsize=8)

    plt.tight_layout()
    plt.savefig(f"{FIGURES_DIR}/training_curves.png", dpi=300, bbox_inches="tight")
    plt.close()
    print(f"Saved {FIGURES_DIR}/training_curves.png")

    # Figure 2: Baseline comparison bar chart
    fig, ax = plt.subplots(figsize=(8, 5))
    env_names = list(results.keys())
    x = np.arange(len(env_names))
    width = 0.25

    random_means = [results[n]["random"]["mean_reward"] for n in env_names]
    heuristic_means = [results[n]["heuristic"]["mean_reward"] for n in env_names]
    ppo_means = [results[n]["ppo"]["mean_reward"] for n in env_names]
    random_stds = [results[n]["random"]["std_reward"] for n in env_names]
    heuristic_stds = [results[n]["heuristic"]["std_reward"] for n in env_names]
    ppo_stds = [results[n]["ppo"]["std_reward"] for n in env_names]

    bars1 = ax.bar(x - width, random_means, width, yerr=random_stds,
                   label="Random", color="#90A4AE", capsize=3)
    bars2 = ax.bar(x, heuristic_means, width, yerr=heuristic_stds,
                   label="Heuristic (TCI)", color="#FFB74D", capsize=3)
    bars3 = ax.bar(x + width, ppo_means, width, yerr=ppo_stds,
                   label="PPO", color="#42A5F5", capsize=3)

    ax.set_xlabel("Environment")
    ax.set_ylabel("Mean Episode Reward")
    ax.set_title("Baseline Comparison Across Environments")
    ax.set_xticks(x)
    short_names = [n.replace("Control", "Ctrl").replace("Dosing", "Dose")
                   .replace("Anesthesia", "Anes.") for n in env_names]
    ax.set_xticklabels(short_names)
    ax.legend()
    ax.grid(True, alpha=0.3, axis="y")

    plt.tight_layout()
    plt.savefig(f"{FIGURES_DIR}/baseline_comparison.png", dpi=300, bbox_inches="tight")
    plt.close()
    print(f"Saved {FIGURES_DIR}/baseline_comparison.png")

    # Figure 3: Architecture diagram (text-based for headless)
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 7)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.set_title("AnestheSim System Architecture", fontsize=14, fontweight="bold")

    # Draw boxes
    boxes = [
        (1, 5.5, "Marsh PK\n(Propofol)", "#E3F2FD"),
        (4, 5.5, "Minto PK\n(Remifentanil)", "#E8F5E9"),
        (7, 5.5, "Hill PD\n(BIS)", "#FFF3E0"),
        (1, 3.5, "PropofolControl\nEnv", "#BBDEFB"),
        (4, 3.5, "RemifentanilDosing\nEnv", "#C8E6C9"),
        (7, 3.5, "CombinedAnesthesia\nEnv", "#FFE0B2"),
        (4, 1.5, "PPO / Heuristic\nAgents", "#F3E5F5"),
    ]

    for x_pos, y_pos, text, color in boxes:
        rect = plt.Rectangle((x_pos - 0.9, y_pos - 0.5), 1.8, 1.0,
                            facecolor=color, edgecolor="black", linewidth=1)
        ax.add_patch(rect)
        ax.text(x_pos, y_pos, text, ha="center", va="center", fontsize=8)

    # Draw arrows
    arrows = [
        (1, 5.0, 1, 4.1), (4, 5.0, 4, 4.1), (7, 5.0, 7, 4.1),
        (1, 3.0, 4, 2.1), (4, 3.0, 4, 2.1), (7, 3.0, 4, 2.1),
        (7, 5.0, 4, 4.1),  # Greco interaction
    ]
    for x1, y1, x2, y2 in arrows:
        ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
                    arrowprops=dict(arrowstyle="->", color="black", lw=1))

    # Label
    ax.text(5, 0.5, "Greco Interaction Model connects propofol + remifentanil PD",
            ha="center", fontsize=8, style="italic")

    plt.savefig(f"{FIGURES_DIR}/architecture.png", dpi=300, bbox_inches="tight")
    plt.close()
    print(f"Saved {FIGURES_DIR}/architecture.png")


if __name__ == "__main__":
    results = train_and_evaluate()
    generate_figures(results)
    print("\nAll training and figure generation complete.")
