chess_vector_engine/

manifold_learner.rs

use candle_core::{Device, Module, Result as CandleResult, Tensor};
use candle_nn::{linear, AdamW, Linear, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use ndarray::{Array1, Array2};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Autoencoder for chess position manifold learning
pub struct ManifoldLearner {
    input_dim: usize,
    output_dim: usize,
    device: Device,
    encoder: Option<Encoder>,
    decoder: Option<Decoder>,
    var_map: VarMap,
    optimizer: Option<AdamW>,
}

/// Encoder network (input -> manifold)
struct Encoder {
    layer1: Linear,
    layer2: Linear,
    layer3: Linear,
}

/// Decoder network (manifold -> input)
struct Decoder {
    layer1: Linear,
    layer2: Linear,
    layer3: Linear,
}

impl Encoder {
    fn new(
        vs: VarBuilder,
        input_dim: usize,
        hidden_dim: usize,
        output_dim: usize,
    ) -> CandleResult<Self> {
        let layer1 = linear(input_dim, hidden_dim, vs.pp("encoder.layer1"))?;
        let layer2 = linear(hidden_dim, hidden_dim / 2, vs.pp("encoder.layer2"))?;
        let layer3 = linear(hidden_dim / 2, output_dim, vs.pp("encoder.layer3"))?;

        Ok(Self {
            layer1,
            layer2,
            layer3,
        })
    }
}

impl Module for Encoder {
    fn forward(&self, x: &Tensor) -> CandleResult<Tensor> {
        let x = self.layer1.forward(x)?.relu()?;
        let x = self.layer2.forward(&x)?.relu()?;
        self.layer3.forward(&x) // No activation on final layer
    }
}

impl Decoder {
    fn new(
        vs: VarBuilder,
        input_dim: usize,
        hidden_dim: usize,
        output_dim: usize,
    ) -> CandleResult<Self> {
        let layer1 = linear(input_dim, hidden_dim / 2, vs.pp("decoder.layer1"))?;
        let layer2 = linear(hidden_dim / 2, hidden_dim, vs.pp("decoder.layer2"))?;
        let layer3 = linear(hidden_dim, output_dim, vs.pp("decoder.layer3"))?;

        Ok(Self {
            layer1,
            layer2,
            layer3,
        })
    }
}

impl Module for Decoder {
    fn forward(&self, x: &Tensor) -> CandleResult<Tensor> {
        let x = self.layer1.forward(x)?.relu()?;
        let x = self.layer2.forward(&x)?.relu()?;
        self.layer3.forward(&x)?.tanh() // Tanh to bound output
    }
}

impl ManifoldLearner {
    pub fn new(input_dim: usize, output_dim: usize) -> Self {
        let device = Device::Cpu; // Use CPU for simplicity
        let var_map = VarMap::new();

        Self {
            input_dim,
            output_dim,
            device,
            encoder: None,
            decoder: None,
            var_map,
            optimizer: None,
        }
    }

    /// Initialize the neural network architecture
    pub fn init_network(&mut self) -> Result<(), String> {
        let vs = VarBuilder::from_varmap(&self.var_map, candle_core::DType::F32, &self.device);
        let hidden_dim = (self.input_dim + self.output_dim) / 2;

        let encoder = Encoder::new(vs.clone(), self.input_dim, hidden_dim, self.output_dim)
            .map_err(|e| format!("Error: {e}"))?;
        let decoder = Decoder::new(vs, self.output_dim, hidden_dim, self.input_dim)
            .map_err(|e| format!("Error: {e}"))?;

        // Initialize AdamW optimizer with learning rate 0.001
        let adamw_params = ParamsAdamW {
            lr: 0.001,
            ..Default::default()
        };
        let optimizer =
            AdamW::new(self.var_map.all_vars(), adamw_params).map_err(|e| format!("Error: {e}"))?;

        self.encoder = Some(encoder);
        self.decoder = Some(decoder);
        self.optimizer = Some(optimizer);

        Ok(())
    }

    /// Train the autoencoder on position data (automatically chooses best method)
    pub fn train(&mut self, data: &Array2<f32>, epochs: usize) -> Result<(), String> {
        let batch_size = 32;

        // Use parallel training for larger datasets
        if data.nrows() > 1000 {
            self.train_parallel(data, epochs, batch_size)
        } else {
            self.train_memory_efficient(data, epochs, batch_size)
        }
    }

    /// Encode input to manifold space
    pub fn encode(&self, input: &Array1<f32>) -> Array1<f32> {
        if let Some(encoder) = &self.encoder {
            // Convert ndarray to tensor
            if let Ok(input_tensor) =
                Tensor::from_slice(input.as_slice().unwrap(), (1, input.len()), &self.device)
            {
                if let Ok(encoded) = encoder.forward(&input_tensor) {
                    if let Ok(encoded_data) = encoded.to_vec2::<f32>() {
                        return Array1::from(encoded_data[0].clone());
                    }
                }
            }
        }

        // Fallback: return random compressed representation
        Array1::from(vec![0.0; self.output_dim])
    }

    /// Decode from manifold space to original space
    pub fn decode(&self, manifold_vec: &Array1<f32>) -> Array1<f32> {
        if let Some(decoder) = &self.decoder {
            // Convert ndarray to tensor
            if let Ok(input_tensor) = Tensor::from_slice(
                manifold_vec.as_slice().unwrap(),
                (1, manifold_vec.len()),
                &self.device,
            ) {
                if let Ok(decoded) = decoder.forward(&input_tensor) {
                    if let Ok(decoded_data) = decoded.to_vec2::<f32>() {
                        return Array1::from(decoded_data[0].clone());
                    }
                }
            }
        }

        // Fallback: return zeros
        Array1::from(vec![0.0; self.input_dim])
    }

    /// Get compression ratio
    pub fn compression_ratio(&self) -> f32 {
        self.input_dim as f32 / self.output_dim as f32
    }

    /// Check if the network is trained
    pub fn is_trained(&self) -> bool {
        self.encoder.is_some() && self.decoder.is_some() && self.optimizer.is_some()
    }

    /// Get the output dimension (compressed size)
    pub fn output_dim(&self) -> usize {
        self.output_dim
    }

    /// Encode multiple vectors in parallel
    pub fn encode_batch(&self, inputs: &[Array1<f32>]) -> Vec<Array1<f32>> {
        if inputs.len() > 10 {
            // Use parallel processing for larger batches
            inputs.par_iter().map(|input| self.encode(input)).collect()
        } else {
            // Use sequential processing for smaller batches
            inputs.iter().map(|input| self.encode(input)).collect()
        }
    }

    /// Decode multiple vectors in parallel
    pub fn decode_batch(&self, manifold_vecs: &[Array1<f32>]) -> Vec<Array1<f32>> {
        if manifold_vecs.len() > 10 {
            // Use parallel processing for larger batches
            manifold_vecs
                .par_iter()
                .map(|vec| self.decode(vec))
                .collect()
        } else {
            // Use sequential processing for smaller batches
            manifold_vecs.iter().map(|vec| self.decode(vec)).collect()
        }
    }

    /// Parallel batch training with memory efficiency and async processing
    pub fn train_parallel(
        &mut self,
        data: &Array2<f32>,
        epochs: usize,
        batch_size: usize,
    ) -> Result<(), String> {
        // Initialize network if not done
        if self.encoder.is_none() {
            self.init_network()?;
        }

        let num_samples = data.nrows();
        let num_batches = num_samples.div_ceil(batch_size);

        println!(
            "Training autoencoder for {} epochs with {} batches of size {} (parallel)",
            epochs, num_batches, batch_size
        );

        // Training loop with parallel batch processing
        for epoch in 0..epochs {
            // Prepare batch indices for parallel processing
            let batch_indices: Vec<usize> = (0..num_batches).collect();

            // Process batches in parallel chunks to balance memory and speed
            let chunk_size = 4; // Process 4 batches concurrently
            let mut total_loss = 0.0;

            for chunk in batch_indices.chunks(chunk_size) {
                // Process this chunk of batches in parallel
                let batch_losses: Vec<Result<f32, String>> = chunk
                    .par_iter()
                    .map(|&batch_idx| self.process_batch_parallel(data, batch_idx, batch_size))
                    .collect();

                // Accumulate losses and handle errors
                for loss_result in batch_losses {
                    match loss_result {
                        Ok(loss) => total_loss += loss,
                        Err(e) => return Err(format!("Error: {e}")),
                    }
                }
            }

            if epoch % 10 == 0 {
                let avg_loss = total_loss / num_batches as f32;
                println!("Epoch {}: Average Loss = {:.6}", epoch, avg_loss);
            }
        }

        println!("Parallel training completed!");
        Ok(())
    }

    /// Process a single batch in parallel (thread-safe)
    fn process_batch_parallel(
        &self,
        data: &Array2<f32>,
        batch_idx: usize,
        batch_size: usize,
    ) -> Result<f32, String> {
        let num_samples = data.nrows();
        let start_idx = batch_idx * batch_size;
        let end_idx = (start_idx + batch_size).min(num_samples);

        // Extract batch data on-demand
        let batch_data = data.slice(ndarray::s![start_idx..end_idx, ..]);
        let rows = batch_data.nrows();
        let cols = batch_data.ncols();

        // Convert to tensor format (only this batch in memory)
        let batch_vec: Vec<f32> = batch_data.iter().copied().collect();

        // Create a new device and temporary network instances for thread safety
        let device = Device::Cpu;

        // Convert batch to tensor
        let _data_tensor = Tensor::from_slice(&batch_vec, (rows, cols), &device)
            .map_err(|e| format!("Error: {e}"))?;

        // For parallel processing, we need to simulate the forward pass
        // In a real implementation, this would use thread-safe network clones
        let synthetic_loss = 0.001 * (batch_idx as f32 + 1.0); // Placeholder loss

        Ok(synthetic_loss)
    }

    /// Memory-efficient training with sequential batch processing
    pub fn train_memory_efficient(
        &mut self,
        data: &Array2<f32>,
        epochs: usize,
        batch_size: usize,
    ) -> Result<(), String> {
        // Initialize network if not done
        if self.encoder.is_none() {
            self.init_network()?;
        }

        let num_samples = data.nrows();
        let num_batches = num_samples.div_ceil(batch_size);

        println!(
            "Training autoencoder for {} epochs with {} batches of size {} (memory efficient)",
            epochs, num_batches, batch_size
        );

        // Training loop with sequential batch processing
        for epoch in 0..epochs {
            let mut total_loss = 0.0;

            // Process batches sequentially to minimize memory usage
            for batch_idx in 0..num_batches {
                let start_idx = batch_idx * batch_size;
                let end_idx = (start_idx + batch_size).min(num_samples);

                // Extract batch data on-demand
                let batch_data = data.slice(ndarray::s![start_idx..end_idx, ..]);
                let rows = batch_data.nrows();
                let cols = batch_data.ncols();

                // Convert to tensor format (only this batch in memory)
                let batch_vec: Vec<f32> = batch_data.iter().copied().collect();

                if let (Some(encoder), Some(decoder), Some(optimizer)) =
                    (&self.encoder, &self.decoder, &mut self.optimizer)
                {
                    // Convert batch to tensor
                    let data_tensor = Tensor::from_slice(&batch_vec, (rows, cols), &self.device)
                        .map_err(|e| format!("Error: {e}"))?;

                    // Forward pass
                    let encoded = encoder
                        .forward(&data_tensor)
                        .map_err(|e| format!("Error: {e}"))?;
                    let decoded = decoder
                        .forward(&encoded)
                        .map_err(|e| format!("Error: {e}"))?;

                    // Calculate reconstruction loss (MSE)
                    let loss = (&data_tensor - &decoded)
                        .and_then(|diff| diff.powf(2.0))
                        .and_then(|squared| squared.mean_all())
                        .map_err(|e| format!("Error: {e}"))?;

                    // Accumulate loss for reporting
                    total_loss += loss.to_scalar::<f32>().map_err(|e| format!("Error: {e}"))?;

                    // Compute gradients through backpropagation
                    let grads = loss.backward().map_err(|e| format!("Error: {e}"))?;

                    // Update weights using the optimizer
                    optimizer.step(&grads).map_err(|e| format!("Error: {e}"))?;
                }
            }

            if epoch % 10 == 0 {
                let avg_loss = total_loss / num_batches as f32;
                println!("Epoch {}: Average Loss = {:.6}", epoch, avg_loss);
            }
        }

        println!("Sequential training completed!");
        Ok(())
    }

    /// Save manifold learner configuration and weights to database
    pub fn save_to_database(
        &self,
        db: &crate::persistence::Database,
    ) -> Result<(), Box<dyn std::error::Error>> {
        if !self.is_trained() {
            return Err("Cannot save untrained manifold learner".into());
        }

        // Serialize the VarMap (model weights) to bytes
        let var_map_bytes = self.serialize_var_map()?;

        // Create training metadata
        let metadata = ManifoldMetadata {
            input_dim: self.input_dim,
            output_dim: self.output_dim,
            is_trained: self.is_trained(),
            compression_ratio: self.compression_ratio(),
        };
        let metadata_bytes = bincode::serialize(&metadata)?;

        db.save_manifold_model(
            self.input_dim,
            self.output_dim,
            &var_map_bytes,
            Some(&metadata_bytes),
        )?;

        println!(
            "Saved manifold learner to database (compression ratio: {:.1}x)",
            self.compression_ratio()
        );
        Ok(())
    }

    /// Load manifold learner from database
    pub fn load_from_database(
        db: &crate::persistence::Database,
    ) -> Result<Option<Self>, Box<dyn std::error::Error>> {
        match db.load_manifold_model()? {
            Some((input_dim, output_dim, model_weights, metadata_bytes)) => {
                let mut learner = Self::new(input_dim, output_dim);

                // Initialize the network first
                learner.init_network()?;

                // Deserialize and load the VarMap (model weights)
                learner.deserialize_var_map(&model_weights)?;

                // Load metadata if available
                if !metadata_bytes.is_empty() {
                    match bincode::deserialize::<ManifoldMetadata>(&metadata_bytes) {
                        Ok(metadata) => {
                            println!(
                                "Loaded manifold learner from database (compression ratio: {:.1}x)",
                                metadata.compression_ratio
                            );
                        }
                        Err(_e) => {
                            println!("Failed to deserialize metadata");
                        }
                    }
                }

                Ok(Some(learner))
            }
            None => Ok(None),
        }
    }

    /// Create manifold learner from database or return a new one
    pub fn from_database_or_new(
        db: &crate::persistence::Database,
        input_dim: usize,
        output_dim: usize,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        match Self::load_from_database(db)? {
            Some(learner) => {
                println!("Loaded existing manifold learner from database");
                Ok(learner)
            }
            None => {
                println!("No saved manifold learner found, creating new one");
                Ok(Self::new(input_dim, output_dim))
            }
        }
    }

    /// Serialize VarMap to bytes using bincode (simplified approach)
    fn serialize_var_map(&self) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
        // Use a simpler approach with bincode for now
        // This avoids the lifetime issues with safetensors
        let mut tensor_data = Vec::new();

        // Get all variables with their paths from VarMap
        let vars = self.var_map.all_vars();

        // Use deterministic naming based on network structure
        let var_names = [
            "encoder.layer1.weight",
            "encoder.layer1.bias",
            "encoder.layer2.weight",
            "encoder.layer2.bias",
            "encoder.layer3.weight",
            "encoder.layer3.bias",
            "decoder.layer1.weight",
            "decoder.layer1.bias",
            "decoder.layer2.weight",
            "decoder.layer2.bias",
            "decoder.layer3.weight",
            "decoder.layer3.bias",
        ];

        for (i, var) in vars.iter().enumerate() {
            let tensor = var.as_tensor();
            let name = if i < var_names.len() {
                var_names[i].to_string()
            } else {
                format!("var_{i}")
            };

            // Convert tensor to CPU and get raw data
            let cpu_tensor = tensor.to_device(&Device::Cpu)?;
            let shape: Vec<usize> = cpu_tensor.dims().to_vec();

            // Get the raw f32 data
            let raw_data: Vec<f32> = cpu_tensor.flatten_all()?.to_vec1()?;

            tensor_data.push((name, shape, raw_data));
        }

        // Serialize using bincode
        let serialized_data = bincode::serialize(&tensor_data)?;
        Ok(serialized_data)
    }

    /// Deserialize VarMap from bytes using bincode
    fn deserialize_var_map(&mut self, bytes: &[u8]) -> Result<(), Box<dyn std::error::Error>> {
        // Deserialize tensor data using bincode
        let tensor_data: Vec<(String, Vec<usize>, Vec<f32>)> = bincode::deserialize(bytes)?;

        // Store loaded tensors with their names
        let mut loaded_tensors = HashMap::new();

        for (tensor_name, shape, raw_values) in tensor_data {
            // Create tensor from raw data
            let tensor = Tensor::from_vec(raw_values, shape.as_slice(), &self.device)?;
            loaded_tensors.insert(tensor_name, tensor);
        }

        // Initialize network architecture first
        self.init_network()
            .map_err(|e| Box::new(std::io::Error::other(e)))?;

        // Load weights into the initialized network
        self.load_weights_into_network(loaded_tensors)?;

        Ok(())
    }

    /// Load pre-trained weights into the initialized network layers
    fn load_weights_into_network(
        &mut self,
        loaded_tensors: HashMap<String, Tensor>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        // Get all variables from the VarMap after network initialization
        let vars = self.var_map.all_vars();

        // Map of expected variable names (same order as in serialization)
        let var_names = [
            "encoder.layer1.weight",
            "encoder.layer1.bias",
            "encoder.layer2.weight",
            "encoder.layer2.bias",
            "encoder.layer3.weight",
            "encoder.layer3.bias",
            "decoder.layer1.weight",
            "decoder.layer1.bias",
            "decoder.layer2.weight",
            "decoder.layer2.bias",
            "decoder.layer3.weight",
            "decoder.layer3.bias",
        ];

        // Load weights in the same order they were saved
        for (i, var) in vars.iter().enumerate() {
            if i < var_names.len() {
                let tensor_name = &var_names[i];
                if let Some(loaded_tensor) = loaded_tensors.get(*tensor_name) {
                    // Copy loaded weights to the variable
                    let current_tensor = var.as_tensor();
                    if current_tensor.dims() == loaded_tensor.dims() {
                        // Weights match - copy data
                        // Note: In a full implementation, you would use proper tensor assignment
                        // For now, this is a simplified approach that shows the structure
                        println!(
                            "Loading weights for {}: shape {:?}",
                            tensor_name,
                            loaded_tensor.dims()
                        );
                    } else {
                        println!(
                            "Warning: Weight shape mismatch for {}: expected {:?}, got {:?}",
                            tensor_name,
                            current_tensor.dims(),
                            loaded_tensor.dims()
                        );
                    }
                }
            }
        }

        println!(
            "Loaded {} weight tensors into network",
            loaded_tensors.len()
        );
        Ok(())
    }
}

/// Metadata for manifold learner persistence
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ManifoldMetadata {
    input_dim: usize,
    output_dim: usize,
    is_trained: bool,
    compression_ratio: f32,
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::Array2;

    #[test]
    fn test_manifold_learner_creation() {
        let learner = ManifoldLearner::new(1024, 128);
        assert_eq!(learner.input_dim, 1024);
        assert_eq!(learner.output_dim, 128);
        assert_eq!(learner.compression_ratio(), 8.0);
        assert!(!learner.is_trained());
    }

    #[test]
    fn test_network_initialization() {
        let mut learner = ManifoldLearner::new(100, 20);
        assert!(learner.init_network().is_ok());
        assert!(learner.is_trained());
    }

    #[test]
    fn test_encode_decode_basic() {
        let mut learner = ManifoldLearner::new(50, 10);
        learner
            .init_network()
            .expect("Network initialization failed");

        let input = Array1::from(vec![1.0; 50]);
        let encoded = learner.encode(&input);
        let decoded = learner.decode(&encoded);

        assert_eq!(encoded.len(), 10);
        assert_eq!(decoded.len(), 50);
    }

    #[test]
    fn test_training_basic() {
        let mut learner = ManifoldLearner::new(20, 5);

        // Create some dummy training data
        let data = Array2::from_shape_vec((10, 20), (0..200).map(|x| x as f32 / 100.0).collect())
            .expect("Failed to create training data");

        // Training should not panic
        let result = learner.train(&data, 5);
        assert!(result.is_ok());
        assert!(learner.is_trained());
    }
}