theory_core 0.1.0

//! Speculate-Verify Routing Logic
//!
//! Decides when to use fast neural surrogate speculation vs
//! slower but accurate SMT solver verification.

use crate::ir::ConstraintIR;
use crate::{ReasoningStrategy, SolutionResult};
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Router for speculation vs verification decisions
pub struct SpeculateVerifyRouter {
    config: RouterConfig,
    performance_history: PerformanceHistory,
}

/// Configuration for routing decisions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterConfig {
    /// Confidence threshold for using speculation only
    pub speculation_threshold: f32,
    /// Confidence threshold for requiring verification
    pub verification_threshold: f32,
    /// Time budget in milliseconds
    pub time_budget_ms: f64,
    /// Accuracy requirements (0.0 to 1.0)
    pub required_accuracy: f32,
    /// Enable adaptive thresholds based on performance
    pub adaptive_thresholds: bool,
}

impl Default for RouterConfig {
    fn default() -> Self {
        Self {
            speculation_threshold: 0.9,
            verification_threshold: 0.5,
            time_budget_ms: 1000.0,
            required_accuracy: 0.95,
            adaptive_thresholds: true,
        }
    }
}

/// Performance history for adaptive routing
#[derive(Debug, Default)]
struct PerformanceHistory {
    /// Track speculation accuracy by problem characteristics
    speculation_accuracy: HashMap<ProblemSignature, AccuracyStats>,
    /// Track solver performance by problem characteristics
    solver_performance: HashMap<ProblemSignature, PerformanceStats>,
}

/// Signature of problem characteristics for routing decisions
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct ProblemSignature {
    num_variables: usize,
    num_constraints: usize,
    theory_tags: Vec<String>, // Simplified theory tag representation
    complexity_class: ComplexityClass,
}

/// Problem complexity classification
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
enum ComplexityClass {
    Simple,  // Linear constraints, few variables
    Medium,  // Moderate non-linearity or many variables
    Complex, // High non-linearity or very many variables
    Unknown, // Not yet classified
}

/// Accuracy statistics for speculation
#[derive(Debug, Default)]
struct AccuracyStats {
    correct_predictions: u64,
    total_predictions: u64,
    average_confidence: f32,
}

/// Performance statistics for solvers
#[derive(Debug, Default)]
struct PerformanceStats {
    average_solve_time_ms: f64,
    success_rate: f32,
    timeout_rate: f32,
    total_attempts: u64,
    total_timeouts: u64,
}

impl PerformanceStats {
    /// Update statistics with a new solve attempt
    pub fn record_attempt(&mut self, solve_time_ms: f64, timed_out: bool) {
        self.total_attempts += 1;
        if timed_out {
            self.total_timeouts += 1;
        }

        // Update timeout rate
        self.timeout_rate = self.total_timeouts as f32 / self.total_attempts as f32;

        // Update success rate (inverse of timeout rate as approximation)
        self.success_rate = 1.0 - self.timeout_rate;

        // Update average solve time using exponential moving average
        let alpha = 0.2; // Learning rate
        if self.total_attempts == 1 {
            self.average_solve_time_ms = solve_time_ms;
        } else {
            self.average_solve_time_ms =
                (1.0 - alpha) * self.average_solve_time_ms + alpha * solve_time_ms;
        }
    }

    /// Get analytics summary
    pub fn analytics_summary(&self) -> AnalyticsSummary {
        AnalyticsSummary {
            avg_solve_time_ms: self.average_solve_time_ms,
            success_rate: self.success_rate,
            timeout_rate: self.timeout_rate,
            total_attempts: self.total_attempts,
            reliability_score: self.calculate_reliability(),
        }
    }

    /// Calculate reliability score (0.0 to 1.0)
    fn calculate_reliability(&self) -> f32 {
        if self.total_attempts == 0 {
            return 0.5; // Unknown
        }

        // Reliability = success_rate * sample_confidence
        let sample_confidence = (self.total_attempts as f32 / 100.0).min(1.0);
        self.success_rate * sample_confidence
    }
}

/// Analytics summary for reporting
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct AnalyticsSummary {
    pub avg_solve_time_ms: f64,
    pub success_rate: f32,
    pub timeout_rate: f32,
    pub total_attempts: u64,
    pub reliability_score: f32,
}

impl SpeculateVerifyRouter {
    pub fn new() -> Self {
        Self {
            config: RouterConfig::default(),
            performance_history: PerformanceHistory::default(),
        }
    }

    pub fn new_with_config(config: RouterConfig) -> Self {
        Self {
            config,
            performance_history: PerformanceHistory::default(),
        }
    }

    /// Route constraint IR to appropriate reasoning strategy
    pub fn route(&self, constraints: &ConstraintIR) -> Result<ReasoningStrategy> {
        let signature = self.compute_problem_signature(constraints);

        // Get estimated performance metrics
        let speculation_confidence = self.estimate_speculation_confidence(&signature);
        let estimated_solver_time = self.estimate_solver_time(&signature);

        // Make routing decision based on multiple factors
        let strategy = if speculation_confidence >= self.config.speculation_threshold {
            // High confidence in speculation
            ReasoningStrategy::Speculate
        } else if speculation_confidence < self.config.verification_threshold
            || estimated_solver_time > self.config.time_budget_ms
        {
            // Low confidence or solver too slow - still try speculation with fallback
            ReasoningStrategy::Speculate
        } else if estimated_solver_time < self.config.time_budget_ms * 0.1 {
            // Solver is very fast - just use it directly
            ReasoningStrategy::Verify
        } else {
            // Use both for comparison and learning
            ReasoningStrategy::Hybrid
        };

        Ok(strategy)
    }

    /// Update routing decisions based on actual performance
    pub fn update_performance(
        &mut self,
        constraints: &ConstraintIR,
        strategy_used: ReasoningStrategy,
        speculation_result: Option<&SolutionResult>,
        verification_result: Option<&SolutionResult>,
        actual_solver_time_ms: f64,
    ) {
        let signature = self.compute_problem_signature(constraints);

        // Update speculation accuracy if we have both results
        if let (Some(spec_result), Some(verify_result)) = (speculation_result, verification_result)
        {
            let correct = self.solutions_match(spec_result, verify_result);
            self.update_speculation_accuracy(&signature, correct, spec_result.confidence);
        }

        // Update solver performance
        if matches!(
            strategy_used,
            ReasoningStrategy::Verify | ReasoningStrategy::Hybrid
        ) {
            self.update_solver_performance(&signature, actual_solver_time_ms, true);
        }
    }

    /// Compute problem signature for routing decisions
    fn compute_problem_signature(&self, constraints: &ConstraintIR) -> ProblemSignature {
        let num_variables = constraints.variables.len();
        let num_constraints = constraints.constraints.len();

        let theory_tags = constraints
            .theory_tags
            .iter()
            .map(|tag| format!("{:?}", tag))
            .collect();

        let complexity_class = self.classify_complexity(constraints);

        ProblemSignature {
            num_variables,
            num_constraints,
            theory_tags,
            complexity_class,
        }
    }

    /// Classify problem complexity
    fn classify_complexity(&self, constraints: &ConstraintIR) -> ComplexityClass {
        let num_vars = constraints.variables.len();
        let num_constraints = constraints.constraints.len();

        // Simple heuristic classification
        if num_vars <= 10 && num_constraints <= 10 {
            ComplexityClass::Simple
        } else if num_vars <= 50 && num_constraints <= 50 {
            ComplexityClass::Medium
        } else if num_vars <= 200 && num_constraints <= 200 {
            ComplexityClass::Complex
        } else {
            ComplexityClass::Complex
        }
    }

    /// Estimate confidence in speculation for this problem type
    fn estimate_speculation_confidence(&self, signature: &ProblemSignature) -> f32 {
        if let Some(stats) = self.performance_history.speculation_accuracy.get(signature) {
            if stats.total_predictions > 0 {
                return stats.correct_predictions as f32 / stats.total_predictions as f32;
            }
        }

        // Default confidence based on complexity
        match signature.complexity_class {
            ComplexityClass::Simple => 0.9,
            ComplexityClass::Medium => 0.7,
            ComplexityClass::Complex => 0.5,
            ComplexityClass::Unknown => 0.6,
        }
    }

    /// Estimate solver time for this problem type
    fn estimate_solver_time(&self, signature: &ProblemSignature) -> f64 {
        if let Some(stats) = self.performance_history.solver_performance.get(signature) {
            return stats.average_solve_time_ms;
        }

        // Default estimates based on problem size and complexity
        let base_time = match signature.complexity_class {
            ComplexityClass::Simple => 10.0,
            ComplexityClass::Medium => 100.0,
            ComplexityClass::Complex => 1000.0,
            ComplexityClass::Unknown => 500.0,
        };

        // Scale by problem size
        let size_factor = (signature.num_variables + signature.num_constraints) as f64;
        base_time + size_factor * 0.1
    }

    /// Check if two solutions match (for accuracy measurement)
    fn solutions_match(
        &self,
        spec_result: &SolutionResult,
        verify_result: &SolutionResult,
    ) -> bool {
        // Simple comparison - both satisfiable/unsatisfiable
        spec_result.satisfiable == verify_result.satisfiable
        // TODO: Compare actual assignments if both satisfiable
    }

    /// Update speculation accuracy statistics
    fn update_speculation_accuracy(
        &mut self,
        signature: &ProblemSignature,
        correct: bool,
        confidence: f32,
    ) {
        let stats = self
            .performance_history
            .speculation_accuracy
            .entry(signature.clone())
            .or_default();

        if correct {
            stats.correct_predictions += 1;
        }
        stats.total_predictions += 1;

        // Update rolling average confidence
        let alpha = 0.1_f32; // Learning rate
        stats.average_confidence =
            stats.average_confidence * (1.0_f32 - alpha) + confidence * alpha;
    }

    /// Update solver performance statistics
    fn update_solver_performance(
        &mut self,
        signature: &ProblemSignature,
        solve_time_ms: f64,
        success: bool,
    ) {
        let stats = self
            .performance_history
            .solver_performance
            .entry(signature.clone())
            .or_default();

        // Update rolling average solve time
        let alpha = 0.1_f64; // Learning rate
        stats.average_solve_time_ms =
            stats.average_solve_time_ms * (1.0_f64 - alpha) + solve_time_ms * alpha;

        // Update success rate
        let alpha_f32 = alpha as f32;
        stats.success_rate = stats.success_rate * (1.0_f32 - alpha_f32)
            + if success { 1.0_f32 } else { 0.0_f32 } * alpha_f32;
    }
}

impl Default for SpeculateVerifyRouter {
    fn default() -> Self {
        Self::new()
    }
}