ClawFoxyVision 0.2.0

// External imports
use crate::daily::lstm::step_1_tensor_preparation::{
    dataframe_to_tensors, impute_missing_values, load_daily_csv, normalize_daily_features,
    split_daily_data,
};
use crate::daily::lstm::step_3_lstm_model_arch::{DailyLSTMModel, DailyLSTMModelConfig};
use anyhow::Result;
use burn::data::{
    dataloader::{batcher::Batcher, DataLoaderBuilder},
    dataset::Dataset,
};
use burn::module::Module;
use burn::optim::{AdamConfig, GradientsParams, Optimizer};
use burn::record::Recorder;
use burn::tensor::backend::AutodiffBackend;
use burn::tensor::cast::ToElement;
use burn::tensor::{backend::Backend, Tensor};
use log::info;
use std::collections::HashMap;
use std::path::Path;

// Define batch structure
#[derive(Debug, Clone)]
pub struct DailyBatch<B: Backend> {
    pub features: Tensor<B, 3>,
    pub targets: Tensor<B, 2>,
}

// Define batcher
pub struct DailyBatcher<B: Backend> {
    _phantom: std::marker::PhantomData<B>,
}

impl<B: Backend> DailyBatcher<B> {
    pub fn new() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<B: Backend> Batcher<B, (Tensor<B, 3>, Tensor<B, 2>), DailyBatch<B>> for DailyBatcher<B> {
    fn batch(
        &self,
        items: Vec<(Tensor<B, 3>, Tensor<B, 2>)>,
        _device: &B::Device,
    ) -> DailyBatch<B> {
        // Extract features and targets
        let (features, targets): (Vec<_>, Vec<_>) = items.into_iter().unzip();

        // Stack tensors
        let features = Tensor::cat(features, 0);
        let targets = Tensor::cat(targets, 0);

        DailyBatch { features, targets }
    }
}

// Define dataset
pub struct DailyDataset {
    features: Vec<Vec<f32>>,
    targets: Vec<Vec<f32>>,
}

impl DailyDataset {
    pub fn new(features: Vec<Vec<f32>>, targets: Vec<Vec<f32>>) -> Self {
        Self { features, targets }
    }
}

impl<B: Backend> Dataset<(Tensor<B, 3>, Tensor<B, 2>)> for DailyDataset {
    fn get(&self, index: usize) -> Option<(Tensor<B, 3>, Tensor<B, 2>)> {
        if index >= self.features.len() {
            return None;
        }

        let device = Default::default();

        // Create tensors using from_data with slices
        let x_shape = [1, self.features[index].len() / 6, 6];
        let y_shape = [1, self.targets[index].len()];

        // Convert data to f32
        let x_data: Vec<f32> = self.features[index].iter().map(|&x| x as f32).collect();
        let y_data: Vec<f32> = self.targets[index].iter().map(|&x| x as f32).collect();

        // Create tensors using from_data with slices
        let x = Tensor::<B, 1>::from_data(x_data.as_slice(), &device).reshape(x_shape);
        let y = Tensor::<B, 1>::from_data(y_data.as_slice(), &device).reshape(y_shape);

        Some((x, y))
    }

    fn len(&self) -> usize {
        self.features.len()
    }
}

// Training function
pub fn train_daily_lstm<B: AutodiffBackend<Gradients = GradientsParams>>(
    csv_path: &str,
    config: DailyLSTMModelConfig,
    device: &B::Device,
    learning_rate: f64,
    batch_size: usize,
    epochs: usize,
    sequence_length: usize,
    forecast_horizon: usize,
    model_save_path: Option<&Path>,
) -> Result<DailyLSTMModel<B>> {
    // Load and preprocess data
    info!("Loading data from {}", csv_path);
    let mut df = load_daily_csv(csv_path)?;

    // Handle missing values
    info!("Handling missing values");
    impute_missing_values(&mut df, "forward")?;

    // Normalize features
    info!("Normalizing features");
    normalize_daily_features(&mut df)?;

    // Split data into training and validation sets
    info!("Splitting data into training and validation sets");
    let (train_df, val_df) = split_daily_data(&df, 0.8)?;

    // Convert to tensors
    info!("Converting data to tensors");
    let (train_x, train_y) =
        dataframe_to_tensors::<B>(&train_df, sequence_length, forecast_horizon, device)?;

    let (val_x, val_y) =
        dataframe_to_tensors::<B>(&val_df, sequence_length, forecast_horizon, device)?;

    // Create model
    info!("Creating model");
    let mut model = config.init(device);

    // Create optimizer
    let optim_config = AdamConfig::new();
    let mut optim = optim_config.init();

    // Training loop
    info!("Starting training");
    let mut best_val_loss = f32::INFINITY;
    let mut metrics = HashMap::new();

    for epoch in 0..epochs {
        // Create training data loader
        let train_dataset = DailyDataset::new(
            vec![vec![0.0; sequence_length * 6]; train_x.dims()[0]],
            vec![vec![0.0; 1]; train_y.dims()[0]],
        );

        let train_loader = DataLoaderBuilder::new(DailyBatcher::new())
            .batch_size(batch_size)
            .shuffle(1) // 1 for true
            .build(train_dataset);

        // Training mode
        let mut epoch_loss = 0.0;
        let num_batches = (train_x.dims()[0] + batch_size - 1) / batch_size;

        // Process each batch
        for batch in train_loader.iter() {
            // Forward pass
            let pred = model.forward(batch.features.clone(), true);

            // Calculate loss (mean squared error)
            let two = Tensor::<B, 2>::ones_like(&pred);
            let loss = ((pred - batch.targets).powf(two)).mean();

            // Backward pass
            let grads = loss.backward();

            // Update parameters
            model = optim.step(learning_rate, model, grads);

            // Record metrics
            epoch_loss += loss.into_scalar().to_f32();
        }

        epoch_loss /= num_batches as f32;

        // Validation mode
        let val_dataset = DailyDataset::new(
            vec![vec![0.0; sequence_length * 6]; val_x.dims()[0]],
            vec![vec![0.0; 1]; val_y.dims()[0]],
        );

        let val_loader = DataLoaderBuilder::new(DailyBatcher::new())
            .batch_size(batch_size)
            .shuffle(0) // 0 for false
            .build(val_dataset);

        let mut val_loss = 0.0;
        let num_val_batches = (val_x.dims()[0] + batch_size - 1) / batch_size;

        // Process each validation batch
        for batch in val_loader.iter() {
            // Forward pass
            let pred = model.predict(batch.features.clone());

            // Calculate loss (mean squared error)
            let two = Tensor::<B, 2>::ones_like(&pred);
            let loss = ((pred - batch.targets).powf(two)).mean();

            // Record metrics
            val_loss += loss.into_scalar().to_f32();
        }

        val_loss /= num_val_batches as f32;

        // Log progress
        println!(
            "Epoch {}/{}: Train Loss = {:.6}, Validation Loss = {:.6}",
            epoch + 1,
            epochs,
            epoch_loss,
            val_loss
        );

        // Save best model
        if val_loss < best_val_loss {
            best_val_loss = val_loss;
            if let Some(path) = model_save_path {
                println!("New best model found! Saving to {:?}...", path);
                // Implement model saving in step_6_model_serialization.rs
            }
        }

        // Update metrics
        metrics.insert(format!("epoch_{}_train_loss", epoch + 1), epoch_loss);
        metrics.insert(format!("epoch_{}_val_loss", epoch + 1), val_loss);
    }

    // Return model
    Ok(model)
}