#!/usr/bin/env python3
"""
Train PPO agent for CFD flow control.

This script trains a Proximal Policy Optimization (PPO) agent to control
flow around a cylinder using the CFD environment.
"""

import argparse
import os
import pickle
from pathlib import Path
import sys

# Set threading layers
os.environ['MKL_THREADING_LAYER'] = 'GNU'
os.environ["OMP_NUM_THREADS"] = "8"
os.environ["MKL_NUM_THREADS"] = "8"

import numpy as np
import torch
from torch.nn import Module
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from torch.utils.tensorboard import SummaryWriter

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))

from environments import CFDFlowControlEnv
from config import (
    load_celeris_configs,
    get_model_path,
    get_tensorboard_logdir,
    get_output_path,
)


class SinActivation(Module):
    """Sine activation function for neural networks."""
    
    def __init__(self):
        super().__init__()
    
    def forward(self, x):
        return torch.sin(x)


class TensorboardCallback(BaseCallback):
    """
    Custom callback for logging additional metrics to TensorBoard.
    """
    
    def __init__(self, check_freq: int = 360, verbose: int = 0):
        super().__init__(verbose)
        self.check_freq = check_freq
        self.episode_rewards = []
        self.episode_cd = []
        self.episode_cl = []
    
    def _on_step(self) -> bool:
        if self.n_calls % self.check_freq == 0:
            # Extract episode info
            if len(self.locals.get('infos', [])) > 0:
                info = self.locals['infos'][0]
                if 'cd' in info:
                    self.logger.record('flow/cd', info['cd'])
                if 'cl' in info:
                    self.logger.record('flow/cl', info['cl'])
                if 'reward_cd' in info:
                    self.logger.record('reward/cd', info['reward_cd'])
                if 'reward_cl' in info:
                    self.logger.record('reward/cl', info['reward_cl'])
                if 'reward_sim' in info:
                    self.logger.record('reward/sim', info['reward_sim'])
        return True


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description='Train PPO for CFD control')
    
    # Environment settings
    parser.add_argument('--device-id', type=int, default=0,
                        help='CUDA device ID for simulation')
    
    # Training hyperparameters
    parser.add_argument('--total-timesteps', type=int, default=100,
                        help='Number of training iterations (each = n_steps)')
    parser.add_argument('--n-steps', type=int, default=3600,
                        help='Steps to collect per training iteration')
    parser.add_argument('--batch-size', type=int, default=360,
                        help='Batch size for PPO updates')
    parser.add_argument('--learning-rate', type=float, default=3e-4,
                        help='Learning rate')
    parser.add_argument('--gamma', type=float, default=0.99,
                        help='Discount factor')
    
    # Model settings  
    parser.add_argument('--activation', choices=['tanh', 'relu', 'sin'], default='sin',
                        help='Activation function for policy network')
    parser.add_argument('--cuda-device', type=int, default=0,
                        help='CUDA device for PyTorch training')
    
    # Experiment settings
    parser.add_argument('--run-name', type=str, default='ppo_cfd_control',
                        help='Name for this training run')
    parser.add_argument('--save-freq', type=int, default=10,
                        help='Save model every N iterations')
    parser.add_argument('--eval-episodes', type=int, default=1,
                        help='Number of episodes to evaluate')
    
    # Resume training
    parser.add_argument('--resume', type=str, default=None,
                        help='Path to model checkpoint to resume from')
    
    return parser.parse_args()


def create_env(device_id: int):
    """Create the CFD environment."""
    config_cuda, config_field = load_celeris_configs()
    
    env = CFDFlowControlEnv(
        device_id=device_id,
        config_cuda=config_cuda,
        config_field=config_field,
    )
    
    return env


def get_activation_fn(name: str):
    """Get activation function by name."""
    if name == 'sin':
        return SinActivation
    elif name == 'tanh':
        return torch.nn.Tanh
    elif name == 'relu':
        return torch.nn.ReLU
    else:
        raise ValueError(f"Unknown activation: {name}")


def evaluate_policy(model, env, n_episodes: int = 1):
    """
    Evaluate the trained policy.
    
    Returns:
        Dictionary of evaluation metrics
    """
    episode_rewards = []
    episode_data = []
    
    for episode in range(n_episodes):
        obs, info = env.reset()
        done = False
        episode_reward = 0
        steps = 0
        
        ep_data = {
            'actions': [],
            'observations': [],
            'rewards': [],
            'cd': [],
            'cl': [],
        }
        
        while not done:
            action, _states = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            
            episode_reward += reward
            steps += 1
            
            # Record data
            ep_data['actions'].append(action)
            ep_data['observations'].append(obs)
            ep_data['rewards'].append(reward)
            ep_data['cd'].append(info.get('cd', 0))
            ep_data['cl'].append(info.get('cl', 0))
        
        episode_rewards.append(episode_reward)
        episode_data.append(ep_data)
        
        print(f"  Episode {episode + 1}/{n_episodes}: "
              f"Reward = {episode_reward:.2f}, "
              f"Steps = {steps}, "
              f"Avg CD = {np.mean(ep_data['cd']):.4f}")
    
    return {
        'mean_reward': np.mean(episode_rewards),
        'std_reward': np.std(episode_rewards),
        'episodes': episode_data,
    }


def main():
    """Main training loop."""
    args = parse_args()
    
    print("=" * 70)
    print(f"DynamisLab - CFD Flow Control Training")
    print(f"Run: {args.run_name}")
    print("=" * 70)
    
    # Create environment
    print(f"\n[1/4] Creating CFD environment on GPU:{args.device_id}...")
    env = create_env(args.device_id)
    print(f"  Action space: {env.action_space}")
    print(f"  Observation space: {env.observation_space}")
    
    # Create or load model
    print(f"\n[2/4] Setting up PPO model...")
    device = torch.device(f"cuda:{args.cuda_device}")
    
    if args.resume:
        print(f"  Resuming from: {args.resume}")
        model = PPO.load(args.resume, env=env, device=device)
    else:
        activation_fn = get_activation_fn(args.activation)
        model = PPO(
            "MlpPolicy",
            env=env,
            learning_rate=args.learning_rate,
            n_steps=args.n_steps,
            batch_size=args.batch_size,
            gamma=args.gamma,
            policy_kwargs=dict(activation_fn=activation_fn),
            device=device,
            verbose=1,
        )
    
    print(f"  Activation: {args.activation}")
    print(f"  Device: {device}")
    print(f"  Learning rate: {args.learning_rate}")
    print(f"  Steps per iteration: {args.n_steps}")
    print(f"  Batch size: {args.batch_size}")
    
    # Setup logging
    tensorboard_dir = get_tensorboard_logdir(args.run_name)
    writer = SummaryWriter(log_dir=str(tensorboard_dir))
    print(f"  TensorBoard: {tensorboard_dir}")
    
    # Training loop
    print(f"\n[3/4] Training for {args.total_timesteps} iterations...")
    
    best_reward = -np.inf
    history_data = []
    
    for iteration in range(args.total_timesteps):
        # Train
        model.learn(total_timesteps=args.n_steps, reset_num_timesteps=False)
        
        # Evaluate
        print(f"\n--- Iteration {iteration + 1}/{args.total_timesteps} ---")
        eval_results = evaluate_policy(model, env, n_episodes=args.eval_episodes)
        
        mean_reward = eval_results['mean_reward']
        std_reward = eval_results['std_reward']
        
        # Log to TensorBoard
        writer.add_scalar('eval/mean_reward', mean_reward, iteration)
        writer.add_scalar('eval/std_reward', std_reward, iteration)
        
        # Extract CD/CL from last episode
        if len(eval_results['episodes']) > 0:
            last_ep = eval_results['episodes'][-1]
            avg_cd = np.mean(last_ep['cd'])
            avg_cl = np.mean(last_ep['cl'])
            writer.add_scalar('eval/avg_cd', avg_cd, iteration)
            writer.add_scalar('eval/avg_cl', avg_cl, iteration)
        
        # Save best model
        if mean_reward > best_reward:
            best_reward = mean_reward
            model_path = get_model_path(f"{args.run_name}_best")
            model.save(str(model_path))
            print(f"  ✓ New best model saved: {model_path} (reward: {mean_reward:.2f})")
        
        # Periodic save
        if (iteration + 1) % args.save_freq == 0:
            model_path = get_model_path(f"{args.run_name}_iter{iteration + 1}")
            model.save(str(model_path))
            print(f"  Checkpoint saved: {model_path}")
        
        # Store history
        history_data.append(eval_results['episodes'])
    
    # Final evaluation
    print(f"\n[4/4] Final evaluation...")
    final_results = evaluate_policy(model, env, n_episodes=5)
    print(f"  Final mean reward: {final_results['mean_reward']:.2f} ± {final_results['std_reward']:.2f}")
    
    # Save final model and history
    final_model_path = get_model_path(f"{args.run_name}_final")
    model.save(str(final_model_path))
    print(f"  Final model saved: {final_model_path}")
    
    history_path = get_output_path(f"{args.run_name}_history.pkl")
    with open(history_path, 'wb') as f:
        pickle.dump(history_data, f)
    print(f"  Training history saved: {history_path}")
    
    # Cleanup
    writer.close()
    env.close()
    
    print("\n" + "=" * 70)
    print("Training complete!")
    print("=" * 70)


if __name__ == '__main__':
    main()