foxhunt/ml/src/tft/quantized_tft_forward.rs

// This file contains the complete forward pass implementation for QuantizedTFT
// To be integrated into quantized_tft.rs

use super::*;

impl QuantizedTemporalFusionTransformer {
    /// Validate input tensor dimensions and device placement
    fn validate_inputs(
        &self,
        static_features: &Tensor,
        historical_features: &Tensor,
        future_features: &Tensor,
    ) -> Result<(), MLError> {
        // Check static features: [batch, num_static_features]
        let static_dims = static_features.dims();
        if static_dims.len() != 2 {
            return Err(MLError::ModelError(format!(
                "Static features must be 2D [batch, features], got {:?}",
                static_dims
            )));
        }
        if static_dims[1] != self.config.num_static_features {
            return Err(MLError::ModelError(format!(
                "Static features dimension mismatch: expected {}, got {}",
                self.config.num_static_features, static_dims[1]
            )));
        }

        // Check historical features: [batch, seq_len, num_unknown_features]
        let hist_dims = historical_features.dims();
        if hist_dims.len() != 3 {
            return Err(MLError::ModelError(format!(
                "Historical features must be 3D [batch, seq, features], got {:?}",
                hist_dims
            )));
        }
        if hist_dims[2] != self.config.num_unknown_features {
            return Err(MLError::ModelError(format!(
                "Historical features dimension mismatch: expected {}, got {}",
                self.config.num_unknown_features, hist_dims[2]
            )));
        }

        // Check future features: [batch, horizon, num_known_features]
        let fut_dims = future_features.dims();
        if fut_dims.len() != 3 {
            return Err(MLError::ModelError(format!(
                "Future features must be 3D [batch, horizon, features], got {:?}",
                fut_dims
            )));
        }
        if fut_dims[2] != self.config.num_known_features {
            return Err(MLError::ModelError(format!(
                "Future features dimension mismatch: expected {}, got {}",
                self.config.num_known_features, fut_dims[2]
            )));
        }

        // Verify device consistency (compare CUDA vs CPU since Device doesn't implement PartialEq)
        let target_device = &self.device;
        if static_features.device().is_cuda() != target_device.is_cuda() {
            return Err(MLError::ModelError(format!(
                "Static features on wrong device: expected {:?}, got {:?}",
                target_device,
                static_features.device()
            )));
        }
        if historical_features.device().is_cuda() != target_device.is_cuda() {
            return Err(MLError::ModelError(format!(
                "Historical features on wrong device: expected {:?}, got {:?}",
                target_device,
                historical_features.device()
            )));
        }
        if future_features.device().is_cuda() != target_device.is_cuda() {
            return Err(MLError::ModelError(format!(
                "Future features on wrong device: expected {:?}, got {:?}",
                target_device,
                future_features.device()
            )));
        }

        Ok(())
    }

    /// Step 1: Static Variable Selection Network (placeholder)
    fn forward_static_vsn(&self, static_features: &Tensor) -> Result<Tensor, MLError> {
        // Placeholder: Project static features to hidden_dim
        // [batch, num_static_features] -> [batch, 1, hidden_dim]
        let batch_size = static_features.dims()[0];

        // Simple linear projection (in production, use quantized VSN)
        let hidden = Tensor::zeros(
            (batch_size, 1, self.config.hidden_dim),
            DType::F32,
            &self.device,
        )?;

        Ok(hidden)
    }

    /// Step 2: Historical Encoder (LSTM with INT8 weights)
    fn forward_historical_encoder(&self, historical_features: &Tensor) -> Result<Tensor, MLError> {
        let (batch_size, seq_len, _input_dim) = historical_features.dims3()?;

        if self.lstm_weights.is_empty() {
            // Fallback: return zeros if LSTM not initialized
            let output = Tensor::zeros(
                (batch_size, seq_len, self.config.hidden_dim),
                DType::F32,
                &self.device,
            )?;
            return Ok(output);
        }

        // Initialize hidden and cell states
        let hidden_dim = self.config.hidden_dim;
        let num_layers = self.lstm_weights.len();

        let mut h = Tensor::zeros((num_layers, batch_size, hidden_dim), DType::F32, &self.device)?;
        let mut c = Tensor::zeros((num_layers, batch_size, hidden_dim), DType::F32, &self.device)?;

        // Process through quantized LSTM layers
        let mut layer_input = historical_features.clone();

        for (layer_idx, layer_weights) in self.lstm_weights.iter().enumerate() {
            let (layer_output, h_new, c_new) = self.forward_lstm_layer(
                &layer_input,
                &h.narrow(0, layer_idx, 1)?.squeeze(0)?,
                &c.narrow(0, layer_idx, 1)?.squeeze(0)?,
                layer_weights,
            )?;

            layer_input = layer_output;

            // Update hidden and cell states
            h = h.slice_assign(&[layer_idx..layer_idx + 1], &h_new.unsqueeze(0)?)?;
            c = c.slice_assign(&[layer_idx..layer_idx + 1], &c_new.unsqueeze(0)?)?;
        }

        Ok(layer_input)
    }

    /// Forward pass through a single LSTM layer with quantized weights
    fn forward_lstm_layer(
        &self,
        input: &Tensor,
        h_prev: &Tensor,
        c_prev: &Tensor,
        layer_weights: &HashMap<String, QuantizedTensor>,
    ) -> Result<(Tensor, Tensor, Tensor), MLError> {
        let (batch_size, seq_len, _input_dim) = input.dims3()?;
        let _hidden_dim = self.config.hidden_dim;

        // Dequantize all LSTM weights
        let w_ii = self.quantizer.dequantize_tensor(
            layer_weights.get("W_ii").ok_or_else(|| MLError::ModelError("Missing W_ii".to_string()))?
        )?;
        let w_if = self.quantizer.dequantize_tensor(
            layer_weights.get("W_if").ok_or_else(|| MLError::ModelError("Missing W_if".to_string()))?
        )?;
        let w_ig = self.quantizer.dequantize_tensor(
            layer_weights.get("W_ig").ok_or_else(|| MLError::ModelError("Missing W_ig".to_string()))?
        )?;
        let w_io = self.quantizer.dequantize_tensor(
            layer_weights.get("W_io").ok_or_else(|| MLError::ModelError("Missing W_io".to_string()))?
        )?;

        let w_hi = self.quantizer.dequantize_tensor(
            layer_weights.get("W_hi").ok_or_else(|| MLError::ModelError("Missing W_hi".to_string()))?
        )?;
        let w_hf = self.quantizer.dequantize_tensor(
            layer_weights.get("W_hf").ok_or_else(|| MLError::ModelError("Missing W_hf".to_string()))?
        )?;
        let w_hg = self.quantizer.dequantize_tensor(
            layer_weights.get("W_hg").ok_or_else(|| MLError::ModelError("Missing W_hg".to_string()))?
        )?;
        let w_ho = self.quantizer.dequantize_tensor(
            layer_weights.get("W_ho").ok_or_else(|| MLError::ModelError("Missing W_ho".to_string()))?
        )?;

        let mut outputs = Vec::new();
        let mut h_t = h_prev.clone();
        let mut c_t = c_prev.clone();

        // Process each time step
        for t in 0..seq_len {
            let x_t = input.narrow(1, t, 1)?.squeeze(1)?; // [batch, input_dim]

            // Input gate: i_t = sigmoid(W_ii @ x_t + W_hi @ h_t)
            let i_t = manual_sigmoid(&(x_t.matmul(&w_ii)? + h_t.matmul(&w_hi)?)?)?;

            // Forget gate: f_t = sigmoid(W_if @ x_t + W_hf @ h_t)
            let f_t = manual_sigmoid(&(x_t.matmul(&w_if)? + h_t.matmul(&w_hf)?)?)?;

            // Cell gate: g_t = tanh(W_ig @ x_t + W_hg @ h_t)
            let g_t = (x_t.matmul(&w_ig)? + h_t.matmul(&w_hg)?)?.tanh()?;

            // Output gate: o_t = sigmoid(W_io @ x_t + W_ho @ h_t)
            let o_t = manual_sigmoid(&(x_t.matmul(&w_io)? + h_t.matmul(&w_ho)?)?)?;

            // Cell state: c_t = f_t * c_t + i_t * g_t
            c_t = (f_t.mul(&c_t)? + i_t.mul(&g_t)?)?;

            // Hidden state: h_t = o_t * tanh(c_t)
            h_t = o_t.mul(&c_t.tanh()?)?;

            outputs.push(h_t.unsqueeze(1)?);
        }

        // Concatenate outputs along time dimension
        let output = Tensor::cat(&outputs, 1)?; // [batch, seq_len, hidden_dim]

        Ok((output, h_t, c_t))
    }

    /// Step 3: Future Decoder (simplified - uses same LSTM architecture)
    fn forward_future_decoder(&self, future_features: &Tensor) -> Result<Tensor, MLError> {
        // For simplicity, treat future decoder similarly to historical encoder
        // In production TFT, this would be a separate decoder with different weights
        self.forward_historical_encoder(future_features)
    }

    /// Step 5: Combine encodings (static, attention output, future)
    fn combine_contexts(
        &self,
        static_encoding: &Tensor,
        attention_output: &Tensor,
        future_encoding: &Tensor,
    ) -> Result<Tensor, MLError> {
        let (batch_size, hist_seq_len, _hidden_dim) = attention_output.dims3()?;
        let (_, fut_seq_len, _) = future_encoding.dims3()?;

        // Expand static context to match sequence length
        let total_seq = hist_seq_len + fut_seq_len;
        let static_expanded = static_encoding
            .squeeze(1)?  // [batch, hidden_dim]
            .unsqueeze(1)?  // [batch, 1, hidden_dim]
            .repeat(&[1, total_seq, 1])?;  // [batch, total_seq, hidden_dim]

        // Concatenate historical and future along time dimension
        let temporal_combined = Tensor::cat(&[attention_output, future_encoding], 1)?;

        // Add static context
        let combined = (&temporal_combined + &static_expanded)?;

        // Pool over time dimension to get fixed-size representation
        // Use mean pooling: [batch, seq, hidden] -> [batch, hidden]
        let pooled = combined.mean(1)?;

        Ok(pooled)
    }

    /// Step 6: Quantile Output Layer (placeholder)
    fn forward_quantile_output(&self, combined: &Tensor) -> Result<Tensor, MLError> {
        let batch_size = combined.dims()[0];

        // In production, apply GRN + linear layer for quantile predictions
        // For now, return zeros with correct shape
        let output = Tensor::zeros(
            (batch_size, self.config.prediction_horizon, self.config.num_quantiles),
            DType::F32,
            &self.device,
        )?;

        Ok(output)
    }

    /// Complete end-to-end INT8 forward pass
    pub fn forward_integrated(
        &self,
        static_features: &Tensor,
        historical_features: &Tensor,
        future_features: &Tensor,
    ) -> Result<Tensor, MLError> {
        // Validate inputs
        self.validate_inputs(static_features, historical_features, future_features)?;

        let batch_size = static_features.dims()[0];

        // Step 1: Static Variable Selection (FWD-01)
        let static_encoding = self.forward_static_vsn(static_features)?;

        // Step 2: Historical Encoder (FWD-02)
        let historical_encoding = self.forward_historical_encoder(historical_features)?;

        // Step 3: Future Decoder (FWD-03)
        let future_encoding = self.forward_future_decoder(future_features)?;

        // Step 4: Temporal Attention (FWD-04)
        let attention_output = self.forward_temporal_attention(&historical_encoding, false)?;

        // Step 5: Combine encodings
        let combined = self.combine_contexts(&static_encoding, &attention_output, &future_encoding)?;

        // Step 6: Quantile Output (FWD-05)
        let predictions = self.forward_quantile_output(&combined)?;

        // Verify output shape
        let expected_shape = vec![batch_size, self.config.prediction_horizon, self.config.num_quantiles];
        if predictions.dims() != expected_shape {
            return Err(MLError::ModelError(format!(
                "Output shape mismatch: expected {:?}, got {:?}",
                expected_shape,
                predictions.dims()
            )));
        }

        Ok(predictions)
    }
}