"""
Adaptive surrogate modeling with confidence-based exploration-exploitation transitions.

This module implements state-of-the-art Bayesian optimization strategies that
dynamically adjust exploration vs exploitation based on surrogate model confidence.
"""

import numpy as np
from typing import Optional, Dict, List
import optuna
from scipy.stats import variation
import json
from pathlib import Path


class SurrogateConfidenceMetrics:
    """Calculate confidence metrics for surrogate model quality.

    STUDY-AWARE: Uses study.trials directly instead of session-based history
    to properly track confidence across multiple optimization runs.
    """

    def __init__(self, min_trials_for_confidence: int = 15):
        self.min_trials = min_trials_for_confidence

    def update(self, study: optuna.Study, trial: optuna.trial.FrozenTrial):
        """Update metrics after each trial (no-op for study-aware implementation)."""
        pass  # Study-aware: we read directly from study.trials

    def calculate_confidence(self, study: optuna.Study) -> Dict[str, float]:
        """
        Calculate comprehensive surrogate confidence metrics.

        STUDY-AWARE: Uses ALL completed trials from the study database,
        not just trials from the current session.

        PROTOCOL 11: Multi-objective studies are NOT supported by adaptive
        characterization. Return immediately with max confidence to skip
        characterization phase.

        Returns:
            Dict with confidence scores:
            - 'overall_confidence': 0-1 score, where 1 = high confidence
            - 'convergence_score': How stable recent improvements are
            - 'exploration_coverage': How well parameter space is covered
            - 'prediction_stability': How consistent the model predictions are
        """
        # [Protocol 11] Multi-objective NOT supported by adaptive characterization
        is_multi_objective = len(study.directions) > 1
        if is_multi_objective:
            return {
                'overall_confidence': 1.0,  # Skip characterization
                'convergence_score': 1.0,
                'exploration_coverage': 1.0,
                'prediction_stability': 1.0,
                'ready_for_exploitation': True,  # Go straight to NSGA-II
                'total_trials': len(study.trials),
                'message': '[Protocol 11] Multi-objective: skipping adaptive characterization, using NSGA-II directly'
            }

        # Get ALL completed trials from study (study-aware)
        all_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]

        if len(all_trials) < self.min_trials:
            return {
                'overall_confidence': 0.0,
                'convergence_score': 0.0,
                'exploration_coverage': 0.0,
                'prediction_stability': 0.0,
                'ready_for_exploitation': False,
                'total_trials': len(all_trials),
                'message': f'Need {self.min_trials - len(all_trials)} more trials for confidence assessment (currently {len(all_trials)} trials)'
            }

        # 1. Convergence Score - are we finding consistent improvements?
        recent_window = 10
        recent_trials = all_trials[-recent_window:]
        recent_values = [t.value for t in recent_trials]  # Safe: single-objective only

        # Calculate improvement rate
        improvements = []
        for i in range(1, len(recent_values)):
            if recent_values[i] < recent_values[i-1]:
                improvement = (recent_values[i-1] - recent_values[i]) / abs(recent_values[i-1])
                improvements.append(improvement)

        # If we're consistently finding improvements, convergence is good
        if improvements:
            avg_improvement = np.mean(improvements)
            improvement_consistency = 1.0 - variation(improvements) if len(improvements) > 1 else 0.5
            convergence_score = min(1.0, avg_improvement * improvement_consistency * 10)
        else:
            convergence_score = 0.0

        # 2. Exploration Coverage - how well have we covered parameter space?
        # Use coefficient of variation for each parameter
        param_names = list(all_trials[0].params.keys()) if all_trials else []

        coverage_scores = []
        for param in param_names:
            values = [t.params[param] for t in all_trials]

            # Get parameter bounds
            distribution = all_trials[0].distributions[param]
            param_range = distribution.high - distribution.low

            # Calculate spread relative to bounds
            spread = max(values) - min(values)
            coverage = spread / param_range
            coverage_scores.append(coverage)

        exploration_coverage = np.mean(coverage_scores)

        # 3. Prediction Stability - are recent trials clustered in good regions?
        recent_best_values = []
        current_best = float('inf')
        for trial in recent_trials:
            current_best = min(current_best, trial.value)
            recent_best_values.append(current_best)

        # If best hasn't improved much recently, model is stable
        if len(recent_best_values) > 1:
            best_stability = 1.0 - (recent_best_values[0] - recent_best_values[-1]) / (recent_best_values[0] + 1e-10)
            prediction_stability = max(0.0, min(1.0, best_stability))
        else:
            prediction_stability = 0.0

        # 4. Overall Confidence - weighted combination
        overall_confidence = (
            0.4 * convergence_score +
            0.3 * exploration_coverage +
            0.3 * prediction_stability
        )

        # Decision: Ready for intensive exploitation?
        ready_for_exploitation = (
            overall_confidence >= 0.65 and
            exploration_coverage >= 0.5 and
            len(all_trials) >= self.min_trials
        )

        message = self._get_confidence_message(overall_confidence, ready_for_exploitation)

        return {
            'overall_confidence': overall_confidence,
            'convergence_score': convergence_score,
            'exploration_coverage': exploration_coverage,
            'prediction_stability': prediction_stability,
            'ready_for_exploitation': ready_for_exploitation,
            'total_trials': len(all_trials),
            'message': message
        }

    def _get_confidence_message(self, confidence: float, ready: bool) -> str:
        """Generate human-readable confidence assessment."""
        if ready:
            return f"HIGH CONFIDENCE ({confidence:.1%}) - Transitioning to exploitation phase"
        elif confidence >= 0.5:
            return f"MEDIUM CONFIDENCE ({confidence:.1%}) - Continue exploration with some exploitation"
        elif confidence >= 0.3:
            return f"LOW CONFIDENCE ({confidence:.1%}) - Focus on exploration"
        else:
            return f"VERY LOW CONFIDENCE ({confidence:.1%}) - Need more diverse exploration"


class AdaptiveExploitationCallback:
    """
    Dynamically adjust sampler behavior based on surrogate confidence.

    This callback monitors surrogate model confidence and adapts the optimization
    strategy from exploration-heavy to exploitation-heavy as confidence increases.

    STUDY-AWARE: Tracks phase transitions across multiple optimization runs
    and persists confidence history to JSON files.
    """

    def __init__(
        self,
        target_value: Optional[float] = None,
        tolerance: float = 0.1,
        min_confidence_for_exploitation: float = 0.65,
        min_trials: int = 15,
        verbose: bool = True,
        tracking_dir: Optional[Path] = None
    ):
        """
        Args:
            target_value: Target objective value (if known)
            tolerance: Acceptable error from target
            min_confidence_for_exploitation: Confidence threshold to enable intensive exploitation
            min_trials: Minimum trials before confidence assessment
            verbose: Print confidence updates
            tracking_dir: Directory to save phase transition tracking files
        """
        self.target_value = target_value
        self.tolerance = tolerance
        self.min_confidence = min_confidence_for_exploitation
        self.verbose = verbose
        self.tracking_dir = tracking_dir

        self.metrics = SurrogateConfidenceMetrics(min_trials_for_confidence=min_trials)
        self.consecutive_successes = 0

        # Initialize phase transition tracking
        self.phase_transition_file = None
        self.confidence_history_file = None
        if tracking_dir:
            self.tracking_dir = Path(tracking_dir)
            self.tracking_dir.mkdir(parents=True, exist_ok=True)
            self.phase_transition_file = self.tracking_dir / "phase_transitions.json"
            self.confidence_history_file = self.tracking_dir / "confidence_history.json"

        # Load existing phase transition data if available
        self.phase_transitions = self._load_phase_transitions()
        self.confidence_history = self._load_confidence_history()

        # Determine current phase from history
        self.phase = self._get_current_phase()

    def _load_phase_transitions(self) -> List[Dict]:
        """Load existing phase transition history from JSON."""
        if self.phase_transition_file and self.phase_transition_file.exists():
            try:
                with open(self.phase_transition_file, 'r') as f:
                    return json.load(f)
            except Exception:
                return []
        return []

    def _load_confidence_history(self) -> List[Dict]:
        """Load existing confidence history from JSON."""
        if self.confidence_history_file and self.confidence_history_file.exists():
            try:
                with open(self.confidence_history_file, 'r') as f:
                    return json.load(f)
            except Exception:
                return []
        return []

    def _get_current_phase(self) -> str:
        """Determine current phase from transition history."""
        if not self.phase_transitions:
            return "exploration"
        # If any transition to exploitation exists, we're in exploitation
        for transition in self.phase_transitions:
            if transition.get('to_phase') == 'exploitation':
                return "exploitation"
        return "exploration"

    def _save_phase_transition(self, trial_number: int, confidence: Dict):
        """Save phase transition event to JSON."""
        if not self.phase_transition_file:
            return

        transition_event = {
            'trial_number': trial_number,
            'from_phase': 'exploration',
            'to_phase': 'exploitation',
            'confidence_metrics': {
                'overall_confidence': confidence['overall_confidence'],
                'convergence_score': confidence['convergence_score'],
                'exploration_coverage': confidence['exploration_coverage'],
                'prediction_stability': confidence['prediction_stability']
            },
            'total_trials': confidence.get('total_trials', trial_number + 1)
        }

        self.phase_transitions.append(transition_event)

        try:
            with open(self.phase_transition_file, 'w') as f:
                json.dump(self.phase_transitions, f, indent=2)
        except Exception as e:
            if self.verbose:
                print(f"  Warning: Failed to save phase transition: {e}")

    def _save_confidence_snapshot(self, trial_number: int, confidence: Dict):
        """Save confidence metrics snapshot to history."""
        if not self.confidence_history_file:
            return

        snapshot = {
            'trial_number': trial_number,
            'phase': self.phase,
            'confidence_metrics': {
                'overall_confidence': confidence['overall_confidence'],
                'convergence_score': confidence['convergence_score'],
                'exploration_coverage': confidence['exploration_coverage'],
                'prediction_stability': confidence['prediction_stability']
            },
            'total_trials': confidence.get('total_trials', trial_number + 1)
        }

        self.confidence_history.append(snapshot)

        try:
            with open(self.confidence_history_file, 'w') as f:
                json.dump(self.confidence_history, f, indent=2)
        except Exception as e:
            if self.verbose:
                print(f"  Warning: Failed to save confidence history: {e}")

    def __call__(self, study: optuna.Study, trial: optuna.trial.FrozenTrial):
        """Called after each trial completes."""
        # Skip failed trials
        if trial.state != optuna.trial.TrialState.COMPLETE:
            return

        # Update metrics (no-op for study-aware implementation)
        self.metrics.update(study, trial)

        # Calculate confidence
        confidence = self.metrics.calculate_confidence(study)

        # Save confidence snapshot every 5 trials
        if trial.number % 5 == 0:
            self._save_confidence_snapshot(trial.number, confidence)

        # Print confidence report
        if self.verbose and trial.number % 5 == 0:  # Every 5 trials
            self._print_confidence_report(trial.number, confidence)

        # Check for phase transition
        if confidence['ready_for_exploitation'] and self.phase == "exploration":
            self.phase = "exploitation"

            # Save transition event
            self._save_phase_transition(trial.number, confidence)

            if self.verbose:
                print(f"\n{'='*60}")
                print(f"  PHASE TRANSITION: EXPLORATION -> EXPLOITATION")
                print(f"  Trial #{trial.number}")
                print(f"  Surrogate confidence: {confidence['overall_confidence']:.1%}")
                print(f"  Now focusing on refining best regions")
                print(f"{'='*60}\n")

        # Check for target achievement
        if self.target_value is not None and trial.value <= self.tolerance:
            self.consecutive_successes += 1

            if self.verbose:
                print(f"  [TARGET] Trial #{trial.number}: {trial.value:.6f} ≤ {self.tolerance:.6f}")
                print(f"  [TARGET] Consecutive successes: {self.consecutive_successes}/3")

            # Stop after 3 consecutive successes in exploitation phase
            if self.consecutive_successes >= 3 and self.phase == "exploitation":
                if self.verbose:
                    print(f"\n{'='*60}")
                    print(f"  TARGET ACHIEVED WITH HIGH CONFIDENCE")
                    print(f"  Best value: {study.best_value:.6f}")
                    print(f"  Stopping optimization")
                    print(f"{'='*60}\n")
                study.stop()
        else:
            self.consecutive_successes = 0

    def _print_confidence_report(self, trial_number: int, confidence: Dict):
        """Print confidence metrics report."""
        print(f"\n  [CONFIDENCE REPORT - Trial #{trial_number}]")
        print(f"  Phase: {self.phase.upper()}")
        print(f"  Overall Confidence: {confidence['overall_confidence']:.1%}")
        print(f"    - Convergence: {confidence['convergence_score']:.1%}")
        print(f"    - Coverage: {confidence['exploration_coverage']:.1%}")
        print(f"    - Stability: {confidence['prediction_stability']:.1%}")
        print(f"  {confidence['message']}")
        print()


def create_adaptive_sampler(
    n_startup_trials: int = 10,
    multivariate: bool = True,
    confidence_threshold: float = 0.65
) -> optuna.samplers.TPESampler:
    """
    Create TPE sampler configured for adaptive exploration-exploitation.

    Args:
        n_startup_trials: Initial random exploration trials
        multivariate: Enable multivariate TPE for correlated parameters
        confidence_threshold: Confidence needed before intensive exploitation

    Returns:
        Configured TPESampler
    """
    # Higher n_ei_candidates = more exploitation
    # Will be used once confidence threshold is reached
    return optuna.samplers.TPESampler(
        n_startup_trials=n_startup_trials,
        n_ei_candidates=24,
        multivariate=multivariate,
        warn_independent_sampling=True
    )