foxhunt/crates/ml-supervised/src/tft/variable_selection.rs

//! Variable Selection Network for TFT
//!
//! Implements learnable feature selection using gated linear units and
//! soft feature selection weights for improved interpretability.

use std::collections::HashMap;
use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use super::GatedResidualNetwork;
use crate::gpu_tensor::{
    gpu_add, gpu_cat_dim1, gpu_narrow_2d, gpu_softmax, GpuLinear, GpuTensor,
};

/// Variable Selection Network for feature importance learning
#[derive(Debug)]
pub struct VariableSelectionNetwork {
    pub input_size: usize,
    pub hidden_size: usize,
    // Gated Linear Units for variable selection
    single_var_grns: Vec<GatedResidualNetwork>,
    // Soft attention weights
    attention_weights: GpuLinear,
    // Feature importance tracking
    importance_scores: HashMap<usize, f64>,
    stream: Arc<CudaStream>,
}

impl VariableSelectionNetwork {
    pub fn new(input_size: usize, hidden_size: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        // Create GRN for flattened inputs (keeps param symmetry with original)
        let _flattened_grn = GatedResidualNetwork::new(input_size, hidden_size, stream)?;

        // Create individual GRNs for each variable
        let mut single_var_grns = Vec::new();
        for _i in 0..input_size {
            let grn = GatedResidualNetwork::new(1, hidden_size, stream)?;
            single_var_grns.push(grn);
        }

        // Attention layer for variable selection
        let attention_weights = GpuLinear::new(hidden_size * input_size, input_size, stream)?;

        Ok(Self {
            input_size,
            hidden_size,
            single_var_grns,
            attention_weights,
            importance_scores: HashMap::new(),
            stream: Arc::clone(stream),
        })
    }

    pub fn forward(
        &mut self,
        inputs: &GpuTensor,
        context: Option<&GpuTensor>,
    ) -> Result<GpuTensor, MLError> {
        let batch_size = inputs.dim(0)?;
        let is_2d = inputs.shape.len() == 2;
        let seq_len = if is_2d { 1 } else { inputs.dim(1)? };
        let input_feat = if is_2d { inputs.dim(1)? } else { inputs.dim(2)? };

        // Reshape to 2D for processing: [batch * seq_len, input_feat]
        let flat_input = if is_2d {
            inputs.reshape(&[batch_size, input_feat])?
        } else {
            inputs.reshape(&[batch_size * seq_len, input_feat])?
        };
        let flat_rows = flat_input.dim(0)?;

        // Process individual variables on GPU
        let mut var_outputs: Vec<GpuTensor> = Vec::with_capacity(self.input_size);

        for (i, grn) in self.single_var_grns.iter().enumerate() {
            // Extract variable i from flat_input: column i -> [flat_rows, 1]
            let var_tensor = gpu_narrow_2d(&flat_input, 1, i, 1)?;
            let var_output = grn.forward(&var_tensor, context)?; // [flat_rows, hidden_size]
            var_outputs.push(var_output);
        }

        // Concatenate all variable outputs along dim1: [flat_rows, hidden_size * input_size]
        let mut cat_tensor = var_outputs.first().ok_or_else(|| {
            MLError::InvalidInput("No variable outputs".to_owned())
        })?.clone();
        for var_out in var_outputs.iter().skip(1) {
            cat_tensor = gpu_cat_dim1(&cat_tensor, var_out)?;
        }

        // Compute attention weights
        let raw_weights = self.attention_weights.forward(&cat_tensor)?; // [flat_rows, input_size]
        let attention = gpu_softmax(&raw_weights)?;

        // Update importance scores (single scalar download per variable -- acceptable for monitoring)
        self.update_importance_scores(&attention)?;

        // Apply variable selection: weighted sum of per-variable GRN outputs
        // attention: [flat_rows, input_size], var_outputs[i]: [flat_rows, hidden_size]
        // result = sum_i(attention[:, i:i+1] * var_outputs[i])
        let mut result = GpuTensor::zeros(&[flat_rows, self.hidden_size], &self.stream)?;
        for (i, var_out) in var_outputs.iter().enumerate() {
            // Extract attention weight for variable i: [flat_rows, 1]
            let attn_col = gpu_narrow_2d(&attention, 1, i, 1)?;
            // Broadcast multiply: [flat_rows, 1] * [flat_rows, hidden_size] -> [flat_rows, hidden_size]
            // We need to broadcast attn_col across hidden_size columns.
            // Use gpu_broadcast_mul_col: a[rows, cols] * b[rows, 1] -> [rows, cols]
            let weighted = crate::gpu_tensor::gpu_broadcast_mul_col(var_out, &attn_col)?;
            result = gpu_add(&result, &weighted)?;
        }

        // Reshape to [batch, seq_len, hidden_size]
        result.reshape(&[batch_size, seq_len, self.hidden_size])
    }

    fn update_importance_scores(&mut self, attention_weights: &GpuTensor) -> Result<(), MLError> {
        // Use gpu_mean_all on per-column slices to avoid bulk download
        let cols = attention_weights.dim(1)?;
        self.importance_scores.clear();
        for c in 0..cols {
            let col = gpu_narrow_2d(attention_weights, 1, c, 1)?;
            let mean_val = crate::gpu_tensor::gpu_mean_all(&col)?;
            self.importance_scores.insert(c, mean_val as f64);
        }
        Ok(())
    }

    pub fn get_importance_scores(&self) -> Result<Vec<f64>, MLError> {
        let mut scores = vec![0.0; self.input_size];
        for (i, &score) in &self.importance_scores {
            if *i < self.input_size {
                scores[*i] = score;
            }
        }

        let sum: f64 = scores.iter().sum();
        if sum == 0.0 {
            let uniform_score = 1.0 / self.input_size as f64;
            scores.fill(uniform_score);
        }

        Ok(scores)
    }

    pub fn get_top_features(&self, k: usize) -> Vec<(usize, f64)> {
        let mut features: Vec<(usize, f64)> = self
            .importance_scores
            .iter()
            .map(|(&idx, &score)| (idx, score))
            .collect();

        features.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        features.truncate(k);
        features
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    #[test]
    fn test_variable_selection_network_creation() -> Result<(), MLError> {
        let stream = test_stream();
        let vsn = VariableSelectionNetwork::new(10, 64, &stream)?;
        assert_eq!(vsn.input_size, 10);
        assert_eq!(vsn.hidden_size, 64);
        Ok(())
    }

    #[test]
    fn test_variable_selection_forward_2d() -> Result<(), MLError> {
        let stream = test_stream();
        let mut vsn = VariableSelectionNetwork::new(5, 32, &stream)?;
        let input = GpuTensor::from_vec(
            vec![1.0, 2.0, 3.0, 4.0, 5.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            &[2, 5],
            &stream,
        )?;
        let output = vsn.forward(&input, None)?;
        assert_eq!(output.shape, vec![2, 1, 32]);
        Ok(())
    }

    #[test]
    fn test_variable_selection_forward_3d() -> Result<(), MLError> {
        let stream = test_stream();
        let mut vsn = VariableSelectionNetwork::new(3, 16, &stream)?;
        let input = GpuTensor::from_vec(vec![1.0_f32; 24], &[2, 4, 3], &stream)?;
        let output = vsn.forward(&input, None)?;
        assert_eq!(output.shape, vec![2, 4, 16]);
        Ok(())
    }

    #[test]
    fn test_importance_scores() -> Result<(), MLError> {
        let stream = test_stream();
        let vsn = VariableSelectionNetwork::new(5, 32, &stream)?;
        let scores = vsn.get_importance_scores()?;
        assert_eq!(scores.len(), 5);
        let sum: f64 = scores.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
        Ok(())
    }
}