"""Train PPO agents on all NephroSim environments and generate results.

This script:
1. Trains PPO on each environment using configs from training/configs.py
2. Evaluates trained models against random and heuristic baselines
3. Saves training_results.json and trained models to results/
4. Generates publication-quality figures
"""

import json
import os
import sys
import time
from pathlib import Path

import gymnasium as gym
import numpy as np

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

try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

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


class RewardLoggerCallback(BaseCallback):
    """Log episode rewards during training for learning curve plots."""

    def __init__(self, verbose=0):
        super().__init__(verbose)
        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):
        if self.current_rewards is None:
            self.current_rewards = np.zeros(self.training_env.num_envs)
        rewards = self.locals.get("rewards", [])
        dones = self.locals.get("dones", [])
        for i, (r, d) in enumerate(zip(rewards, dones)):
            self.current_rewards[i] += r
            if d:
                self.episode_rewards.append(float(self.current_rewards[i]))
                self.current_rewards[i] = 0.0
        return True


def smooth(values, window=50):
    """Simple moving average for plotting."""
    if len(values) < window:
        return values
    smoothed = []
    for i in range(len(values)):
        start = max(0, i - window + 1)
        smoothed.append(np.mean(values[start:i + 1]))
    return smoothed


def generate_training_curves(all_rewards, output_path):
    """Generate training curve figure."""
    if not HAS_MATPLOTLIB:
        print("WARNING: matplotlib not installed, skipping figure generation")
        return

    plt.rcParams.update({
        "font.family": "serif",
        "font.size": 11,
        "axes.labelsize": 12,
        "axes.titlesize": 13,
        "figure.dpi": 300,
    })

    fig, axes = plt.subplots(2, 2, figsize=(12, 9))
    axes = axes.flatten()

    env_names = [
        ("nephrosim/UreaClearing-v0", "Urea Clearing"),
        ("nephrosim/UltrafiltrationControl-v0", "Ultrafiltration Control"),
        ("nephrosim/PhosphateMgmt-v0", "Phosphate Management"),
        ("nephrosim/FullDialysisSession-v0", "Full Dialysis Session"),
    ]

    for ax, (env_id, name) in zip(axes, env_names):
        if env_id in all_rewards and len(all_rewards[env_id]) > 0:
            rewards = all_rewards[env_id]
            smoothed = smooth(rewards, window=50)
            ax.plot(rewards, alpha=0.15, color="#3B82F6", linewidth=0.5)
            ax.plot(smoothed, color="#1E40AF", linewidth=2.0)
            ax.set_xlabel("Episode")
            ax.set_ylabel("Episode Reward")
            ax.set_title(name)
            ax.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close()
    print(f"  Saved training curves: {output_path}")


def generate_baseline_comparison(results, output_path):
    """Generate baseline comparison bar chart."""
    if not HAS_MATPLOTLIB:
        return

    plt.rcParams.update({
        "font.family": "serif",
        "font.size": 11,
        "axes.labelsize": 12,
        "axes.titlesize": 13,
        "figure.dpi": 300,
    })

    env_labels = []
    random_means = []
    heuristic_means = []
    ppo_means = []

    short_names = {
        "UreaClearing-v0": "Urea\nClearing",
        "UltrafiltrationControl-v0": "UF\nControl",
        "PhosphateMgmt-v0": "Phosphate\nMgmt",
        "FullDialysisSession-v0": "Full\nSession",
    }

    for env_name, data in results.items():
        env_labels.append(short_names.get(env_name, env_name))
        random_means.append(data["random"]["mean_reward"])
        heuristic_means.append(data["heuristic"]["mean_reward"])
        ppo_means.append(data["ppo"]["mean_reward"])

    x = np.arange(len(env_labels))
    width = 0.25

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(x - width, random_means, width, label="Random", color="#94A3B8")
    ax.bar(x, heuristic_means, width, label="Heuristic", color="#F59E0B")
    ax.bar(x + width, ppo_means, width, label="PPO", color="#3B82F6")

    ax.set_ylabel("Mean Episode Reward")
    ax.set_title("Agent Performance Comparison Across Environments")
    ax.set_xticks(x)
    ax.set_xticklabels(env_labels)
    ax.legend()
    ax.grid(True, axis="y", alpha=0.3)

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close()
    print(f"  Saved baseline comparison: {output_path}")


def generate_architecture_diagram(output_path):
    """Generate architecture overview diagram."""
    if not HAS_MATPLOTLIB:
        return

    plt.rcParams.update({
        "font.family": "serif",
        "font.size": 10,
        "figure.dpi": 300,
    })

    fig, ax = plt.subplots(figsize=(12, 7))
    ax.set_xlim(0, 12)
    ax.set_ylim(0, 7)
    ax.axis("off")

    # Title
    ax.text(6, 6.6, "NephroSim Architecture", ha="center", fontsize=16,
            fontweight="bold", fontfamily="serif")

    # Physiological Models layer
    models_box = plt.Rectangle((0.5, 4.0), 11, 2.0, fill=True,
                                facecolor="#DBEAFE", edgecolor="#1E40AF",
                                linewidth=2, zorder=2)
    ax.add_patch(models_box)
    ax.text(6, 5.6, "Physiological Models", ha="center", fontsize=13,
            fontweight="bold", color="#1E40AF")
    # Sub-models
    for i, (name, desc) in enumerate([
        ("Urea Kinetics", "Two-compartment\nGotch-Sargent"),
        ("Cardiovascular", "Baroreceptor reflex\nPlasma refilling"),
        ("Phosphate", "Two-compartment\nBinder effects"),
        ("Patient", "Profile + difficulty\ntiers"),
    ]):
        x_pos = 1.5 + i * 2.7
        box = plt.Rectangle((x_pos - 0.9, 4.2), 2.2, 1.2, fill=True,
                             facecolor="white", edgecolor="#3B82F6",
                             linewidth=1.5, zorder=3)
        ax.add_patch(box)
        ax.text(x_pos + 0.2, 5.0, name, ha="center", fontsize=9,
                fontweight="bold", zorder=4)
        ax.text(x_pos + 0.2, 4.55, desc, ha="center", fontsize=7,
                color="#64748B", zorder=4)

    # Gymnasium Environments layer
    envs_box = plt.Rectangle((0.5, 1.8), 11, 1.8, fill=True,
                              facecolor="#FEF3C7", edgecolor="#D97706",
                              linewidth=2, zorder=2)
    ax.add_patch(envs_box)
    ax.text(6, 3.3, "Gymnasium Environments", ha="center", fontsize=13,
            fontweight="bold", color="#92400E")
    for i, name in enumerate([
        "UreaClearing-v0", "UltrafiltrationControl-v0",
        "PhosphateMgmt-v0", "FullDialysisSession-v0",
    ]):
        x_pos = 1.5 + i * 2.7
        box = plt.Rectangle((x_pos - 0.9, 2.0), 2.2, 0.9, fill=True,
                             facecolor="white", edgecolor="#F59E0B",
                             linewidth=1.5, zorder=3)
        ax.add_patch(box)
        short = name.replace("-v0", "").replace("Ultrafiltration", "UF")
        ax.text(x_pos + 0.2, 2.45, short, ha="center", fontsize=8,
                fontweight="bold", zorder=4)

    # Agents layer
    agents_box = plt.Rectangle((0.5, 0.2), 11, 1.2, fill=True,
                                facecolor="#DCFCE7", edgecolor="#16A34A",
                                linewidth=2, zorder=2)
    ax.add_patch(agents_box)
    ax.text(6, 1.1, "Agents", ha="center", fontsize=13,
            fontweight="bold", color="#166534")
    for i, name in enumerate(["Random", "Heuristic\n(Clinical Protocol)", "PPO\n(Stable-Baselines3)"]):
        x_pos = 2.0 + i * 3.5
        box = plt.Rectangle((x_pos - 1.2, 0.3), 2.8, 0.7, fill=True,
                             facecolor="white", edgecolor="#22C55E",
                             linewidth=1.5, zorder=3)
        ax.add_patch(box)
        ax.text(x_pos + 0.2, 0.65, name, ha="center", fontsize=8,
                fontweight="bold", zorder=4)

    # Arrows
    for x_pos in [1.5 + i * 2.7 + 0.2 for i in range(4)]:
        ax.annotate("", xy=(x_pos, 3.0), xytext=(x_pos, 4.0),
                     arrowprops=dict(arrowstyle="->", color="#64748B",
                                     linewidth=1.5))

    plt.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close()
    print(f"  Saved architecture diagram: {output_path}")


def main():
    results_dir = Path("results")
    results_dir.mkdir(exist_ok=True)
    (results_dir / "models").mkdir(exist_ok=True)
    (results_dir / "figures").mkdir(exist_ok=True)

    all_rewards = {}
    training_results = {}
    overall_start = time.time()

    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"  Timesteps: {config['total_timesteps']:,}")
        print(f"{'='*60}")

        # Create wrapped environment
        env = ClipAction(NormalizeObservation(gym.make(env_id)))

        # Create callback
        callback = RewardLoggerCallback()

        # Train
        start = time.time()
        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,
        )
        model.learn(
            total_timesteps=config["total_timesteps"],
            callback=callback,
        )
        elapsed = time.time() - start
        env.close()

        # Save model
        model_path = str(results_dir / "models" / env_name)
        model.save(model_path)
        print(f"  Training time: {elapsed:.1f}s")
        print(f"  Model saved: {model_path}")

        # Store training rewards
        all_rewards[env_id] = callback.episode_rewards

        # Evaluate all agents
        print("  Evaluating agents...")
        random_result = evaluate_random(env_id, n_episodes=50, seed=42)
        heuristic_result = evaluate_heuristic(env_id, n_episodes=50, seed=42)
        ppo_result = evaluate_agent(model, env_id, n_episodes=50, seed=42)

        ratio = (
            ppo_result["mean_reward"] / random_result["mean_reward"]
            if random_result["mean_reward"] != 0
            else float("inf")
        )

        training_results[env_name] = {
            "env_id": env_id,
            "total_timesteps": config["total_timesteps"],
            "training_time_s": round(elapsed, 1),
            "training_episodes": len(callback.episode_rewards),
            "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_result["mean_reward"], 2),
                "std_reward": round(ppo_result["std_reward"], 2),
            },
            "ppo_vs_random_ratio": round(ratio, 2),
            "hyperparameters": {
                "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"],
            },
        }

        print(f"  Random:    {random_result['mean_reward']:.2f} +/- {random_result['std_reward']:.2f}")
        print(f"  Heuristic: {heuristic_result['mean_reward']:.2f} +/- {heuristic_result['std_reward']:.2f}")
        print(f"  PPO:       {ppo_result['mean_reward']:.2f} +/- {ppo_result['std_reward']:.2f}")
        print(f"  PPO/Random ratio: {ratio:.2f}x")

    # Save results JSON
    total_time = time.time() - overall_start
    output = {
        "package": "nephrosim",
        "version": "0.1.0",
        "total_training_time_s": round(total_time, 1),
        "environments": training_results,
    }

    results_path = results_dir / "training_results.json"
    with open(results_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved: {results_path}")

    # Generate figures
    print("\nGenerating figures...")
    generate_training_curves(all_rewards, str(results_dir / "figures" / "training_curves.png"))
    generate_baseline_comparison(training_results, str(results_dir / "figures" / "baseline_comparison.png"))
    generate_architecture_diagram(str(results_dir / "figures" / "architecture.png"))

    # Discovery JSON
    discovery = {
        "domain": "nephrology-hemodialysis",
        "package_name": "nephrosim",
        "environments": [
            {"id": env_id, "category": env_id.split("/")[-1]}
            for env_id in ENV_CONFIGS
        ],
        "models_implemented": [
            "two-compartment urea kinetics (Gotch-Sargent)",
            "cardiovascular response with baroreceptor reflex",
            "two-compartment phosphate kinetics",
        ],
        "clinical_targets": {
            "kt_v_minimum": 1.2,
            "phosphate_range_mg_dl": [3.5, 5.5],
            "hypotension_threshold_mmhg": 90,
        },
    }
    with open(results_dir / "discovery.json", "w") as f:
        json.dump(discovery, f, indent=2)
    print(f"Discovery saved: {results_dir / 'discovery.json'}")

    print(f"\nTotal time: {total_time:.1f}s")
    print("Done.")


if __name__ == "__main__":
    main()
