#!/usr/bin/env python3
"""
GRU-based curriculum training script. Trains a masked DQN with a GRU-augmented
feature extractor; GRU remains statically enabled throughout training (curriculum
controls only environment difficulty, not recurrence).
"""
# ====================================================
import sys
import os
import argparse
from pathlib import Path
import numpy as np
import torch
import gymnasium as gym
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
from gru_belief_policy_v3_agnostic import GRUAugmentedFeaturesExtractor
# Custom objects for proper serialization
CUSTOM_OBJECTS = {
    "features_extractor_class": GRUAugmentedFeaturesExtractor,
}
# -----------------------------------------------------------------------------
# Project imports
# -----------------------------------------------------------------------------
sys.path.insert(0, str(Path(__file__).parent))
from masked_dqn_policy import MaskedDQN
from mems_adaptive_inspection_env_curriculum_v5_SOFT_RESET_STABLE import (
    ResolutionAgnosticInspectionEnv,
    InspectionConfig,
    NaNSafetyWrapper
)
from gru_belief_policy_v3_agnostic import create_gru_policy_kwargs
from gru_env_wrappers import GRUStateManager
# ============================================================================
# CURRICULUM (NO GRU CONTROL)
# ============================================================================
class GRUCurriculum:
    """
    Curriculum controls ONLY environment difficulty.
    GRU is always active.
    """
    def __init__(self, total_steps: int):
        self.total_steps = total_steps
        self.current_step = 0
        self.warmup_end = 0 # Skip warmup
        self.easy_end = total_steps # Easy for entire run
    def step(self):
        self.current_step += 1
    def get_phase(self) -> str:
        return "easy" # Always easy
    def get_randomization_prob(self) -> float:
        return 0.7 # 70% randomization (consistent easy mode)
# ============================================================================
# CALLBACK
# ============================================================================
class GRUTrainingCallback(BaseCallback):
    def __init__(self, curriculum, env_wrapper, log_dir, save_freq, resume_step):
        super().__init__()
        self.curriculum = curriculum
        self.env_wrapper = env_wrapper
        self.log_dir = Path(log_dir)
        self.save_freq = save_freq
        self.resume_step = resume_step
        self.last_save = 0
        for _ in range(resume_step):
            self.curriculum.step()
        self.best_catch_rate = 0.0
        self.episode_count = 0
        self.phase_stats = {"warmup": [], "easy": [], "hard": []}
        self.last_log_step = -1
        self.last_phase = self.curriculum.get_phase()
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.models_dir = self.log_dir / "models"
        self.models_dir.mkdir(exist_ok=True)
    def _on_step(self) -> bool:
        self.curriculum.step()
        phase = self.curriculum.get_phase()
        display_step = self.num_timesteps + self.resume_step
        # Update env randomization
        if hasattr(self.env_wrapper, "randomization_prob"):
            self.env_wrapper.randomization_prob = self.curriculum.get_randomization_prob()
        if phase != self.last_phase:
            print(f"\n{'='*80}")
            print(f"🔄 PHASE TRANSITION → {phase.upper()}")
            print(f"Step: {display_step:,}")
            print(f"Randomization: {self.curriculum.get_randomization_prob():.0%}")
            print(f"{'='*80}\n")
            self.last_phase = phase
        if self.num_timesteps % 1000 == 0 and self.num_timesteps != self.last_log_step:
            eps = getattr(self.model, "exploration_rate", 0.0)
            print(f"[{display_step:7,}] {phase:6s} GRU ε={eps:.3f}")
            self.last_log_step = self.num_timesteps
        infos = self.locals.get("infos", [])
        for info in infos:
            if "episode" in info:
                self.episode_count += 1
                catch_rate = info.get("catch_rate", 0.0)
                reward = info["episode"]["r"]
                self.phase_stats[phase].append(catch_rate)
                if catch_rate > self.best_catch_rate:
                    self.best_catch_rate = catch_rate
                    print(f"\n🎯 NEW BEST: {catch_rate:.4f} @ step {display_step:,}\n")
                    self.model.save(self.models_dir / "best_model_gru.zip")
        if self.num_timesteps - self.last_save >= self.save_freq:
            self.last_save = self.num_timesteps
            ckpt = self.models_dir / f"checkpoint_{display_step}.zip"
            self.model.save(ckpt)
            print(f"💾 checkpoint saved: {ckpt.name}")
            checkpoints = sorted(self.models_dir.glob("checkpoint_*.zip"))
            if len(checkpoints) > 3:
                for old in checkpoints[:-3]:
                    old.unlink()
                    print(f"🗑️ deleted {old.name}")
        return True
# ============================================================================
# MAIN
# ============================================================================
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str, default=None)
    parser.add_argument("--steps", type=int, default=200_000)
    parser.add_argument("--resume-from-step", type=int, default=0)
    parser.add_argument("--hidden-size", type=int, default=64)
    parser.add_argument("--output-dir", type=str, default="./gru_training")
    parser.add_argument("--learning-rate", type=float, default=1e-4)
    parser.add_argument("--save-freq", type=int, default=25_000)
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--from-scratch", action="store_true")
    args = parser.parse_args()
    print("\n" + "=" * 80)
    print("GRU CURRICULUM TRAINING — GRU ALWAYS ON (v2)")
    print("=" * 80)
    print(f"🔧 Memory-safe settings: buffer=50k, batch=32")
    print(f"⚡ Exploration: fraction=0.7 (floor at {int(args.steps * 0.7):,} steps), ε_min=0.10")
    print(f"💾 Saves every {args.save_freq:,} steps")
    print("=" * 80 + "\n")
    # ------------------------------------------------------------------
    # ENV
    # ------------------------------------------------------------------
    config = InspectionConfig(
        grid_size=64,
        inspection_budget=3000,
        soft_reset=True,
        seed=args.seed,
        prior_belief=0.1,
    economic_randomization=True,
    clean_episode_ratio=0.7  # CHANGED (was 1.0)
    )
    base_env = ResolutionAgnosticInspectionEnv(config=config, use_gpu=True)
    base_env = NaNSafetyWrapper(base_env)
    base_env = Monitor(base_env)
    # ------------------------------------------------------------------
    # MODEL
    # ------------------------------------------------------------------
    if args.model and not args.from_scratch:
        print(f"Loading base model: {args.model}")
        model = MaskedDQN.load(
            args.model,
            env=base_env,
            custom_objects=CUSTOM_OBJECTS,
            device=args.device
        )
        print(f"✓ Model loaded: {model.num_timesteps:,} steps, ε={model.exploration_rate:.3f}")
        # Override LR for Phase 2 conservative refinement
        model.learning_rate = 0.00007
        model.policy.optimizer.param_groups[0]['lr'] = 0.00007
        print(f"✓ LR overridden to 0.00007")
    else:
        print("Training from scratch with GRU")
        policy_kwargs = create_gru_policy_kwargs(
            hidden_size=args.hidden_size,
            use_recurrence=True
        )
        model = MaskedDQN(
            policy="MultiInputPolicy",
            env=base_env,
            policy_kwargs=policy_kwargs,
            learning_rate=0.00007,
            buffer_size=200000,
            learning_starts=2000,
            batch_size=32,
            gamma=0.99,                 # CHANGED (was 0.995)
            train_freq=4,
            gradient_steps=1,
            exploration_fraction=0.5,   # CHANGED (was 0.5) # KEY FIX: was 0.3 — epsilon now reaches
                                             # floor at 100k instead of 60k, giving
                                             # richer mid-training exploration
            exploration_initial_eps=0.15,
            exploration_final_eps=0.10, # CHANGED (was 0.05) # KEY FIX: was 0.02 — higher floor prevents
                                             # distribution narrowing and value homogenization
            device=args.device,
            verbose=1
        )
    # ------------------------------------------------------------------
    # WRAP ENV
    # ------------------------------------------------------------------
    env = GRUStateManager(base_env)
    model.set_env(env)
    env.set_policy(model.policy)
    # SANITY CHECK
    print("\n" + "="*80)
    print("GRU ARCHITECTURE VERIFICATION")
    print("="*80)
    fe = model.policy.q_net.features_extractor
    print(f"Feature Extractor Type: {type(fe).__name__}")
    print(f"GRU Enabled: {fe.use_recurrence}")
    print(f"GRU Hidden Size: {fe.hidden_size}")
    if hasattr(fe, 'belief_encoder') and fe.belief_encoder is not None:
        print(f"✓ GRU belief_encoder exists")
    else:
        print(f"✗ WARNING: No belief_encoder found!")
    print("="*80 + "\n")
    # ------------------------------------------------------------------
    # TRAIN
    # ------------------------------------------------------------------
    remaining_steps = args.steps - args.resume_from_step
    # If resuming, set epsilon to floor immediately (don't re-explore)
    if args.model and not args.from_scratch and args.resume_from_step > 0:
        model.exploration_rate = 0.10
        model.exploration_schedule = lambda _: 0.10
        print(f"✓ Resuming: epsilon fixed at 0.10")
    curriculum = GRUCurriculum(args.steps)
    callback = GRUTrainingCallback(
        curriculum,
        env,
        args.output_dir,
        args.save_freq,
        args.resume_from_step
    )
    print(f"🚀 Starting training: {remaining_steps:,} steps")
    print(f"📊 Exploration schedule:")
    print(f" 0 steps: ε=1.00 (100% random)")
    print(f" {int(args.steps * 0.5):,} steps: ε=0.10 (10% random) ← Exploitation begins")
    print(f" {args.steps:,} steps: ε=0.10 (stays at floor)")
    print()
    try:
        model.learn(
            total_timesteps=remaining_steps,
            callback=callback,
            reset_num_timesteps=False
        )
    finally:
        print("\n💾 Saving shutdown checkpoint...")
        shutdown_path = callback.models_dir / "last_shutdown_model.zip"
        model.save(str(shutdown_path))
        print(f"✅ Shutdown checkpoint saved: {shutdown_path}")
    final_path = Path(args.output_dir) / "models" / "final_model_gru.zip"
    model.save(final_path)
    print(f"\n✅ TRAINING COMPLETE – saved {final_path}")
if __name__ == "__main__":
    main()