#!/usr/bin/env python3
"""
Agent 49: Comprehensive Hyperparameter Optimization Runner

This script orchestrates hyperparameter optimization for all 4 ML models:
- DQN (Deep Q-Network)
- PPO (Proximal Policy Optimization)
- MAMBA-2 (State Space Model)
- TFT (Temporal Fusion Transformer)

Features:
- Sequential execution (one model at a time for GPU safety)
- Comprehensive results aggregation and reporting
- Performance improvement tracking
- Best hyperparameter extraction per model
- Visualization generation (optional)

Usage:
    python3 run_hyperparameter_optimization.py \
        --num-trials 50 \
        --config tuning_config_optimized.yaml \
        --data-path /path/to/market_data.parquet \
        --output-dir ./hyperparameter_results \
        --use-gpu
"""

import argparse
import json
import logging
import os
import sys
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Any, Optional
import subprocess
import uuid

import pandas as pd
import yaml

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)


class HyperparameterOptimizationRunner:
    """Orchestrates hyperparameter optimization for all models."""

    # Model priority order (from fastest to slowest training)
    MODEL_ORDER = ["DQN", "PPO", "MAMBA_2", "TFT"]

    def __init__(
        self,
        num_trials: int,
        config_path: str,
        data_path: str,
        output_dir: str,
        use_gpu: bool,
        grpc_host: str = "localhost",
        grpc_port: int = 50054,
        tuner_script: str = "./hyperparameter_tuner.py"
    ):
        self.num_trials = num_trials
        self.config_path = Path(config_path).resolve()
        self.data_path = Path(data_path).resolve()
        self.output_dir = Path(output_dir).resolve()
        self.use_gpu = use_gpu
        self.grpc_host = grpc_host
        self.grpc_port = grpc_port
        self.tuner_script = Path(tuner_script).resolve()

        # Load configuration
        with open(self.config_path, 'r') as f:
            self.config = yaml.safe_load(f)

        # Create output directories
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.results_dir = self.output_dir / "results"
        self.results_dir.mkdir(exist_ok=True)
        self.studies_dir = self.output_dir / "studies"
        self.studies_dir.mkdir(exist_ok=True)

        # Results tracking
        self.results: Dict[str, Dict[str, Any]] = {}
        self.start_time = datetime.now()

        logger.info(f"Hyperparameter Optimization Runner initialized")
        logger.info(f"  Trials per model: {num_trials}")
        logger.info(f"  Models to optimize: {', '.join(self.MODEL_ORDER)}")
        logger.info(f"  GPU enabled: {use_gpu}")
        logger.info(f"  Output directory: {self.output_dir}")

    def prepare_data_source(self) -> Dict[str, Any]:
        """Prepare data source configuration for training."""
        # Use sample data window (last 7 days of available data)
        # In production, this should be parameterized
        data_source = {
            "file_path": str(self.data_path),
            "start_time": 1633046400,  # Sample: 2021-10-01 (Unix timestamp)
            "end_time": 1633651200     # Sample: 2021-10-08 (7 days)
        }
        return data_source

    def run_optimization_for_model(self, model_type: str) -> Dict[str, Any]:
        """Run hyperparameter optimization for a single model."""
        logger.info(f"\n{'='*80}")
        logger.info(f"Starting hyperparameter optimization for {model_type}")
        logger.info(f"{'='*80}\n")

        job_id = str(uuid.uuid4())
        storage_path = self.studies_dir / f"study_{model_type}_{job_id}.log"
        data_source = self.prepare_data_source()

        # Build command
        cmd = [
            "python3",
            str(self.tuner_script),
            "--job-id", job_id,
            "--model-type", model_type,
            "--num-trials", str(self.num_trials),
            "--config", str(self.config_path),
            "--data-source-json", json.dumps(data_source),
            "--storage-path", str(storage_path),
            "--grpc-host", self.grpc_host,
            "--grpc-port", str(self.grpc_port)
        ]

        if self.use_gpu:
            cmd.append("--use-gpu")

        # Run optimization subprocess
        model_start_time = datetime.now()
        logger.info(f"Executing: {' '.join(cmd)}")

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=7200  # 2 hour timeout per model
            )

            model_end_time = datetime.now()
            duration = (model_end_time - model_start_time).total_seconds()

            # Log output
            logger.info(f"\n--- {model_type} STDOUT ---")
            logger.info(result.stdout)

            if result.stderr:
                logger.warning(f"\n--- {model_type} STDERR ---")
                logger.warning(result.stderr)

            # Parse results from Optuna study file
            results = self.parse_optuna_results(storage_path, model_type)
            results["duration_seconds"] = duration
            results["job_id"] = job_id
            results["success"] = result.returncode == 0

            # Save model-specific results
            self.save_model_results(model_type, results)

            logger.info(f"\n{model_type} optimization completed in {duration/60:.1f} minutes")
            logger.info(f"Best Sharpe ratio: {results.get('best_sharpe', 0.0):.4f}")

            return results

        except subprocess.TimeoutExpired:
            logger.error(f"{model_type} optimization timed out after 2 hours")
            return {
                "model_type": model_type,
                "success": False,
                "error": "Timeout after 2 hours",
                "best_sharpe": 0.0,
                "best_params": {},
                "improvement": 0.0
            }
        except Exception as e:
            logger.error(f"{model_type} optimization failed: {e}", exc_info=True)
            return {
                "model_type": model_type,
                "success": False,
                "error": str(e),
                "best_sharpe": 0.0,
                "best_params": {},
                "improvement": 0.0
            }

    def parse_optuna_results(self, storage_path: Path, model_type: str) -> Dict[str, Any]:
        """Parse Optuna JournalStorage results."""
        try:
            import optuna
            from optuna.storages import JournalStorage, JournalFileStorage

            # Load study
            file_storage = JournalFileStorage(str(storage_path))
            storage = JournalStorage(file_storage)
            study_name = f"study_{os.path.basename(storage_path).split('_')[1]}"  # Extract job_id

            # Note: Study name format is "study_{job_id}"
            # We need to find the actual study name from the file
            # For simplicity, let's assume the study name matches our pattern

            # Try to load with pattern matching
            studies = []
            try:
                # Optuna 3.0+ API
                studies = storage.get_all_study_summaries()
                if studies:
                    study_name = studies[0].study_name
            except:
                # Fallback: assume standard naming
                pass

            study = optuna.load_study(
                study_name=study_name,
                storage=storage
            )

            # Extract best trial
            best_trial = study.best_trial

            # Calculate performance improvement (baseline = 0.0)
            baseline_sharpe = 0.0  # Default baseline
            best_sharpe = best_trial.value
            improvement = ((best_sharpe - baseline_sharpe) / max(abs(baseline_sharpe), 1e-6)) * 100

            # Count trial outcomes
            completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
            pruned_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]
            failed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.FAIL]

            return {
                "model_type": model_type,
                "best_sharpe": best_sharpe,
                "best_params": best_trial.params,
                "best_trial_number": best_trial.number,
                "num_completed_trials": len(completed_trials),
                "num_pruned_trials": len(pruned_trials),
                "num_failed_trials": len(failed_trials),
                "improvement_percent": improvement,
                "baseline_sharpe": baseline_sharpe,
                "study_name": study_name,
                "storage_path": str(storage_path)
            }

        except Exception as e:
            logger.error(f"Failed to parse Optuna results for {model_type}: {e}")
            return {
                "model_type": model_type,
                "best_sharpe": 0.0,
                "best_params": {},
                "improvement_percent": 0.0,
                "num_completed_trials": 0,
                "num_pruned_trials": 0,
                "num_failed_trials": 0,
                "error": str(e)
            }

    def save_model_results(self, model_type: str, results: Dict[str, Any]):
        """Save model-specific results to JSON."""
        output_file = self.results_dir / f"{model_type}_results.json"
        with open(output_file, 'w') as f:
            json.dump(results, f, indent=2)
        logger.info(f"Results saved to {output_file}")

    def run_all_optimizations(self):
        """Run hyperparameter optimization for all models sequentially."""
        logger.info("\n" + "="*80)
        logger.info("STARTING COMPREHENSIVE HYPERPARAMETER OPTIMIZATION")
        logger.info("="*80 + "\n")

        for model_type in self.MODEL_ORDER:
            results = self.run_optimization_for_model(model_type)
            self.results[model_type] = results

            # Brief pause between models (GPU memory cleanup)
            if self.use_gpu and model_type != self.MODEL_ORDER[-1]:
                logger.info("Pausing 30 seconds for GPU memory cleanup...")
                time.sleep(30)

    def generate_summary_report(self) -> str:
        """Generate comprehensive summary report."""
        total_duration = (datetime.now() - self.start_time).total_seconds()

        report = []
        report.append("\n" + "="*80)
        report.append("HYPERPARAMETER OPTIMIZATION SUMMARY REPORT")
        report.append("Agent 49 - Foxhunt HFT Trading System")
        report.append("="*80 + "\n")

        report.append(f"Execution Time: {total_duration/60:.1f} minutes")
        report.append(f"Trials per Model: {self.num_trials}")
        report.append(f"GPU Enabled: {self.use_gpu}")
        report.append(f"Timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")

        # Model-by-model results
        report.append("="*80)
        report.append("MODEL-BY-MODEL RESULTS")
        report.append("="*80 + "\n")

        for model_type in self.MODEL_ORDER:
            if model_type not in self.results:
                continue

            result = self.results[model_type]
            report.append(f"### {model_type} ###")
            report.append(f"Status: {'SUCCESS' if result.get('success', False) else 'FAILED'}")
            report.append(f"Best Sharpe Ratio: {result.get('best_sharpe', 0.0):.4f}")
            report.append(f"Performance Improvement: {result.get('improvement_percent', 0.0):.2f}%")
            report.append(f"Completed Trials: {result.get('num_completed_trials', 0)}")
            report.append(f"Pruned Trials: {result.get('num_pruned_trials', 0)}")
            report.append(f"Failed Trials: {result.get('num_failed_trials', 0)}")
            report.append(f"Duration: {result.get('duration_seconds', 0)/60:.1f} minutes")

            # Best hyperparameters
            best_params = result.get('best_params', {})
            if best_params:
                report.append("Best Hyperparameters:")
                for param_name, param_value in sorted(best_params.items()):
                    report.append(f"  - {param_name}: {param_value}")
            report.append("")

        # Aggregate statistics
        report.append("="*80)
        report.append("AGGREGATE STATISTICS")
        report.append("="*80 + "\n")

        successful_models = [m for m, r in self.results.items() if r.get('success', False)]
        failed_models = [m for m, r in self.results.items() if not r.get('success', False)]

        report.append(f"Successful Optimizations: {len(successful_models)}/{len(self.MODEL_ORDER)}")
        report.append(f"Failed Optimizations: {len(failed_models)}/{len(self.MODEL_ORDER)}")

        if successful_models:
            avg_sharpe = sum(self.results[m]['best_sharpe'] for m in successful_models) / len(successful_models)
            report.append(f"Average Best Sharpe Ratio: {avg_sharpe:.4f}")

            best_model = max(successful_models, key=lambda m: self.results[m]['best_sharpe'])
            report.append(f"Best Performing Model: {best_model} (Sharpe: {self.results[best_model]['best_sharpe']:.4f})")

        report.append("")

        # Search space coverage
        report.append("="*80)
        report.append("SEARCH SPACE ANALYSIS")
        report.append("="*80 + "\n")

        report.append("Agent 49 Specified Search Spaces:")
        report.append("- DQN: Learning rate [1e-5, 1e-4, 1e-3], Batch [64, 128, 256], Gamma [0.95, 0.99, 0.999]")
        report.append("- PPO: Learning rate [3e-5, 1e-4, 3e-4], Entropy [0.01, 0.05, 0.1], Clip [0.1, 0.2, 0.3]")
        report.append("- MAMBA-2: Learning rate [1e-5, 1e-4, 1e-3], State size [16, 32, 64], Layers [4, 6, 8]")
        report.append("- TFT: Learning rate [1e-5, 1e-4, 1e-3], Attention heads [4, 8, 16], Hidden [128, 256, 512]")
        report.append("")
        report.append("Search Method: Bayesian Optimization (TPE Sampler)")
        report.append("Grid Combinations per Model: 27 (3^3)")
        report.append(f"Actual Trials per Model: {self.num_trials} (explores beyond grid)")
        report.append("")

        # Recommendations
        report.append("="*80)
        report.append("RECOMMENDATIONS")
        report.append("="*80 + "\n")

        if failed_models:
            report.append("⚠️  FAILED MODELS:")
            for model in failed_models:
                error = self.results[model].get('error', 'Unknown error')
                report.append(f"  - {model}: {error}")
            report.append("")

        if successful_models:
            report.append("✅ PRODUCTION RECOMMENDATIONS:")
            for model in successful_models:
                sharpe = self.results[model]['best_sharpe']
                if sharpe > 1.5:
                    report.append(f"  - {model}: EXCELLENT performance (Sharpe {sharpe:.2f}) - Deploy immediately")
                elif sharpe > 1.0:
                    report.append(f"  - {model}: GOOD performance (Sharpe {sharpe:.2f}) - Deploy with monitoring")
                elif sharpe > 0.5:
                    report.append(f"  - {model}: MODERATE performance (Sharpe {sharpe:.2f}) - Further tuning recommended")
                else:
                    report.append(f"  - {model}: LOW performance (Sharpe {sharpe:.2f}) - Investigate data/features")
            report.append("")

        report.append("="*80)
        report.append("END OF REPORT")
        report.append("="*80 + "\n")

        return "\n".join(report)

    def save_summary_report(self):
        """Save summary report to file and print to console."""
        report = self.generate_summary_report()

        # Save to file
        report_file = self.output_dir / "hyperparameter_optimization_report.txt"
        with open(report_file, 'w') as f:
            f.write(report)

        # Save aggregate results as JSON
        aggregate_file = self.output_dir / "aggregate_results.json"
        with open(aggregate_file, 'w') as f:
            json.dump(self.results, f, indent=2)

        # Print to console
        print(report)

        logger.info(f"\nReports saved:")
        logger.info(f"  - {report_file}")
        logger.info(f"  - {aggregate_file}")
        logger.info(f"  - Individual model results in {self.results_dir}/")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Agent 49: Comprehensive Hyperparameter Optimization Runner"
    )

    parser.add_argument(
        "--num-trials",
        type=int,
        default=50,
        help="Number of optimization trials per model (default: 50)"
    )

    parser.add_argument(
        "--config",
        type=str,
        default="tuning_config_optimized.yaml",
        help="Path to tuning configuration file (default: tuning_config_optimized.yaml)"
    )

    parser.add_argument(
        "--data-path",
        type=str,
        required=True,
        help="Path to market data Parquet file for training"
    )

    parser.add_argument(
        "--output-dir",
        type=str,
        default="./hyperparameter_results",
        help="Output directory for results (default: ./hyperparameter_results)"
    )

    parser.add_argument(
        "--use-gpu",
        action="store_true",
        help="Enable GPU acceleration"
    )

    parser.add_argument(
        "--grpc-host",
        type=str,
        default="localhost",
        help="ML Training Service gRPC host (default: localhost)"
    )

    parser.add_argument(
        "--grpc-port",
        type=int,
        default=50054,
        help="ML Training Service gRPC port (default: 50054)"
    )

    parser.add_argument(
        "--tuner-script",
        type=str,
        default="./hyperparameter_tuner.py",
        help="Path to hyperparameter tuner script (default: ./hyperparameter_tuner.py)"
    )

    args = parser.parse_args()

    # Validate inputs
    if not Path(args.data_path).exists():
        logger.error(f"Data path not found: {args.data_path}")
        sys.exit(1)

    if not Path(args.config).exists():
        logger.error(f"Config file not found: {args.config}")
        sys.exit(1)

    if not Path(args.tuner_script).exists():
        logger.error(f"Tuner script not found: {args.tuner_script}")
        sys.exit(1)

    # Create runner
    runner = HyperparameterOptimizationRunner(
        num_trials=args.num_trials,
        config_path=args.config,
        data_path=args.data_path,
        output_dir=args.output_dir,
        use_gpu=args.use_gpu,
        grpc_host=args.grpc_host,
        grpc_port=args.grpc_port,
        tuner_script=args.tuner_script
    )

    # Run optimizations
    try:
        runner.run_all_optimizations()
        runner.save_summary_report()
        logger.info("\n✅ Hyperparameter optimization completed successfully!")
        sys.exit(0)
    except Exception as e:
        logger.error(f"\n❌ Hyperparameter optimization failed: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()