#!/usr/bin/env python3
"""
Optuna Hyperparameter Tuner for Foxhunt ML Models

This subprocess is spawned by the ML Training Service to perform hyperparameter optimization
using Optuna 3.0+ with JournalStorage for crash recovery. Coordinates with the Rust service
via gRPC to train models with sampled hyperparameters and optimize for Sharpe ratio.

Requirements:
- Optuna 3.0+ with JournalStorage (file-based persistence)
- Sequential trials (n_jobs=1) for GPU memory safety (4GB VRAM constraint)
- MedianPruner for early stopping (30-50% time savings)
- pynvml for GPU memory monitoring
- Graceful shutdown on SIGTERM

MedianPruner Configuration (tuning_config.yaml):
- n_startup_trials: 5 (no pruning for first 5 trials to establish baseline median)
- n_warmup_steps: 10 (wait 10 epochs before starting to prune, gives models time to converge)
- interval_steps: 5 (check for pruning every 5 epochs, balance overhead vs responsiveness)

Pruning Logic:
- After warmup period (10 epochs), MedianPruner compares trial's intermediate Sharpe ratio
  to the median of completed trials at the same step
- If current_sharpe < median_sharpe → trial is pruned (stopped early)
- Expected impact: 30-50% reduction in total tuning time by eliminating unpromising trials

Current Limitations:
- gRPC TrainModel returns only final metrics (no streaming intermediate values)
- MedianPruner compares final Sharpe ratios across trials (inter-trial pruning)
- For intra-trial early stopping, TrainModel would need to support:
  1. Streaming responses with epoch-by-epoch metrics OR
  2. Callback mechanism for intermediate reporting OR
  3. Status polling endpoint for querying training progress

Architecture:
- Reads search spaces from tuning_config.yaml
- Calls TrainModel gRPC endpoint (localhost:50054) for each trial
- Reports final Sharpe ratio via trial.report() for MedianPruner
- Persists study state to MinIO mount after each trial
- Reports progress via stdout (captured by Rust service)

Usage:
    python3 hyperparameter_tuner.py \
        --job-id <tuning_job_id> \
        --model-type TLOB \
        --num-trials 50 \
        --config tuning_config.yaml \
        --data-source-json '{"file_path": "data.parquet", "start_time": 1633046400, "end_time": 1633132800}' \
        --use-gpu \
        --storage-path /minio/studies/study_<job_id>.log
"""

import argparse
import json
import logging
import os
import signal
import sys
import time
from typing import Dict, Any, Optional

import grpc
import optuna
from optuna.pruners import MedianPruner
from optuna.storages import JournalStorage, JournalFileStorage
import yaml

# GPU monitoring
try:
    import pynvml
    pynvml.nvmlInit()
    GPU_AVAILABLE = True
except Exception:
    GPU_AVAILABLE = False

# 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__)

# Global shutdown flag
shutdown_requested = False


def signal_handler(signum, frame):
    """Handle SIGTERM/SIGINT for graceful shutdown."""
    global shutdown_requested
    logger.info(f"Received signal {signum}, initiating graceful shutdown...")
    shutdown_requested = True


# Register signal handlers
signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler)


class GPUMonitor:
    """Monitor GPU memory usage using pynvml (not subprocess)."""

    def __init__(self):
        self.enabled = GPU_AVAILABLE
        if self.enabled:
            try:
                self.device_count = pynvml.nvmlDeviceGetCount()
                logger.info(f"GPU monitoring enabled: {self.device_count} device(s) detected")
            except Exception as e:
                logger.warning(f"GPU monitoring initialization failed: {e}")
                self.enabled = False

    def get_memory_usage(self, device_id: int = 0) -> Dict[str, float]:
        """Get GPU memory usage in GB."""
        if not self.enabled:
            return {"used_gb": 0.0, "total_gb": 0.0, "percent": 0.0}

        try:
            handle = pynvml.nvmlDeviceGetHandleByIndex(device_id)
            mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)

            used_gb = mem_info.used / (1024 ** 3)
            total_gb = mem_info.total / (1024 ** 3)
            percent = (mem_info.used / mem_info.total) * 100.0

            return {
                "used_gb": used_gb,
                "total_gb": total_gb,
                "percent": percent
            }
        except Exception as e:
            logger.warning(f"Failed to read GPU memory: {e}")
            return {"used_gb": 0.0, "total_gb": 0.0, "percent": 0.0}

    def check_memory_available(self, device_id: int = 0, required_gb: float = 2.0) -> bool:
        """Check if sufficient GPU memory is available."""
        if not self.enabled:
            return False

        mem_usage = self.get_memory_usage(device_id)
        available_gb = mem_usage["total_gb"] - mem_usage["used_gb"]

        if available_gb < required_gb:
            logger.warning(
                f"Insufficient GPU memory: {available_gb:.2f} GB available, "
                f"{required_gb:.2f} GB required"
            )
            return False

        return True


class GRPCModelTrainer:
    """Client for calling TrainModel gRPC endpoint."""

    def __init__(self, grpc_host: str = "localhost", grpc_port: int = 50054):
        self.address = f"{grpc_host}:{grpc_port}"
        self.channel = None
        self.stub = None
        logger.info(f"gRPC client initialized for {self.address}")

    def connect(self):
        """Establish gRPC channel."""
        try:
            self.channel = grpc.insecure_channel(self.address)
            # Import generated proto stubs
            from proto import ml_training_pb2
            from proto import ml_training_pb2_grpc

            self.stub = ml_training_pb2_grpc.MLTrainingServiceStub(self.channel)

            # Test connection with health check
            request = ml_training_pb2.HealthCheckRequest()
            response = self.stub.HealthCheck(request, timeout=5.0)

            if response.healthy:
                logger.info("gRPC connection established successfully")
            else:
                logger.warning(f"gRPC service reports unhealthy: {response.message}")
        except Exception as e:
            logger.error(f"Failed to connect to gRPC service: {e}")
            raise

    def train_model(
        self,
        model_type: str,
        hyperparameters: Dict[str, float],
        data_source: Dict[str, Any],
        use_gpu: bool,
        trial_id: str,
        trial: Optional[optuna.Trial] = None
    ) -> Dict[str, Any]:
        """
        Call TrainModel gRPC endpoint and return results.

        Args:
            trial: Optional Optuna trial for intermediate reporting and pruning.
                   If provided, will report intermediate Sharpe ratios and check for pruning.

        Returns:
            dict: {
                "success": bool,
                "sharpe_ratio": float,
                "training_loss": float,
                "validation_metrics": dict,
                "error_message": str,
                "training_duration_seconds": int,
                "was_pruned": bool  # True if trial was pruned early
            }
        """
        from proto import ml_training_pb2

        # Build DataSource message
        data_source_msg = ml_training_pb2.DataSource()
        if "file_path" in data_source:
            data_source_msg.file_path = data_source["file_path"]
        elif "historical_db_query" in data_source:
            data_source_msg.historical_db_query = data_source["historical_db_query"]
        elif "real_time_stream_topic" in data_source:
            data_source_msg.real_time_stream_topic = data_source["real_time_stream_topic"]

        if "start_time" in data_source:
            data_source_msg.start_time = int(data_source["start_time"])
        if "end_time" in data_source:
            data_source_msg.end_time = int(data_source["end_time"])

        # Build TrainModelRequest
        request = ml_training_pb2.TrainModelRequest(
            model_type=model_type,
            hyperparameters=hyperparameters,
            data_source=data_source_msg,
            use_gpu=use_gpu,
            trial_id=trial_id
        )

        try:
            logger.info(f"Trial {trial_id}: Calling TrainModel gRPC with {len(hyperparameters)} params")

            # Call gRPC with generous timeout (training can take minutes)
            response = self.stub.TrainModel(request, timeout=3600.0)

            result = {
                "success": response.success,
                "sharpe_ratio": response.sharpe_ratio,
                "training_loss": response.training_loss,
                "validation_metrics": dict(response.validation_metrics),
                "error_message": response.error_message,
                "training_duration_seconds": response.training_duration_seconds,
                "was_pruned": False
            }

            # Handle intermediate reporting and pruning (future enhancement)
            # Note: Current gRPC interface returns final metrics only.
            # For intermediate reporting, the TrainModel endpoint would need to:
            # 1. Support streaming responses OR
            # 2. Accept a callback URL/endpoint OR
            # 3. Store intermediate values that we poll
            #
            # For now, we report the final Sharpe ratio at the end of training.
            # This still benefits from MedianPruner because it compares final values
            # across trials, though we lose intra-trial early stopping.
            if trial is not None and result["success"]:
                # Report final Sharpe ratio (treated as intermediate value for last epoch)
                # Optuna's MedianPruner will use this for future trial comparisons
                total_epochs = int(hyperparameters.get("epochs", 100))
                trial.report(result["sharpe_ratio"], step=total_epochs)

                # Note: We can't prune at this point since training is already complete.
                # True intra-trial pruning requires the Rust service to support callbacks
                # or streaming intermediate metrics.

            if result["success"]:
                logger.info(
                    f"Trial {trial_id}: Training succeeded - "
                    f"Sharpe={result['sharpe_ratio']:.4f}, Loss={result['training_loss']:.6f}, "
                    f"Duration={result['training_duration_seconds']}s"
                )
            else:
                logger.error(f"Trial {trial_id}: Training failed - {result['error_message']}")

            return result

        except grpc.RpcError as e:
            logger.error(f"Trial {trial_id}: gRPC error - {e.code()}: {e.details()}")
            return {
                "success": False,
                "sharpe_ratio": 0.0,
                "training_loss": float('inf'),
                "validation_metrics": {},
                "error_message": f"gRPC error: {e.code()} - {e.details()}",
                "training_duration_seconds": 0,
                "was_pruned": False
            }
        except Exception as e:
            logger.error(f"Trial {trial_id}: Unexpected error - {e}")
            return {
                "success": False,
                "sharpe_ratio": 0.0,
                "training_loss": float('inf'),
                "validation_metrics": {},
                "error_message": f"Unexpected error: {str(e)}",
                "training_duration_seconds": 0,
                "was_pruned": False
            }

    def close(self):
        """Close gRPC channel."""
        if self.channel:
            self.channel.close()
            logger.info("gRPC connection closed")


class HyperparameterTuner:
    """Optuna-based hyperparameter optimization coordinator."""

    def __init__(
        self,
        job_id: str,
        model_type: str,
        num_trials: int,
        config_path: str,
        data_source: Dict[str, Any],
        use_gpu: bool,
        storage_path: str,
        grpc_host: str = "localhost",
        grpc_port: int = 50054
    ):
        self.job_id = job_id
        self.model_type = model_type
        self.num_trials = num_trials
        self.config_path = config_path
        self.data_source = data_source
        self.use_gpu = use_gpu
        self.storage_path = storage_path

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

        # Initialize components
        self.gpu_monitor = GPUMonitor()
        self.grpc_client = GRPCModelTrainer(grpc_host, grpc_port)
        self.study = None

        logger.info(f"Tuner initialized: job_id={job_id}, model={model_type}, trials={num_trials}")

    def create_study(self):
        """Create or load Optuna study with JournalStorage."""
        global_config = self.config.get("global", {})

        # Configure JournalStorage for crash recovery
        file_storage = JournalFileStorage(self.storage_path)
        storage = JournalStorage(file_storage)

        # Configure MedianPruner
        pruner_config = global_config.get("median_pruner", {})
        pruner = MedianPruner(
            n_startup_trials=pruner_config.get("n_startup_trials", 5),
            n_warmup_steps=pruner_config.get("n_warmup_steps", 0),
            interval_steps=pruner_config.get("interval_steps", 1)
        )

        # Create or load study
        direction = global_config.get("optimization_direction", "maximize")
        study_name = f"study_{self.job_id}"

        self.study = optuna.create_study(
            study_name=study_name,
            storage=storage,
            load_if_exists=True,  # Resume from crash
            direction=direction,
            pruner=pruner,
            sampler=optuna.samplers.TPESampler()
        )

        logger.info(
            f"Study created: name={study_name}, direction={direction}, "
            f"storage={self.storage_path}"
        )

    def suggest_hyperparameters(self, trial: optuna.Trial) -> Dict[str, float]:
        """Sample hyperparameters from search space defined in config."""
        model_config = self.config["models"].get(self.model_type)
        if not model_config:
            raise ValueError(f"No search space defined for model type: {self.model_type}")

        params = {}

        for param_name, param_spec in model_config.items():
            param_type = param_spec["type"]

            if param_type == "int":
                params[param_name] = float(trial.suggest_int(
                    param_name,
                    param_spec["low"],
                    param_spec["high"],
                    step=param_spec.get("step", 1)
                ))

            elif param_type == "float":
                log_scale = param_spec.get("log", False)
                step = param_spec.get("step")

                if step:
                    params[param_name] = trial.suggest_float(
                        param_name,
                        param_spec["low"],
                        param_spec["high"],
                        step=step,
                        log=log_scale
                    )
                else:
                    params[param_name] = trial.suggest_float(
                        param_name,
                        param_spec["low"],
                        param_spec["high"],
                        log=log_scale
                    )

            elif param_type == "categorical":
                choices = param_spec["choices"]
                selected = trial.suggest_categorical(param_name, choices)

                # Convert booleans to float for gRPC map<string, float>
                if isinstance(selected, bool):
                    params[param_name] = 1.0 if selected else 0.0
                else:
                    params[param_name] = float(selected)

            else:
                raise ValueError(f"Unknown parameter type: {param_type}")

        return params

    def objective(self, trial: optuna.Trial) -> float:
        """
        Optuna objective function: train model and return Sharpe ratio.

        This function is called by Optuna for each trial.
        """
        global shutdown_requested

        # Check for shutdown signal
        if shutdown_requested:
            logger.info("Shutdown requested, aborting trial")
            raise optuna.TrialPruned()

        # Check GPU memory if needed
        if self.use_gpu and not self.gpu_monitor.check_memory_available(required_gb=2.0):
            logger.warning("Insufficient GPU memory, pruning trial")
            raise optuna.TrialPruned()

        # Sample hyperparameters
        trial_id = f"{self.job_id}_trial_{trial.number}"
        hyperparameters = self.suggest_hyperparameters(trial)

        logger.info(f"Trial {trial.number}/{self.num_trials}: {hyperparameters}")

        # Train model via gRPC (pass trial for intermediate reporting)
        result = self.grpc_client.train_model(
            model_type=self.model_type,
            hyperparameters=hyperparameters,
            data_source=self.data_source,
            use_gpu=self.use_gpu,
            trial_id=trial_id,
            trial=trial  # Pass trial for intermediate reporting and pruning
        )

        # Check if trial was pruned
        if result.get("was_pruned", False):
            logger.info(f"Trial {trial.number} was pruned early by MedianPruner")
            raise optuna.TrialPruned()

        # Report GPU usage
        if self.use_gpu:
            gpu_mem = self.gpu_monitor.get_memory_usage()
            logger.info(
                f"Trial {trial.number}: GPU memory: "
                f"{gpu_mem['used_gb']:.2f}/{gpu_mem['total_gb']:.2f} GB "
                f"({gpu_mem['percent']:.1f}%)"
            )

        # Handle failure
        if not result["success"]:
            logger.error(f"Trial {trial.number} failed: {result['error_message']}")
            # Return worst possible Sharpe ratio
            return -999.0

        # Report additional metrics as user attributes
        trial.set_user_attr("training_loss", result["training_loss"])
        trial.set_user_attr("duration_seconds", result["training_duration_seconds"])

        for metric_name, metric_value in result["validation_metrics"].items():
            trial.set_user_attr(f"val_{metric_name}", metric_value)

        # Return Sharpe ratio (optimization objective)
        sharpe_ratio = result["sharpe_ratio"]
        logger.info(f"Trial {trial.number} completed: Sharpe ratio = {sharpe_ratio:.4f}")

        return sharpe_ratio

    def run_optimization(self):
        """Execute hyperparameter optimization."""
        global shutdown_requested

        logger.info(f"Starting optimization: {self.num_trials} trials, sequential execution (n_jobs=1)")

        # Connect to gRPC service
        self.grpc_client.connect()

        # Create study
        self.create_study()

        try:
            # Run optimization with sequential trials (n_jobs=1 for GPU safety)
            self.study.optimize(
                self.objective,
                n_trials=self.num_trials,
                n_jobs=1,  # CRITICAL: Sequential execution for 4GB VRAM constraint
                catch=(Exception,),  # Continue on trial failures
                show_progress_bar=True
            )

            if shutdown_requested:
                logger.info("Optimization stopped by shutdown signal")
            else:
                logger.info("Optimization completed successfully")

            # Report best results
            best_trial = self.study.best_trial
            logger.info(f"Best trial: {best_trial.number}")
            logger.info(f"Best Sharpe ratio: {best_trial.value:.4f}")
            logger.info(f"Best hyperparameters: {best_trial.params}")

            # Print summary statistics
            completed_trials = [t for t in self.study.trials if t.state == optuna.trial.TrialState.COMPLETE]
            pruned_trials = [t for t in self.study.trials if t.state == optuna.trial.TrialState.PRUNED]
            failed_trials = [t for t in self.study.trials if t.state == optuna.trial.TrialState.FAIL]

            logger.info(
                f"Trial summary: {len(completed_trials)} completed, "
                f"{len(pruned_trials)} pruned, {len(failed_trials)} failed"
            )

        except KeyboardInterrupt:
            logger.info("Optimization interrupted by user")
        except Exception as e:
            logger.error(f"Optimization failed: {e}", exc_info=True)
            raise
        finally:
            # Cleanup
            self.grpc_client.close()

            # Final GPU cleanup check
            if self.use_gpu:
                gpu_mem = self.gpu_monitor.get_memory_usage()
                logger.info(
                    f"Final GPU memory: {gpu_mem['used_gb']:.2f}/{gpu_mem['total_gb']:.2f} GB"
                )


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Optuna Hyperparameter Tuner for Foxhunt ML Models"
    )

    parser.add_argument(
        "--job-id",
        type=str,
        required=True,
        help="Tuning job identifier"
    )

    parser.add_argument(
        "--model-type",
        type=str,
        required=True,
        choices=["TLOB", "MAMBA_2", "DQN", "PPO", "LIQUID", "TFT"],
        help="Model type to optimize"
    )

    parser.add_argument(
        "--num-trials",
        type=int,
        required=True,
        help="Number of optimization trials to run"
    )

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

    parser.add_argument(
        "--data-source-json",
        type=str,
        required=True,
        help="Data source configuration as JSON string"
    )

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

    parser.add_argument(
        "--storage-path",
        type=str,
        required=True,
        help="Path to Optuna JournalStorage file (for crash recovery)"
    )

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

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

    args = parser.parse_args()

    # Parse data source JSON
    try:
        data_source = json.loads(args.data_source_json)
    except json.JSONDecodeError as e:
        logger.error(f"Invalid data source JSON: {e}")
        sys.exit(1)

    # Create tuner
    tuner = HyperparameterTuner(
        job_id=args.job_id,
        model_type=args.model_type,
        num_trials=args.num_trials,
        config_path=args.config,
        data_source=data_source,
        use_gpu=args.use_gpu,
        storage_path=args.storage_path,
        grpc_host=args.grpc_host,
        grpc_port=args.grpc_port
    )

    # Run optimization
    try:
        tuner.run_optimization()
        logger.info("Tuner subprocess completed successfully")
        sys.exit(0)
    except Exception as e:
        logger.error(f"Tuner subprocess failed: {e}", exc_info=True)
        sys.exit(1)


if __name__ == "__main__":
    main()