//! Training system for box embeddings in coreference resolution.
//!
//! This module implements the mathematical framework for learning box embeddings
//! from coreference annotations, based on research from BERE, BoxTE, and UKGE.
//!
//! # Mathematical Foundation
//!
//! ## Objective Function
//!
//! We maximize the conditional probability for positive pairs (entities that corefer)
//! and minimize it for negative pairs (entities that don't corefer):
//!
//! ```
//! L = -Σ log P(box_i | box_j) for positive pairs (i,j)
//!     + λ * Σ max(0, margin - P(box_i | box_j)) for negative pairs (i,j)
//! ```
//!
//! Where:
//! - P(box_i | box_j) = Vol(box_i ∩ box_j) / Vol(box_j) (conditional probability)
//! - Positive pairs: entities in the same coreference chain
//! - Negative pairs: entities in different chains
//! - λ: negative sampling weight
//! - margin: minimum separation for negative pairs
//!
//! ## Gradient Computation
//!
//! We use **analytical gradients** (not finite differences) for efficiency:
//!
//! For a box with min = μ - exp(δ)/2, max = μ + exp(δ)/2:
//!
//! ∂Vol/∂μᵢ = 0 (volume doesn't depend on center position)
//! ∂Vol/∂δᵢ = Vol * 1 (since Vol = ∏exp(δᵢ))
//!
//! For intersection volume, gradients flow through each dimension's overlap.
//!
//! ## Box Constraints
//!
//! We enforce mᵢ ≤ Mᵢ for all dimensions using:
//! - Reparameterization: Mᵢ = mᵢ + exp(δᵢ) where δᵢ is the learned parameter

use crate::backends::box_embeddings::BoxEmbedding;
use crate::eval::coref::{CorefChain, CorefDocument};
use crate::Entity;
use std::collections::HashMap;
use serde::{Deserialize, Serialize};

// =============================================================================
// Trainable Box Embedding
// =============================================================================

/// A trainable box embedding with learnable parameters.
///
/// Uses reparameterization to ensure min ≤ max:
/// - min[i] = μᵢ - exp(δᵢ)/2 (learnable via μᵢ, δᵢ)
/// - max[i] = μᵢ + exp(δᵢ)/2 (learnable via μᵢ, δᵢ)
///
/// This ensures boxes are always valid (min ≤ max).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainableBox {
    /// Mean position in each dimension [d]
    pub mu: Vec<f32>,
    /// Log-width in each dimension [d] (width = exp(delta))
    pub delta: Vec<f32>,
    /// Dimension
    pub dim: usize,
}

impl TrainableBox {
    /// Create a new trainable box.
    ///
    /// # Arguments
    ///
    /// * `mu` - Mean position (center of box)
    /// * `delta` - Log-width (width = exp(delta))
    ///
    /// The box will have:
    /// - min[i] = mu[i] - exp(delta[i]) / 2
    /// - max[i] = mu[i] + exp(delta[i]) / 2
    #[must_use]
    pub fn new(mu: Vec<f32>, delta: Vec<f32>) -> Self {
        assert_eq!(mu.len(), delta.len(), "mu and delta must have same dimension");
        let dim = mu.len();
        Self { mu, delta, dim }
    }

    /// Initialize from a vector embedding.
    ///
    /// Creates a small box around the vector with initial width `init_width`.
    #[must_use]
    pub fn from_vector(vector: &[f32], init_width: f32) -> Self {
        let mu = vector.to_vec();
        let delta: Vec<f32> = vec![init_width.ln(); mu.len()];
        Self::new(mu, delta)
    }

    /// Convert to a BoxEmbedding (for inference).
    #[must_use]
    pub fn to_box(&self) -> BoxEmbedding {
        let min: Vec<f32> = self
            .mu
            .iter()
            .zip(self.delta.iter())
            .map(|(&m, &d)| m - (d.exp() / 2.0))
            .collect();
        let max: Vec<f32> = self
            .mu
            .iter()
            .zip(self.delta.iter())
            .map(|(&m, &d)| m + (d.exp() / 2.0))
            .collect();
        BoxEmbedding::new(min, max)
    }
}

// =============================================================================
// Training Data Structures
// =============================================================================

/// A single training example (one document with coreference chains).
///
/// Each example contains:
/// - Entities: All mentions in the document
/// - Chains: Groups of entities that corefer
#[derive(Debug, Clone)]
pub struct TrainingExample {
    /// All entity mentions in the document
    pub entities: Vec<Entity>,
    /// Coreference chains (groups of entity IDs that refer to the same entity)
    pub chains: Vec<CorefChain>,
}

/// Convert a CorefDocument to a TrainingExample.
///
/// This extracts all mentions as entities and preserves the coreference chains.
impl From<&CorefDocument> for TrainingExample {
    fn from(doc: &CorefDocument) -> Self {
        // Collect all mentions as entities
        let mut entities = Vec::new();
        let mut mention_to_entity_id = HashMap::new();

        for chain in &doc.chains {
            for mention in &chain.mentions {
                // Use character offset as entity ID (unique per document)
                let entity_id = mention.start;
                
                // Convert mention entity_type (Option<String>) to EntityType
                // Default to Person if type is unknown
                let entity_type = mention.entity_type.as_ref()
                    .and_then(|s| {
                        match s.as_str() {
                            "PER" | "Person" | "person" => Some(crate::EntityType::Person),
                            "ORG" | "Organization" | "organization" => Some(crate::EntityType::Organization),
                            "LOC" | "Location" | "location" => Some(crate::EntityType::Location),
                            _ => None,
                        }
                    })
                    .unwrap_or(crate::EntityType::Person);
                
                // Create Entity from Mention
                let entity = Entity::new(
                    mention.text.clone(),
                    entity_type,
                    entity_id,
                    mention.end,
                    1.0,
                );
                
                entities.push(entity);
                mention_to_entity_id.insert((mention.start, mention.end), entity_id);
            }
        }

        // Rebuild chains with same mentions (they already have correct offsets)
        let chains = doc.chains.clone();

        Self { entities, chains }
    }
}

/// Convert multiple CorefDocuments to TrainingExamples.
pub fn coref_documents_to_training_examples(docs: &[CorefDocument]) -> Vec<TrainingExample> {
    docs.iter().map(TrainingExample::from).collect()
}

// =============================================================================
// Training Configuration
// =============================================================================

/// Training configuration for box embeddings.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Learning rate
    pub learning_rate: f32,
    /// Weight for negative pairs
    pub negative_weight: f32,
    /// Margin for negative pairs
    pub margin: f32,
    /// L2 regularization weight
    pub regularization: f32,
    /// Number of training epochs
    pub epochs: usize,
    /// Batch size (for mini-batch training)
    pub batch_size: usize,
    /// Warmup epochs (linear increase from 0.1*lr to lr)
    pub warmup_epochs: usize,
    /// Use self-adversarial negative sampling
    pub use_self_adversarial: bool,
    /// Temperature for self-adversarial sampling
    pub adversarial_temperature: f32,
    /// Early stopping patience (stop if no improvement for N epochs)
    pub early_stopping_patience: Option<usize>,
    /// Minimum improvement for early stopping (relative)
    pub early_stopping_min_delta: f32,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            learning_rate: 0.001,
            negative_weight: 0.5,
            margin: 0.3,
            regularization: 0.0001,
            epochs: 100,
            batch_size: 32,
            warmup_epochs: 10,
            use_self_adversarial: true,
            adversarial_temperature: 1.0,
            early_stopping_patience: Some(10),
            early_stopping_min_delta: 0.001,
        }
    }
}

// =============================================================================
// AMSGrad Optimizer State
// =============================================================================

/// AMSGrad optimizer state for a single box.
#[derive(Debug, Clone)]
pub struct AMSGradState {
    /// First moment estimate (m)
    pub m: Vec<f32>,
    /// Second moment estimate (v)
    pub v: Vec<f32>,
    /// Max second moment estimate (v_hat)
    pub v_hat: Vec<f32>,
    /// Iteration counter
    pub t: usize,
    /// Learning rate
    pub lr: f32,
    /// Beta1 (momentum)
    pub beta1: f32,
    /// Beta2 (RMSprop)
    pub beta2: f32,
    /// Epsilon (numerical stability)
    pub epsilon: f32,
}

impl AMSGradState {
    /// Create new AMSGrad state.
    pub fn new(dim: usize, learning_rate: f32) -> Self {
        Self {
            m: vec![0.0; dim],
            v: vec![0.0; dim],
            v_hat: vec![0.0; dim],
            t: 0,
            lr: learning_rate,
            beta1: 0.9,
            beta2: 0.999,
            epsilon: 1e-8,
        }
    }

    /// Set learning rate.
    pub fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }
}

// =============================================================================
// Trainer
// =============================================================================

/// Trainer for box embeddings.
pub struct BoxEmbeddingTrainer {
    /// Training configuration
    config: TrainingConfig,
    /// Entity ID → TrainableBox mapping
    boxes: HashMap<usize, TrainableBox>,
    /// Entity ID → AMSGradState mapping
    optimizer_states: HashMap<usize, AMSGradState>,
    /// Embedding dimension
    dim: usize,
}

impl BoxEmbeddingTrainer {
    /// Create a new trainer.
    ///
    /// # Arguments
    ///
    /// * `config` - Training configuration
    /// * `dim` - Embedding dimension
    /// * `initial_embeddings` - Optional initial vector embeddings (entity_id → vector)
    pub fn new(
        config: TrainingConfig,
        dim: usize,
        initial_embeddings: Option<HashMap<usize, Vec<f32>>>,
    ) -> Self {
        let mut boxes = HashMap::new();
        let mut optimizer_states = HashMap::new();

        if let Some(embeddings) = initial_embeddings {
            // Initialize from vector embeddings
            for (entity_id, vector) in embeddings {
                assert_eq!(vector.len(), dim);
                let box_embedding = TrainableBox::from_vector(&vector, 0.1);
                boxes.insert(entity_id, box_embedding.clone());
                optimizer_states.insert(entity_id, AMSGradState::new(dim, config.learning_rate));
            }
        }

        Self {
            config,
            boxes,
            optimizer_states,
            dim,
        }
    }

    /// Initialize boxes from entities.
    ///
    /// Creates trainable boxes for all entities, either from provided
    /// vector embeddings or random initialization.
    ///
    /// **Key insight**: For positive pairs (entities that corefer), initialize
    /// boxes to overlap so gradients can flow from the start.
    ///
    /// # Arguments
    ///
    /// * `examples` - Training examples with entities and coreference chains
    /// * `initial_embeddings` - Optional pre-computed vector embeddings (entity_id → vector)
    ///   If provided, boxes are initialized around these vectors. If None, uses smart
    ///   random initialization with shared centers for coreferent entities.
    pub fn initialize_boxes(
        &mut self,
        examples: &[TrainingExample],
        initial_embeddings: Option<&HashMap<usize, Vec<f32>>>,
    ) {
        // Collect all unique entity IDs and build coreference groups
        let mut entity_ids = std::collections::HashSet::new();
        let mut coref_groups: Vec<Vec<usize>> = Vec::new();
        
        for example in examples {
            for entity in &example.entities {
                let entity_id = entity.start;
                entity_ids.insert(entity_id);
            }
            
            // Build groups of entities that corefer
            for chain in &example.chains {
                let group: Vec<usize> = chain.mentions.iter().map(|m| m.start).collect();
                if group.len() > 1 {
                    coref_groups.push(group);
                }
            }
        }

        // Initialize boxes
        for &entity_id in &entity_ids {
            // If vector embeddings provided, use them (better initialization)
            if let Some(embeddings) = initial_embeddings {
                if let Some(vector) = embeddings.get(&entity_id) {
                    // Normalize vector to unit length for better initialization
                    let norm: f32 = vector.iter().map(|&x| x * x).sum::<f32>().sqrt();
                    let normalized: Vec<f32> = if norm > 0.0 {
                        vector.iter().map(|&x| x / norm).collect()
                    } else {
                        vector.clone()
                    };
                    
                    // Use larger initial width (0.2) when starting from vectors
                    // This ensures boxes can overlap even if vectors are slightly different
                    let box_embedding = TrainableBox::from_vector(&normalized, 0.2);
                    self.boxes.insert(entity_id, box_embedding.clone());
                    self.optimizer_states.insert(entity_id, AMSGradState::new(self.dim, self.config.learning_rate));
                    continue;
                }
            }

            // Check if this entity is in a coreference group
            let mut group_center: Option<Vec<f32>> = None;
            let mut in_coref_group = false;
            
            for group in &coref_groups {
                if group.contains(&entity_id) {
                    // Use a shared center for all entities in the group
                    if group_center.is_none() {
                        group_center = Some(
                            (0..self.dim)
                                .map(|_| (simple_random() - 0.5) * 0.3) // Smaller region for better overlap
                                .collect()
                        );
                    }
                    in_coref_group = true;
                    break;
                }
            }

            // Initialize: if in coref group, use shared center; otherwise random
            let mu = if let Some(ref center) = group_center {
                // Add very small random offset to shared center (ensures overlap)
                center.iter()
                    .map(|&c| c + (simple_random() - 0.5) * 0.05) // Very small offset
                    .collect()
            } else {
                // Random center, but spread out more to avoid accidental overlap
                (0..self.dim)
                    .map(|_| (simple_random() - 0.5) * 1.0)
                    .collect()
            };
            
            // Balanced initialization: coreferent entities should overlap, but not too much
            // Non-coreferent entities should be distinct
            let initial_width = if in_coref_group {
                0.6_f32 // Moderate width for coreferent entities (ensures overlap but not excessive)
            } else {
                0.2_f32 // Smaller width for non-coreferent (keeps them distinct)
            };
            let delta: Vec<f32> = vec![initial_width.ln(); self.dim];
            let box_embedding = TrainableBox::new(mu, delta);
            self.boxes.insert(entity_id, box_embedding.clone());
            self.optimizer_states.insert(entity_id, AMSGradState::new(self.dim, self.config.learning_rate));
        }
    }

    /// Train on a single example.
    fn train_example(&mut self, example: &TrainingExample, epoch: usize) -> f32 {
        let mut total_loss = 0.0;
        let mut num_pairs = 0;

        // Update learning rate with warmup and decay
        let current_lr = get_learning_rate(
            epoch,
            self.config.epochs,
            self.config.learning_rate,
            self.config.warmup_epochs,
        );
        for state in self.optimizer_states.values_mut() {
            state.set_lr(current_lr);
        }

        // Build positive pairs (entities in same chain)
        let mut positive_pairs = Vec::new();
        for chain in &example.chains {
            let mentions: Vec<usize> = chain
                .mentions
                .iter()
                .map(|m| m.start)
                .collect();
            for i in 0..mentions.len() {
                for j in (i + 1)..mentions.len() {
                    positive_pairs.push((mentions[i], mentions[j]));
                }
            }
        }

        // Build negative pairs (entities in different chains)
        let mut negative_pairs = Vec::new();
        for i in 0..example.chains.len() {
            for j in (i + 1)..example.chains.len() {
                let chain_i: Vec<usize> = example.chains[i]
                    .mentions
                    .iter()
                    .map(|m| m.start)
                    .collect();
                let chain_j: Vec<usize> = example.chains[j]
                    .mentions
                    .iter()
                    .map(|m| m.start)
                    .collect();
                for &id_i in &chain_i {
                    for &id_j in &chain_j {
                        negative_pairs.push((id_i, id_j));
                    }
                }
            }
        }

        // Accumulate gradients for all pairs
        let mut gradients: HashMap<usize, (Vec<f32>, Vec<f32>)> = HashMap::new();

        // Process positive pairs
        for &(id_a, id_b) in &positive_pairs {
            // Clone boxes for gradient computation
            let box_a = self.boxes.get(&id_a).cloned();
            let box_b = self.boxes.get(&id_b).cloned();
            
            if let (Some(box_a_ref), Some(box_b_ref)) = (box_a.as_ref(), box_b.as_ref()) {
                let loss = compute_pair_loss(box_a_ref, box_b_ref, true, &self.config);
                total_loss += loss;
                num_pairs += 1;

                // Compute analytical gradients
                let (grad_mu_a, grad_delta_a, grad_mu_b, grad_delta_b) =
                    compute_analytical_gradients(box_a_ref, box_b_ref, true, &self.config);

                // Skip if gradients are invalid
                if grad_mu_a.iter().any(|&x| !x.is_finite()) 
                    || grad_delta_a.iter().any(|&x| !x.is_finite())
                    || grad_mu_b.iter().any(|&x| !x.is_finite())
                    || grad_delta_b.iter().any(|&x| !x.is_finite()) {
                    continue;
                }

                // Accumulate gradients
                let entry_a = gradients.entry(id_a).or_insert_with(|| (vec![0.0; self.dim], vec![0.0; self.dim]));
                for i in 0..self.dim {
                    entry_a.0[i] += grad_mu_a[i];
                    entry_a.1[i] += grad_delta_a[i];
                }
                
                let entry_b = gradients.entry(id_b).or_insert_with(|| (vec![0.0; self.dim], vec![0.0; self.dim]));
                for i in 0..self.dim {
                    entry_b.0[i] += grad_mu_b[i];
                    entry_b.1[i] += grad_delta_b[i];
                }
            }
        }

        // Process negative pairs (with self-adversarial sampling if enabled)
        let negative_samples: Vec<(usize, usize)> = if self.config.use_self_adversarial && !negative_pairs.is_empty() {
            // Sample based on current predictions
            let num_samples = positive_pairs.len().min(negative_pairs.len());
            let sampled_indices = sample_self_adversarial_negatives(
                &negative_pairs,
                &self.boxes,
                num_samples,
                self.config.adversarial_temperature,
            );
            sampled_indices.iter().map(|&idx| negative_pairs[idx]).collect()
        } else {
            // Uniform sampling
            let num_samples = positive_pairs.len().min(negative_pairs.len());
            negative_pairs.into_iter().take(num_samples).collect()
        };

        for &(id_a, id_b) in &negative_samples {
            // Clone boxes for gradient computation
            let box_a = self.boxes.get(&id_a).cloned();
            let box_b = self.boxes.get(&id_b).cloned();
            
            if let (Some(box_a_ref), Some(box_b_ref)) = (box_a.as_ref(), box_b.as_ref()) {
                let loss = compute_pair_loss(box_a_ref, box_b_ref, false, &self.config);
                total_loss += loss;
                num_pairs += 1;

                // Compute analytical gradients
                let (grad_mu_a, grad_delta_a, grad_mu_b, grad_delta_b) =
                    compute_analytical_gradients(box_a_ref, box_b_ref, false, &self.config);

                // Skip if gradients are invalid
                if grad_mu_a.iter().any(|&x| !x.is_finite()) 
                    || grad_delta_a.iter().any(|&x| !x.is_finite())
                    || grad_mu_b.iter().any(|&x| !x.is_finite())
                    || grad_delta_b.iter().any(|&x| !x.is_finite()) {
                    continue;
                }

                // Accumulate gradients
                let entry_a = gradients.entry(id_a).or_insert_with(|| (vec![0.0; self.dim], vec![0.0; self.dim]));
                for i in 0..self.dim {
                    entry_a.0[i] += grad_mu_a[i];
                    entry_a.1[i] += grad_delta_a[i];
                }
                
                let entry_b = gradients.entry(id_b).or_insert_with(|| (vec![0.0; self.dim], vec![0.0; self.dim]));
                for i in 0..self.dim {
                    entry_b.0[i] += grad_mu_b[i];
                    entry_b.1[i] += grad_delta_b[i];
                }
            }
        }

        // Apply accumulated gradients using AMSGrad
        for (entity_id, (grad_mu, grad_delta)) in gradients {
            if let (Some(box_mut), Some(state)) = (
                self.boxes.get_mut(&entity_id),
                self.optimizer_states.get_mut(&entity_id),
            ) {
                box_mut.update_amsgrad(&grad_mu, &grad_delta, state);
            }
        }

        if num_pairs > 0 {
            total_loss / num_pairs as f32
        } else {
            0.0
        }
    }

    /// Train on a dataset with mini-batching and early stopping.
    pub fn train(&mut self, examples: &[TrainingExample]) -> Vec<f32> {
        let mut losses = Vec::new();
        let mut best_loss = f32::INFINITY;
        let mut patience_counter = 0;

        for epoch in 0..self.config.epochs {
            // Shuffle examples for better training (simple Fisher-Yates)
            let mut shuffled_indices: Vec<usize> = (0..examples.len()).collect();
            for i in (1..shuffled_indices.len()).rev() {
                let j = (simple_random() * (i + 1) as f32) as usize;
                shuffled_indices.swap(i, j);
            }

            let mut epoch_loss = 0.0;
            let mut num_batches = 0;

            // Mini-batch training
            for batch_start in (0..examples.len()).step_by(self.config.batch_size) {
                let batch_end = (batch_start + self.config.batch_size).min(examples.len());
                let batch_indices = &shuffled_indices[batch_start..batch_end];
                
                let mut batch_loss = 0.0;
                let mut batch_pairs = 0;

                // Process batch
                for &idx in batch_indices {
                    let example = &examples[idx];
                    let loss = self.train_example(example, epoch);
                    batch_loss += loss;
                    batch_pairs += 1;
                }

                if batch_pairs > 0 {
                    epoch_loss += batch_loss / batch_pairs as f32;
                    num_batches += 1;
                }
            }

            let avg_loss = if num_batches > 0 {
                epoch_loss / num_batches as f32
            } else {
                0.0
            };
            losses.push(avg_loss);

            let current_lr = get_learning_rate(
                epoch,
                self.config.epochs,
                self.config.learning_rate,
                self.config.warmup_epochs,
            );

            // Early stopping check
            let improved = avg_loss < best_loss - self.config.early_stopping_min_delta;
            if improved {
                best_loss = avg_loss;
                patience_counter = 0;
            } else {
                patience_counter += 1;
            }

            // Show overlap stats periodically
            if epoch % 10 == 0 || epoch == self.config.epochs - 1 || improved {
                let (avg_pos, avg_neg, overlap_rate) = self.get_overlap_stats(examples);
                let status = if improved { "✓" } else { " " };
                let patience_info = if let Some(patience) = self.config.early_stopping_patience {
                    format!(", patience={}/{}", patience_counter, patience)
                } else {
                    String::new()
                };
                let loss_reduction = if losses.len() > 1 {
                    format!(" ({:.1}%↓)", (1.0 - avg_loss / losses[0]) * 100.0)
                } else {
                    String::new()
                };
                let score_gap = avg_pos - avg_neg; // Positive should be higher than negative
                println!("Epoch {}: loss = {:.4}{}, lr = {:.6}, best = {:.4} {} ({} batches{})", 
                    epoch, avg_loss, loss_reduction, current_lr, best_loss, status, num_batches, patience_info);
                println!("  Overlap: {:.1}%, Pos: {:.4}, Neg: {:.4}, Gap: {:.4} {}", 
                    overlap_rate * 100.0, avg_pos, avg_neg, score_gap,
                    if score_gap > 0.0 { "✓" } else { "⚠" });
            }

            // Early stopping
            if let Some(patience) = self.config.early_stopping_patience {
                if patience_counter >= patience {
                    println!("Early stopping at epoch {} (no improvement for {} epochs)", epoch, patience);
                    break;
                }
            }
        }

        losses
    }

    /// Get trained boxes for inference.
    pub fn get_boxes(&self) -> HashMap<usize, BoxEmbedding> {
        self.boxes
            .iter()
            .map(|(id, trainable)| (*id, trainable.to_box()))
            .collect()
    }

    /// Get diagnostic statistics about box overlaps.
    ///
    /// Returns (avg_positive_score, avg_negative_score, overlap_rate)
    pub fn get_overlap_stats(&self, examples: &[TrainingExample]) -> (f32, f32, f32) {
        let mut positive_scores = Vec::new();
        let mut negative_scores = Vec::new();
        let mut overlapping_pairs = 0;
        let mut total_pairs = 0;

        for example in examples {
            // Positive pairs
            for chain in &example.chains {
                let mentions: Vec<usize> = chain.mentions.iter().map(|m| m.start).collect();
                for i in 0..mentions.len() {
                    for j in (i + 1)..mentions.len() {
                        if let (Some(box_a), Some(box_b)) = (self.boxes.get(&mentions[i]), self.boxes.get(&mentions[j])) {
                            let box_a_embed = box_a.to_box();
                            let box_b_embed = box_b.to_box();
                            let score = box_a_embed.coreference_score(&box_b_embed);
                            positive_scores.push(score);
                            if score > 0.01 {
                                overlapping_pairs += 1;
                            }
                            total_pairs += 1;
                        }
                    }
                }
            }

            // Negative pairs
            for i in 0..example.chains.len() {
                for j in (i + 1)..example.chains.len() {
                    let chain_i: Vec<usize> = example.chains[i].mentions.iter().map(|m| m.start).collect();
                    let chain_j: Vec<usize> = example.chains[j].mentions.iter().map(|m| m.start).collect();
                    for &id_i in &chain_i {
                        for &id_j in &chain_j {
                            if let (Some(box_a), Some(box_b)) = (self.boxes.get(&id_i), self.boxes.get(&id_j)) {
                                let box_a_embed = box_a.to_box();
                                let box_b_embed = box_b.to_box();
                                let score = box_a_embed.coreference_score(&box_b_embed);
                                negative_scores.push(score);
                            }
                        }
                    }
                }
            }
        }

        let avg_positive = if !positive_scores.is_empty() {
            positive_scores.iter().sum::<f32>() / positive_scores.len() as f32
        } else {
            0.0
        };

        let avg_negative = if !negative_scores.is_empty() {
            negative_scores.iter().sum::<f32>() / negative_scores.len() as f32
        } else {
            0.0
        };

        let overlap_rate = if total_pairs > 0 {
            overlapping_pairs as f32 / total_pairs as f32
        } else {
            0.0
        };

        (avg_positive, avg_negative, overlap_rate)
    }

    /// Evaluate coreference accuracy on a test set.
    ///
    /// Returns (accuracy, precision, recall, f1) where:
    /// - Accuracy: fraction of pairs correctly classified
    /// - Precision: fraction of predicted positives that are correct
    /// - Recall: fraction of true positives that are predicted
    /// - F1: harmonic mean of precision and recall
    pub fn evaluate(&self, examples: &[TrainingExample], threshold: f32) -> (f32, f32, f32, f32) {
        let mut true_positives = 0;
        let mut false_positives = 0;
        let mut false_negatives = 0;
        let mut total_pairs = 0;

        for example in examples {
            // Build positive pairs (should corefer)
            let mut positive_pairs = Vec::new();
            for chain in &example.chains {
                let mentions: Vec<usize> = chain
                    .mentions
                    .iter()
                    .map(|m| m.start)
                    .collect();
                for i in 0..mentions.len() {
                    for j in (i + 1)..mentions.len() {
                        positive_pairs.push((mentions[i], mentions[j]));
                    }
                }
            }

            // Build negative pairs (shouldn't corefer)
            let mut negative_pairs = Vec::new();
            for i in 0..example.chains.len() {
                for j in (i + 1)..example.chains.len() {
                    let chain_i: Vec<usize> = example.chains[i]
                        .mentions
                        .iter()
                        .map(|m| m.start)
                        .collect();
                    let chain_j: Vec<usize> = example.chains[j]
                        .mentions
                        .iter()
                        .map(|m| m.start)
                        .collect();
                    for &id_i in &chain_i {
                        for &id_j in &chain_j {
                            negative_pairs.push((id_i, id_j));
                        }
                    }
                }
            }

            // Evaluate positive pairs
            for &(id_a, id_b) in &positive_pairs {
                total_pairs += 1;
                if let (Some(box_a), Some(box_b)) = (self.boxes.get(&id_a), self.boxes.get(&id_b)) {
                    let box_a_embed = box_a.to_box();
                    let box_b_embed = box_b.to_box();
                    let score = box_a_embed.coreference_score(&box_b_embed);
                    if score >= threshold {
                        true_positives += 1;
                    } else {
                        false_negatives += 1;
                    }
                }
            }

            // Evaluate negative pairs
            for &(id_a, id_b) in &negative_pairs {
                total_pairs += 1;
                if let (Some(box_a), Some(box_b)) = (self.boxes.get(&id_a), self.boxes.get(&id_b)) {
                    let box_a_embed = box_a.to_box();
                    let box_b_embed = box_b.to_box();
                    let score = box_a_embed.coreference_score(&box_b_embed);
                    if score >= threshold {
                        false_positives += 1;
                    }
                }
            }
        }

        // Compute metrics
        let precision = if true_positives + false_positives > 0 {
            true_positives as f32 / (true_positives + false_positives) as f32
        } else {
            0.0
        };

        let recall = if true_positives + false_negatives > 0 {
            true_positives as f32 / (true_positives + false_negatives) as f32
        } else {
            0.0
        };

        let f1 = if precision + recall > 0.0 {
            2.0 * precision * recall / (precision + recall)
        } else {
            0.0
        };

        let accuracy = if total_pairs > 0 {
            (true_positives + (total_pairs - true_positives - false_positives - false_negatives)) as f32 / total_pairs as f32
        } else {
            0.0
        };

        (accuracy, precision, recall, f1)
    }

    /// Save trained boxes to a file (JSON format).
    ///
    /// # Arguments
    ///
    /// * `path` - File path to save to
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// trainer.save_boxes("trained_boxes.json")?;
    /// ```
    pub fn save_boxes(&self, path: &str) -> Result<(), Box<dyn std::error::Error>> {
        use std::fs::File;
        use std::io::Write;
        
        let serialized = serde_json::to_string_pretty(&self.boxes)?;
        let mut file = File::create(path)?;
        file.write_all(serialized.as_bytes())?;
        Ok(())
    }

    /// Load trained boxes from a file (JSON format).
    ///
    /// # Arguments
    ///
    /// * `path` - File path to load from
    /// * `dim` - Expected embedding dimension
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let boxes = BoxEmbeddingTrainer::load_boxes("trained_boxes.json", 32)?;
    /// ```
    pub fn load_boxes(path: &str, dim: usize) -> Result<HashMap<usize, TrainableBox>, Box<dyn std::error::Error>> {
        use std::fs::File;
        use std::io::Read;
        
        let mut file = File::open(path)?;
        let mut contents = String::new();
        file.read_to_string(&mut contents)?;
        let boxes: HashMap<usize, TrainableBox> = serde_json::from_str(&contents)?;
        
        // Verify dimensions
        for (id, box_embedding) in &boxes {
            if box_embedding.dim != dim {
                return Err(format!("Box for entity {} has dimension {}, expected {}", id, box_embedding.dim, dim).into());
            }
        }
        
        Ok(boxes)
    }
}

/// Split training examples into train/validation sets.
///
/// # Arguments
///
/// * `examples` - All training examples
/// * `val_ratio` - Fraction of examples to use for validation (0.0-1.0)
///
/// # Returns
///
/// (train_examples, val_examples)
pub fn split_train_val(examples: &[TrainingExample], val_ratio: f32) -> (Vec<TrainingExample>, Vec<TrainingExample>) {
    let val_size = (examples.len() as f32 * val_ratio) as usize;
    let mut shuffled: Vec<TrainingExample> = examples.to_vec();
    
    // Simple shuffle
    for i in (1..shuffled.len()).rev() {
        let j = (simple_random() * (i + 1) as f32) as usize;
        shuffled.swap(i, j);
    }
    
    let val_examples = shuffled.split_off(val_size);
    (shuffled, val_examples)
}

// =============================================================================
// Loss and Gradient Computation
// =============================================================================

/// Compute loss for a pair of boxes.
fn compute_pair_loss(
    box_a: &TrainableBox,
    box_b: &TrainableBox,
    is_positive: bool,
    config: &TrainingConfig,
) -> f32 {
    let box_a_embed = box_a.to_box();
    let box_b_embed = box_b.to_box();

    if is_positive {
        // Positive pair: maximize conditional probability
        let p_a_b = box_a_embed.conditional_probability(&box_b_embed);
        let p_b_a = box_b_embed.conditional_probability(&box_a_embed);
        
        // Clamp probabilities to avoid log(0)
        let p_a_b = p_a_b.max(1e-8);
        let p_b_a = p_b_a.max(1e-8);
        
        // Use symmetric score: min of both conditional probabilities
        // This ensures both boxes must overlap significantly
        let min_prob = p_a_b.min(p_b_a);
        let neg_log_prob = -min_prob.ln();
        
        // Also add penalty if boxes are too far apart (encourages movement)
        let vol_intersection = box_a_embed.intersection_volume(&box_b_embed);
        let distance_penalty = if vol_intersection < 1e-10 {
            // Boxes don't overlap - add distance penalty
            let center_a = box_a_embed.center();
            let center_b = box_b_embed.center();
            let dist: f32 = center_a.iter()
                .zip(center_b.iter())
                .map(|(a, b)| (a - b).powi(2))
                .sum::<f32>()
                .sqrt();
            0.3 * dist // Stronger penalty to encourage boxes to move closer
        } else {
            // Even if overlapping, encourage more overlap
            // Penalize if overlap is too small
            let vol_a = box_a_embed.volume();
            let vol_b = box_b_embed.volume();
            let overlap_ratio = vol_intersection / vol_a.min(vol_b).max(1e-10);
            if overlap_ratio < 0.5 {
                // Encourage more overlap
                0.1 * (0.5 - overlap_ratio)
            } else {
                0.0
            }
        };
        
        // Regularization: penalize large volumes (stronger for positive pairs)
        let vol_a = box_a_embed.volume();
        let vol_b = box_b_embed.volume();
        // Use stronger regularization to prevent boxes from growing too large
        let reg = config.regularization * 2.0 * (vol_a + vol_b);
        
        (neg_log_prob + reg + distance_penalty).max(0.0)
    } else {
        // Negative pair: enforce separation
        // We want conditional probability to be LOW (boxes should be disjoint)
        let p_a_b = box_a_embed.conditional_probability(&box_b_embed);
        let p_b_a = box_b_embed.conditional_probability(&box_a_embed);
        
        // Use max of both conditional probabilities
        let max_prob = p_a_b.max(p_b_a);
        
        // Loss: penalize high conditional probability
        // If max_prob > margin, we have a violation - penalize it
        // Loss should be high when max_prob is high
        let margin_loss = if max_prob > config.margin {
            (max_prob - config.margin).powi(2) // Quadratic penalty for exceeding margin
        } else {
            0.0 // No loss if already below margin (good!)
        };
        
        // Add extra penalty if boxes overlap significantly
        let vol_intersection = box_a_embed.intersection_volume(&box_b_embed);
        let vol_a = box_a_embed.volume();
        let vol_b = box_b_embed.volume();
        let overlap_penalty = if vol_intersection > 1e-10 {
            // Boxes overlap - add very strong penalty
            let overlap_ratio = vol_intersection / vol_a.min(vol_b).max(1e-10);
            2.0 * overlap_ratio // Very strong penalty for overlapping negative pairs
        } else {
            0.0
        };
        
        // Base loss: penalize any non-zero conditional probability
        // Even if below margin, we want it as low as possible
        let base_loss = max_prob * 0.5; // Small penalty for any probability
        
        config.negative_weight * (margin_loss + overlap_penalty + base_loss)
    }
}

/// Compute analytical gradients for a pair of boxes.
fn compute_analytical_gradients(
    box_a: &TrainableBox,
    box_b: &TrainableBox,
    is_positive: bool,
    config: &TrainingConfig,
) -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>) {
    let box_a_embed = box_a.to_box();
    let box_b_embed = box_b.to_box();
    let dim = box_a.dim;

    // Initialize gradients
    let mut grad_mu_a = vec![0.0; dim];
    let mut grad_delta_a = vec![0.0; dim];
    let mut grad_mu_b = vec![0.0; dim];
    let mut grad_delta_b = vec![0.0; dim];

    // Compute intersection volume and individual volumes
    let vol_a = box_a_embed.volume();
    let vol_b = box_b_embed.volume();
    let vol_intersection = box_a_embed.intersection_volume(&box_b_embed);

    if is_positive {
        // Positive pair: L = -log(P(A|B)) - log(P(B|A)) + reg
        // P(A|B) = Vol(A ∩ B) / Vol(B)
        // P(B|A) = Vol(A ∩ B) / Vol(A)
        
        let p_a_b = if vol_b > 0.0 { vol_intersection / vol_b } else { 0.0 };
        let p_b_a = if vol_a > 0.0 { vol_intersection / vol_a } else { 0.0 };
        
        // Clamp to avoid division by zero
        let p_a_b = p_a_b.max(1e-8);
        let p_b_a = p_b_a.max(1e-8);
        
        // Gradients through -log(P(A|B)) = -log(Vol_intersection) + log(Vol_B)
        // For positive pairs, we want to maximize overlap
        // If boxes don't overlap, we need gradients to move them together
        
        // Check if boxes overlap
        let vol_intersection = box_a_embed.intersection_volume(&box_b_embed);
        let has_overlap = vol_intersection > 1e-10;
        
        if !has_overlap {
            // Boxes don't overlap - add gradient to move centers closer
            let center_a = box_a_embed.center();
            let center_b = box_b_embed.center();
            for i in 0..dim {
                let diff = center_b[i] - center_a[i];
                // Moderate gradient to move boxes together
                grad_mu_a[i] += 0.2 * diff;
                grad_mu_b[i] += -0.2 * diff;
                
                // Small increase in box sizes to help them overlap (but not too much)
                grad_delta_a[i] += 0.03; // Very small growth for box A
                grad_delta_b[i] += 0.03; // Very small growth for box B
            }
        }
        
        for i in 0..dim {
            // Gradient w.r.t. box A
            // ∂(-log(P(A|B)))/∂δ_A = -1/P(A|B) * ∂P(A|B)/∂δ_A
            // ∂P(A|B)/∂δ_A = (1/Vol_B) * ∂Vol_intersection/∂δ_A
            // ∂Vol_intersection/∂δ_A = Vol_intersection (if boxes overlap in dim i)
            
            let overlap_i = if box_a_embed.min[i] < box_b_embed.max[i] && box_b_embed.min[i] < box_a_embed.max[i] {
                // Boxes overlap in dimension i
                let min_overlap = box_a_embed.min[i].max(box_b_embed.min[i]);
                let max_overlap = box_a_embed.max[i].min(box_b_embed.max[i]);
                (max_overlap - min_overlap).max(0.0)
            } else {
                0.0
            };
            
            if overlap_i > 0.0 && vol_intersection > 0.0 {
                // Gradient through intersection volume
                // When boxes overlap, focus on improving overlap ratio (not just growing boxes)
                let overlap_ratio_a = vol_intersection / vol_a.max(1e-10);
                let overlap_ratio_b = vol_intersection / vol_b.max(1e-10);
                
                // If overlap ratio is low, encourage slight growth; if high, maintain
                if overlap_ratio_a < 0.2 {
                    // Very low overlap - small growth
                    grad_delta_a[i] += 0.02;
                } else if overlap_ratio_a < 0.5 {
                    // Moderate overlap - tiny growth
                    grad_delta_a[i] += 0.01;
                }
                // If overlap_ratio_a >= 0.5, don't grow (good enough)
                
                if overlap_ratio_b < 0.2 {
                    // Very low overlap - small growth
                    grad_delta_b[i] += 0.02;
                } else if overlap_ratio_b < 0.5 {
                    // Moderate overlap - tiny growth
                    grad_delta_b[i] += 0.01;
                }
                
                // Gradient through conditional probability (main signal)
                let grad_vol_intersection_delta_a = vol_intersection * 0.15; // Conservative
                let grad_p_a_b_delta_a = grad_vol_intersection_delta_a / vol_b.max(1e-8);
                grad_delta_a[i] += -grad_p_a_b_delta_a / p_a_b.max(1e-8) * 0.3; // Scaled down
                
                let grad_vol_intersection_delta_b = vol_intersection * 0.15;
                let grad_p_b_a_delta_b = grad_vol_intersection_delta_b / vol_a.max(1e-8);
                grad_delta_b[i] += -grad_p_b_a_delta_b / p_b_a.max(1e-8) * 0.3;
            } else {
                // Boxes don't overlap in this dimension - tiny growth to achieve overlap
                grad_delta_a[i] += 0.01; // Very small growth for box A
                grad_delta_b[i] += 0.01; // Very small growth for box B
            }
            
            // Regularization gradient: ∂(λ * Vol)/∂δ = λ * Vol
            // Strong regularization to prevent boxes from growing too large
            // Apply to both positive and negative pairs
            grad_delta_a[i] += config.regularization * 2.0 * vol_a;
            grad_delta_b[i] += config.regularization * 2.0 * vol_b;
        }
    } else {
        // Negative pair: L = max(0, margin - max(P(A|B), P(B|A))) * λ_neg + overlap_penalty
        let p_a_b = if vol_b > 0.0 { vol_intersection / vol_b } else { 0.0 };
        let p_b_a = if vol_a > 0.0 { vol_intersection / vol_a } else { 0.0 };
        let max_prob = p_a_b.max(p_b_a);
        
        if max_prob < config.margin || vol_intersection > 1e-10 {
            // Active margin loss or boxes overlap (shouldn't for negatives)
            for i in 0..dim {
                // Check if boxes overlap in this dimension
                let overlap_i = if box_a_embed.min[i] < box_b_embed.max[i] && box_b_embed.min[i] < box_a_embed.max[i] {
                    let min_overlap = box_a_embed.min[i].max(box_b_embed.min[i]);
                    let max_overlap = box_a_embed.max[i].min(box_b_embed.max[i]);
                    (max_overlap - min_overlap).max(0.0)
                } else {
                    0.0
                };
                
                if overlap_i > 0.0 {
                    // Boxes overlap - strong gradient to separate
                    // Move centers apart
                    let center_a = box_a_embed.center();
                    let center_b = box_b_embed.center();
                    let diff = center_b[i] - center_a[i];
                    
                    // Gradient to push boxes apart
                    if diff > 0.0 {
                        grad_mu_a[i] += -config.negative_weight * 0.5 * diff; // Move A left
                        grad_mu_b[i] += config.negative_weight * 0.5 * diff;  // Move B right
                    } else {
                        grad_mu_a[i] += -config.negative_weight * 0.5 * diff; // Move A right
                        grad_mu_b[i] += config.negative_weight * 0.5 * diff;  // Move B left
                    }
                    
                    // Shrink boxes to reduce overlap (stronger)
                    grad_delta_a[i] += -config.negative_weight * 0.5; // Shrink box A
                    grad_delta_b[i] += -config.negative_weight * 0.5; // Shrink box B
                } else {
                    // Boxes don't overlap - just ensure they stay separated
                    // Moderate gradient to shrink boxes (prevent growth)
                    grad_delta_a[i] += -config.negative_weight * 0.1;
                    grad_delta_b[i] += -config.negative_weight * 0.1;
                }
            }
        }
    }

    // Clip gradients to prevent explosion
    for i in 0..dim {
        grad_mu_a[i] = (grad_mu_a[i] as f32).clamp(-10.0_f32, 10.0_f32);
        grad_delta_a[i] = (grad_delta_a[i] as f32).clamp(-10.0_f32, 10.0_f32);
        grad_mu_b[i] = (grad_mu_b[i] as f32).clamp(-10.0_f32, 10.0_f32);
        grad_delta_b[i] = (grad_delta_b[i] as f32).clamp(-10.0_f32, 10.0_f32);
    }

    (grad_mu_a, grad_delta_a, grad_mu_b, grad_delta_b)
}

/// Sample negative pairs using self-adversarial sampling.
fn sample_self_adversarial_negatives(
    negative_pairs: &[(usize, usize)],
    boxes: &HashMap<usize, TrainableBox>,
    num_samples: usize,
    temperature: f32,
) -> Vec<usize> {
    // Compute scores for all negative pairs
    let mut scores: Vec<(usize, f32)> = negative_pairs
        .iter()
        .enumerate()
        .filter_map(|(idx, &(id_a, id_b))| {
            if let (Some(box_a), Some(box_b)) = (boxes.get(&id_a), boxes.get(&id_b)) {
                let box_a_embed = box_a.to_box();
                let box_b_embed = box_b.to_box();
                let score = box_a_embed.coreference_score(&box_b_embed);
                Some((idx, score / temperature))
            } else {
                None
            }
        })
        .collect();

    // Sort by score (descending) - higher scores are "harder" negatives
    scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    // Sample top-k (hardest negatives)
    scores.into_iter()
        .take(num_samples)
        .map(|(idx, _)| idx)
        .collect()
}

/// Get learning rate with warmup and cosine decay.
fn get_learning_rate(epoch: usize, total_epochs: usize, base_lr: f32, warmup_epochs: usize) -> f32 {
    if epoch < warmup_epochs {
        // Linear warmup: 0.1 * lr → lr
        let warmup_lr = base_lr * 0.1;
        warmup_lr + (base_lr - warmup_lr) * (epoch as f32 / warmup_epochs as f32)
    } else {
        // Cosine decay: lr → 0.1 * lr
        let progress = (epoch - warmup_epochs) as f32 / (total_epochs - warmup_epochs).max(1) as f32;
        let min_lr = base_lr * 0.1;
        min_lr + (base_lr - min_lr) * (1.0 + (std::f32::consts::PI * progress).cos()) / 2.0
    }
}

// =============================================================================
// TrainableBox AMSGrad Update
// =============================================================================

impl TrainableBox {
    /// Update box parameters using AMSGrad optimizer.
    pub fn update_amsgrad(
        &mut self,
        grad_mu: &[f32],
        grad_delta: &[f32],
        state: &mut AMSGradState,
    ) {
        state.t += 1;
        let t = state.t as f32;

        // Update first moment (m)
        for i in 0..self.dim {
            state.m[i] = state.beta1 * state.m[i] + (1.0 - state.beta1) * grad_mu[i];
        }

        // Update second moment (v) and max (v_hat)
        for i in 0..self.dim {
            let v_new = state.beta2 * state.v[i] + (1.0 - state.beta2) * grad_mu[i] * grad_mu[i];
            state.v[i] = v_new;
            state.v_hat[i] = state.v_hat[i].max(v_new);
        }

        // Bias correction for first moment
        let m_hat: Vec<f32> = state.m.iter().map(|&m| m / (1.0 - state.beta1.powf(t))).collect();

        // Update mu
        for i in 0..self.dim {
            let update = state.lr * m_hat[i] / (state.v_hat[i].sqrt() + state.epsilon);
            self.mu[i] -= update;
            
            // Ensure finite
            if !self.mu[i].is_finite() {
                self.mu[i] = 0.0;
            }
        }

        // Similar for delta
        let mut m_delta = vec![0.0; self.dim];
        let mut v_delta = vec![0.0; self.dim];
        let mut v_hat_delta = vec![0.0; self.dim];

        for i in 0..self.dim {
            m_delta[i] = state.beta1 * m_delta[i] + (1.0 - state.beta1) * grad_delta[i];
            let v_new: f32 = state.beta2 * v_delta[i] + (1.0 - state.beta2) * grad_delta[i] * grad_delta[i];
            v_delta[i] = v_new;
            v_hat_delta[i] = (v_hat_delta[i] as f32).max(v_new);
        }

        let m_hat_delta: Vec<f32> = m_delta.iter().map(|&m| m / (1.0 - state.beta1.powf(t))).collect();

        for i in 0..self.dim {
            let update = state.lr * m_hat_delta[i] / ((v_hat_delta[i] as f32).sqrt() + state.epsilon);
            self.delta[i] -= update;
            
            // Clamp delta to reasonable range (width between 0.01 and 10.0)
            self.delta[i] = self.delta[i].clamp(0.01_f32.ln(), 10.0_f32.ln());
            
            // Ensure finite
            if !self.delta[i].is_finite() {
                self.delta[i] = 0.5_f32.ln();
            }
        }
    }
}

// =============================================================================
// Simple Random Number Generator
// =============================================================================

/// Simple random number generator (for when rand feature is not available).
fn simple_random() -> f32 {
    use std::collections::hash_map::DefaultHasher;
    use std::hash::{Hash, Hasher};
    use std::time::{SystemTime, UNIX_EPOCH};
    
    static mut COUNTER: usize = 0;
    unsafe {
        COUNTER += 1;
        let mut hasher = DefaultHasher::new();
        SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_nanos().hash(&mut hasher);
        COUNTER.hash(&mut hasher);
        let hash = hasher.finish();
        (hash as f32) / (u64::MAX as f32)
    }
}
