//! Knowledge Graph
//!
//! Meta-learning storage, pattern relationships, and success prediction

use crate::{ExecutionResult, ProposalSpec, Result};
use petgraph::Graph;
use petgraph::graph::{NodeIndex, UnGraph};
use std::collections::HashMap;
use uuid::Uuid;

/// Knowledge graph for meta-learning
pub struct KnowledgeGraph {
    graph: UnGraph<Pattern, Relationship>,
    pattern_index: HashMap<String, NodeIndex>,
    success_history: Vec<SuccessRecord>,
    ready: bool,
}

/// Pattern stored in knowledge graph
#[derive(Debug, Clone)]
pub struct Pattern {
    pub id: String,
    pub name: String,
    pub pattern_type: PatternType,
    pub success_rate: f64,
    pub average_improvement: f64,
    pub usage_count: u64,
    pub confidence: f64,
}

/// Type of pattern in the knowledge graph
#[derive(Debug, Clone, PartialEq)]
pub enum PatternType {
    OptimizationTechnique,  // General optimization approach
    ParameterSetting,       // Specific parameter configurations
    AlgorithmChoice,        // Algorithm selection patterns
    ResourceUtilization,    // Resource usage patterns
    PerformanceCorrelation, // Performance correlation patterns
}

/// Relationship between patterns
#[derive(Debug, Clone)]
pub struct Relationship {
    pub kind: RelationshipKind,
    pub strength: f64,
    pub evidence_count: u64,
}

/// Types of relationships between patterns
#[derive(Debug, Clone, PartialEq)]
pub enum RelationshipKind {
    Synergy,      // Patterns work well together
    Conflict,     // Patterns interfere with each other
    Prerequisite, // One pattern requires another
    Alternative,  // Patterns are alternatives to each other
    Correlation,  // Patterns often co-occur
}

/// Record of successful optimizations
#[derive(Debug, Clone)]
struct SuccessRecord {
    proposal_id: Uuid,
    patterns_used: Vec<String>,
    performance_improvement: f64,
    timestamp: std::time::SystemTime,
}

impl KnowledgeGraph {
    /// Create new knowledge graph
    pub fn new() -> Self {
        let mut kg = Self {
            graph: Graph::new_undirected(),
            pattern_index: HashMap::new(),
            success_history: Vec::new(),
            ready: true,
        };

        // Initialize with some basic patterns
        kg.initialize_base_patterns();

        kg
    }

    /// Check if knowledge graph is ready
    pub fn is_ready(&self) -> bool {
        self.ready
    }

    /// Update knowledge graph from execution results
    pub async fn update_from_execution(
        &mut self,
        proposal: &ProposalSpec,
        result: &ExecutionResult,
    ) -> Result<()> {
        if result.safety_check_passed && result.performance_delta > 0.0 {
            // Extract patterns from successful proposal
            let patterns = self.extract_patterns_from_proposal(proposal).await?;

            // Update pattern success rates
            for pattern_name in &patterns {
                self.update_pattern_success(pattern_name, result.performance_delta)
                    .await?;
            }

            // Learn relationships between patterns
            self.learn_pattern_relationships(&patterns, result.performance_delta)
                .await?;

            // Record successful execution
            self.record_success(proposal, result, patterns).await?;
        } else {
            // Learn from failures too
            let patterns = self.extract_patterns_from_proposal(proposal).await?;
            for pattern_name in &patterns {
                self.update_pattern_failure(pattern_name).await?;
            }
        }

        Ok(())
    }

    /// Initialize knowledge graph with base patterns
    fn initialize_base_patterns(&mut self) {
        let base_patterns = vec![
            Pattern {
                id: "kernel_tile_optimization".to_string(),
                name: "Kernel Tile Size Optimization".to_string(),
                pattern_type: PatternType::ParameterSetting,
                success_rate: 0.7,
                average_improvement: 0.12,
                usage_count: 0,
                confidence: 0.8,
            },
            Pattern {
                id: "memory_layout_optimization".to_string(),
                name: "Memory Layout Optimization".to_string(),
                pattern_type: PatternType::OptimizationTechnique,
                success_rate: 0.6,
                average_improvement: 0.08,
                usage_count: 0,
                confidence: 0.7,
            },
            Pattern {
                id: "compiler_aggressive_opts".to_string(),
                name: "Aggressive Compiler Optimizations".to_string(),
                pattern_type: PatternType::AlgorithmChoice,
                success_rate: 0.5,
                average_improvement: 0.05,
                usage_count: 0,
                confidence: 0.6,
            },
            Pattern {
                id: "algorithm_switching".to_string(),
                name: "Algorithm Switching".to_string(),
                pattern_type: PatternType::AlgorithmChoice,
                success_rate: 0.8,
                average_improvement: 0.15,
                usage_count: 0,
                confidence: 0.85,
            },
            // AI-powered optimization patterns
            Pattern {
                id: "ai_powered_optimization".to_string(),
                name: "AI-Powered Code Optimization".to_string(),
                pattern_type: PatternType::OptimizationTechnique,
                success_rate: 0.85,
                average_improvement: 0.22,
                usage_count: 0,
                confidence: 0.90,
            },
            // RTX 5090 hardware-specific patterns
            Pattern {
                id: "hardware_specific_optimization".to_string(),
                name: "Hardware-Specific RTX 5090 Optimization".to_string(),
                pattern_type: PatternType::OptimizationTechnique,
                success_rate: 0.90,
                average_improvement: 0.28,
                usage_count: 0,
                confidence: 0.95,
            },
            // CUDA 13.0 feature enablement patterns
            Pattern {
                id: "cuda_feature_enablement".to_string(),
                name: "CUDA 13.0 Feature Enablement".to_string(),
                pattern_type: PatternType::OptimizationTechnique,
                success_rate: 0.82,
                average_improvement: 0.20,
                usage_count: 0,
                confidence: 0.88,
            },
        ];

        for pattern in base_patterns {
            let node_index = self.graph.add_node(pattern.clone());
            self.pattern_index.insert(pattern.id, node_index);
        }

        // Add some initial relationships
        self.add_initial_relationships();
    }

    /// Add initial known relationships between patterns
    fn add_initial_relationships(&mut self) {
        // Kernel tile optimization and memory layout often work well together (synergy)
        if let (Some(&tile_idx), Some(&memory_idx)) = (
            self.pattern_index.get("kernel_tile_optimization"),
            self.pattern_index.get("memory_layout_optimization"),
        ) {
            self.graph.add_edge(
                tile_idx,
                memory_idx,
                Relationship {
                    kind: RelationshipKind::Synergy,
                    strength: 0.7,
                    evidence_count: 5,
                },
            );
        }

        // Aggressive compiler opts and algorithm switching can conflict
        if let (Some(&compiler_idx), Some(&algorithm_idx)) = (
            self.pattern_index.get("compiler_aggressive_opts"),
            self.pattern_index.get("algorithm_switching"),
        ) {
            self.graph.add_edge(
                compiler_idx,
                algorithm_idx,
                Relationship {
                    kind: RelationshipKind::Conflict,
                    strength: 0.4,
                    evidence_count: 3,
                },
            );
        }
    }

    /// Extract patterns from a proposal
    async fn extract_patterns_from_proposal(&self, proposal: &ProposalSpec) -> Result<Vec<String>> {
        let mut patterns = Vec::new();

        for change in &proposal.changes {
            match change {
                crate::Change::KernelParameter { .. } => {
                    patterns.push("kernel_tile_optimization".to_string());
                }
                crate::Change::MemoryLayout { .. } => {
                    patterns.push("memory_layout_optimization".to_string());
                }
                crate::Change::CompilerFlag { .. } => {
                    patterns.push("compiler_aggressive_opts".to_string());
                }
                crate::Change::AlgorithmSwitch { .. } => {
                    patterns.push("algorithm_switching".to_string());
                }
                crate::Change::CodeOptimization {
                    optimization_type, ..
                } => {
                    patterns.push(format!("ai_code_optimization_{}", optimization_type));
                    patterns.push("ai_powered_optimization".to_string());
                }
                crate::Change::Rtx5090Optimization {
                    optimization_type, ..
                } => {
                    patterns.push(format!("rtx5090_optimization_{}", optimization_type));
                    patterns.push("hardware_specific_optimization".to_string());
                }
                crate::Change::Cuda13Feature { feature_name, .. } => {
                    patterns.push(format!("cuda13_feature_{}", feature_name));
                    patterns.push("cuda_feature_enablement".to_string());
                }
            }
        }

        Ok(patterns)
    }

    /// Update pattern success statistics
    async fn update_pattern_success(&mut self, pattern_name: &str, improvement: f64) -> Result<()> {
        if let Some(&node_idx) = self.pattern_index.get(pattern_name) {
            if let Some(pattern) = self.graph.node_weight_mut(node_idx) {
                // Update success rate using exponential moving average
                let alpha = 0.1; // Learning rate
                pattern.success_rate = (1.0 - alpha) * pattern.success_rate + alpha * 1.0;

                // Update average improvement
                pattern.average_improvement =
                    (1.0 - alpha) * pattern.average_improvement + alpha * improvement;

                // Increment usage count
                pattern.usage_count += 1;

                // Update confidence based on usage count
                pattern.confidence =
                    (pattern.usage_count as f64 / (pattern.usage_count as f64 + 10.0)).min(0.95);
            }
        }

        Ok(())
    }

    /// Update pattern failure statistics
    async fn update_pattern_failure(&mut self, pattern_name: &str) -> Result<()> {
        if let Some(&node_idx) = self.pattern_index.get(pattern_name) {
            if let Some(pattern) = self.graph.node_weight_mut(node_idx) {
                // Update success rate with failure
                let alpha = 0.1;
                pattern.success_rate = (1.0 - alpha) * pattern.success_rate + alpha * 0.0;

                // Increment usage count
                pattern.usage_count += 1;

                // Adjust confidence down slightly for failures
                pattern.confidence *= 0.98;
            }
        }

        Ok(())
    }

    /// Learn relationships between patterns based on co-occurrence
    async fn learn_pattern_relationships(
        &mut self,
        patterns: &[String],
        improvement: f64,
    ) -> Result<()> {
        // Learn relationships between all pairs of patterns used together
        for i in 0..patterns.len() {
            for j in (i + 1)..patterns.len() {
                let pattern_a = &patterns[i];
                let pattern_b = &patterns[j];

                if let (Some(&idx_a), Some(&idx_b)) = (
                    self.pattern_index.get(pattern_a),
                    self.pattern_index.get(pattern_b),
                ) {
                    self.update_relationship(idx_a, idx_b, improvement).await?;
                }
            }
        }

        Ok(())
    }

    /// Update relationship strength between two patterns
    async fn update_relationship(
        &mut self,
        idx_a: NodeIndex,
        idx_b: NodeIndex,
        improvement: f64,
    ) -> Result<()> {
        // Find existing edge or create new one
        let edge_idx = self.graph.find_edge(idx_a, idx_b);

        match edge_idx {
            Some(edge_idx) => {
                // Update existing relationship
                if let Some(relationship) = self.graph.edge_weight_mut(edge_idx) {
                    relationship.evidence_count += 1;

                    // Determine relationship kind based on improvement
                    let relationship_kind = if improvement > 0.1 {
                        RelationshipKind::Synergy
                    } else if improvement < 0.0 {
                        RelationshipKind::Conflict
                    } else {
                        RelationshipKind::Correlation
                    };

                    // Update strength using exponential moving average
                    let alpha = 0.2;
                    let new_strength = improvement.abs();
                    relationship.strength =
                        (1.0 - alpha) * relationship.strength + alpha * new_strength;

                    // Update relationship kind if evidence is strong
                    if relationship.evidence_count >= 3 {
                        relationship.kind = relationship_kind;
                    }
                }
            }
            None => {
                // Create new relationship
                let relationship_kind = if improvement > 0.1 {
                    RelationshipKind::Synergy
                } else if improvement < 0.0 {
                    RelationshipKind::Conflict
                } else {
                    RelationshipKind::Correlation
                };

                let new_relationship = Relationship {
                    kind: relationship_kind,
                    strength: improvement.abs(),
                    evidence_count: 1,
                };

                self.graph.add_edge(idx_a, idx_b, new_relationship);
            }
        }

        Ok(())
    }

    /// Record successful execution in history
    async fn record_success(
        &mut self,
        proposal: &ProposalSpec,
        result: &ExecutionResult,
        patterns: Vec<String>,
    ) -> Result<()> {
        let success_record = SuccessRecord {
            proposal_id: proposal.id,
            patterns_used: patterns,
            performance_improvement: result.performance_delta,
            timestamp: std::time::SystemTime::now(),
        };

        // In a real implementation, this would be mutable
        // For now, we simulate recording the success
        tracing::info!(
            "Recording success: proposal {} with {:.2}% improvement using {} patterns",
            proposal.id,
            result.performance_delta * 100.0,
            success_record.patterns_used.len()
        );

        Ok(())
    }

    /// Get recommendations for optimization patterns
    pub async fn get_recommendations(
        &self,
        context: &OptimizationContext,
    ) -> Result<Vec<PatternRecommendation>> {
        let mut recommendations = Vec::new();

        // Analyze all patterns and rank by potential success
        for (pattern_id, &node_idx) in &self.pattern_index {
            if let Some(pattern) = self.graph.node_weight(node_idx) {
                let score = self
                    .calculate_recommendation_score(pattern, context)
                    .await?;

                if score > 0.3 {
                    // Threshold for recommendations
                    recommendations.push(PatternRecommendation {
                        pattern_id: pattern_id.clone(),
                        pattern_name: pattern.name.clone(),
                        confidence: pattern.confidence,
                        expected_improvement: pattern.average_improvement,
                        score,
                        rationale: self.generate_rationale(pattern, context).await?,
                    });
                }
            }
        }

        // Sort by score (highest first)
        recommendations.sort_by(|a, b| b.score.total_cmp(&a.score));

        // Limit to top 5 recommendations
        recommendations.truncate(5);

        Ok(recommendations)
    }

    /// Calculate recommendation score for a pattern given context
    async fn calculate_recommendation_score(
        &self,
        pattern: &Pattern,
        context: &OptimizationContext,
    ) -> Result<f64> {
        let base_score = pattern.success_rate * pattern.confidence;

        // Adjust based on context
        let context_multiplier = match (&pattern.pattern_type, &context.optimization_target) {
            (PatternType::ParameterSetting, OptimizationTarget::Performance) => 1.2,
            (PatternType::AlgorithmChoice, OptimizationTarget::Performance) => 1.5,
            (PatternType::OptimizationTechnique, OptimizationTarget::Memory) => 1.3,
            (PatternType::ResourceUtilization, OptimizationTarget::Power) => 1.4,
            _ => 1.0,
        };

        let final_score = base_score * context_multiplier;
        Ok(final_score.max(0.0).min(1.0))
    }

    /// Generate rationale for why a pattern is recommended
    async fn generate_rationale(
        &self,
        pattern: &Pattern,
        context: &OptimizationContext,
    ) -> Result<String> {
        let success_percentage = (pattern.success_rate * 100.0) as i32;
        let avg_improvement_percentage = (pattern.average_improvement * 100.0) as i32;

        let rationale = format!(
            "Pattern '{}' has {:.0}% success rate with average {:.0}% improvement. Used successfully {} times. Confidence: {:.0}%",
            pattern.name,
            success_percentage,
            avg_improvement_percentage,
            pattern.usage_count,
            pattern.confidence * 100.0
        );

        Ok(rationale)
    }
}

/// Context for optimization recommendations
#[derive(Debug, Clone)]
pub struct OptimizationContext {
    pub optimization_target: OptimizationTarget,
    pub current_metrics: HashMap<String, f64>,
    pub constraints: Vec<String>,
}

/// Target for optimization
#[derive(Debug, Clone, PartialEq)]
pub enum OptimizationTarget {
    Performance,
    Memory,
    Power,
    Latency,
}

/// Pattern recommendation from knowledge graph
#[derive(Debug, Clone)]
pub struct PatternRecommendation {
    pub pattern_id: String,
    pub pattern_name: String,
    pub confidence: f64,
    pub expected_improvement: f64,
    pub score: f64,
    pub rationale: String,
}