torsh-quantization 0.2.0

//! Post-training quantization

use crate::{Observer, QScheme, QuantBackend, QuantConfig, TorshResult};
use std::collections::{HashMap, HashSet};
use torsh_core::{DType, TorshError};
// use torsh_nn::Module; // Temporarily disabled due to autograd compilation issues
use torsh_tensor::Tensor;

/// Sink invoked once per quantizable layer during a calibration forward pass.
///
/// The first argument is the layer name (matching the observer keys produced by
/// [`extract_layer_name`]), the second is that layer's own output activation.
pub type ActivationSink<'a> = dyn FnMut(&str, &Tensor) -> TorshResult<()> + 'a;

// Temporary placeholder trait for Module (to allow compilation)
#[allow(dead_code)]
pub trait Module {
    fn forward(&self, input: &Tensor) -> TorshResult<Tensor>;
    fn parameters(&self) -> Vec<&Tensor>;
    fn parameters_mut(&mut self) -> Vec<&mut Tensor>;
    fn named_parameters(&self) -> Vec<(String, &Tensor)>;
    fn train(&mut self, mode: bool);
    fn eval(&mut self) {
        self.train(false);
    }

    /// Forward pass that reports every quantizable layer's output activation.
    ///
    /// This is the calibration hook used by [`calibrate_model`]: an
    /// implementation must call `sink(layer_name, &activation)` for each layer
    /// it wants calibrated, passing **that layer's own output**, and return the
    /// model output.
    ///
    /// The default implementation reports no activations at all. Calibration
    /// then fails with an explicit error instead of silently filling every
    /// observer with the model input, which would produce quantization
    /// parameters derived from the wrong distribution.
    fn forward_with_activations(
        &self,
        input: &Tensor,
        sink: &mut ActivationSink<'_>,
    ) -> TorshResult<Tensor> {
        let _ = sink;
        self.forward(input)
    }
}

/// Calibration dataset for post-training quantization
#[derive(Debug)]
pub struct CalibrationDataset {
    pub data: Vec<Tensor>,
    pub batch_size: usize,
    pub num_batches: Option<usize>,
}

impl CalibrationDataset {
    /// Create a new calibration dataset
    pub fn new(data: Vec<Tensor>) -> Self {
        Self {
            batch_size: data.len(),
            num_batches: Some(1),
            data,
        }
    }

    /// Create dataset with specified batch size
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    /// Limit number of calibration batches
    pub fn with_num_batches(mut self, num_batches: usize) -> Self {
        self.num_batches = Some(num_batches);
        self
    }

    /// Get iterator over calibration batches
    pub fn iter(&self) -> impl Iterator<Item = &[Tensor]> {
        let max_batches = self.num_batches.unwrap_or(usize::MAX);
        self.data.chunks(self.batch_size).take(max_batches)
    }
}

/// Layer quantization parameters
#[derive(Debug, Clone)]
pub struct LayerQuantParams {
    pub scales: Vec<f32>,
    pub zero_points: Vec<i32>,
    pub scheme: QScheme,
    pub dtype: DType,
    pub ch_axis: Option<usize>,
}

/// Post-training quantization state
#[derive(Debug)]
pub struct PTQState {
    pub config: QuantConfig,
    pub observers: HashMap<String, Observer>,
    pub layer_params: HashMap<String, LayerQuantParams>,
    pub calibrated: bool,
    pub num_calibration_samples: usize,
}

impl PTQState {
    pub fn new(config: QuantConfig) -> Self {
        Self {
            config,
            observers: HashMap::new(),
            layer_params: HashMap::new(),
            calibrated: false,
            num_calibration_samples: 0,
        }
    }

    /// Add observer for a layer
    pub fn add_observer(&mut self, layer_name: String, observer: Observer) {
        self.observers.insert(layer_name, observer);
    }

    /// Get observer for a layer
    pub fn get_observer(&self, layer_name: &str) -> Option<&Observer> {
        self.observers.get(layer_name)
    }

    /// Get mutable observer for a layer
    pub fn get_observer_mut(&mut self, layer_name: &str) -> Option<&mut Observer> {
        self.observers.get_mut(layer_name)
    }

    /// Calculate quantization parameters for all layers
    pub fn calculate_all_qparams(&mut self) -> TorshResult<()> {
        for (layer_name, observer) in &self.observers {
            let (scale, zero_point) = observer.calculate_qparams(self.config.dtype)?;

            let layer_params = LayerQuantParams {
                scales: vec![scale],
                zero_points: vec![zero_point],
                scheme: self.config.scheme,
                dtype: self.config.dtype,
                ch_axis: self.config.ch_axis,
            };

            self.layer_params.insert(layer_name.clone(), layer_params);
        }

        self.calibrated = true;
        Ok(())
    }

    /// Get quantization parameters for a layer
    pub fn get_layer_params(&self, layer_name: &str) -> Option<&LayerQuantParams> {
        self.layer_params.get(layer_name)
    }

    /// Check if calibration is complete
    pub fn is_calibrated(&self) -> bool {
        self.calibrated && !self.layer_params.is_empty()
    }
}

/// Perform post-training quantization on a module
///
/// This function:
/// 1. Calibrates the model using representative data
/// 2. Computes quantization parameters
/// 3. Converts the model to use quantized operations
pub fn quantize_post_training(module: &mut dyn Module) -> TorshResult<()> {
    // Initialize PTQ state
    let mut ptq_state = PTQState::new(QuantConfig::default());

    // Put model in evaluation mode
    module.eval();

    // Attach observers to quantizable layers
    attach_observers(module, &mut ptq_state)?;

    // Note: In a real implementation, we would:
    // 1. Run calibration with representative data
    // 2. Collect statistics from observers
    // 3. Calculate quantization parameters
    // 4. Replace floating-point ops with quantized versions

    Ok(())
}

/// Calibrate the model with representative data
///
/// Each observer is updated with the activations reported by
/// [`Module::forward_with_activations`] for the layer it is attached to. A
/// model that does not implement that hook cannot be calibrated and this
/// function returns [`TorshError::NotImplemented`] rather than filling the
/// observers with the model input.
pub fn calibrate_model(
    module: &mut dyn Module,
    dataset: &CalibrationDataset,
    ptq_state: &mut PTQState,
) -> TorshResult<()> {
    // Validate configuration
    ptq_state.config.validate()?;

    if ptq_state.observers.is_empty() {
        return Err(TorshError::InvalidArgument(
            "No observers attached: nothing to calibrate".to_string(),
        ));
    }

    // Put model in evaluation mode
    module.eval();

    let mut processed_samples = 0;
    let mut observed_layers: HashSet<String> = HashSet::new();

    // Process each calibration batch
    for batch in dataset.iter() {
        for sample in batch {
            // Forward pass through the model. Every quantizable layer reports
            // its own output activation through the sink, so each observer sees
            // the distribution it is actually responsible for.
            {
                let observers = &mut ptq_state.observers;
                let seen = &mut observed_layers;
                let mut sink = |layer_name: &str, activation: &Tensor| -> TorshResult<()> {
                    if let Some(observer) = observers.get_mut(layer_name) {
                        observer.update(activation)?;
                        seen.insert(layer_name.to_string());
                    }
                    Ok(())
                };
                module.forward_with_activations(sample, &mut sink)?;
            }

            processed_samples += 1;
        }

        // Early stopping if we have enough samples
        if let Some(max_batches) = dataset.num_batches {
            if processed_samples >= max_batches * dataset.batch_size {
                break;
            }
        }
    }

    // Every attached observer must have received real activations, otherwise
    // its quantization parameters would be fabricated.
    let missing: Vec<&str> = ptq_state
        .observers
        .keys()
        .filter(|name| !observed_layers.contains(*name))
        .map(String::as_str)
        .collect();
    if !missing.is_empty() {
        return Err(TorshError::NotImplemented(format!(
            "Module::forward_with_activations reported no activations for layer(s): {}. \
             Implement the calibration hook so each observer receives its own layer's output.",
            missing.join(", ")
        )));
    }

    // Calculate quantization parameters from collected statistics
    ptq_state.calculate_all_qparams()?;
    ptq_state.num_calibration_samples = processed_samples;

    Ok(())
}

/// Attach observers to quantizable layers
fn attach_observers(module: &dyn Module, ptq_state: &mut PTQState) -> TorshResult<()> {
    let params = module.named_parameters();

    for (name, _param) in params {
        if is_quantizable_parameter(&name) {
            // Extract layer name from parameter name
            let layer_name = extract_layer_name(&name);
            if !ptq_state.observers.contains_key(&layer_name) {
                let observer = Observer::new(ptq_state.config.observer_type);
                ptq_state.observers.insert(layer_name, observer);
            }
        }
    }

    Ok(())
}

/// Check if a parameter belongs to a quantizable layer
fn is_quantizable_parameter(param_name: &str) -> bool {
    param_name.contains("linear")
        || param_name.contains("conv")
        || param_name.contains("Linear")
        || param_name.contains("Conv")
        || param_name.contains("batch_norm")
        || param_name.contains("BatchNorm")
}

/// Extract layer name from parameter name
fn extract_layer_name(param_name: &str) -> String {
    // Simple heuristic: take everything before ".weight" or ".bias"
    param_name
        .split(".weight")
        .next()
        .unwrap_or(param_name)
        .split(".bias")
        .next()
        .unwrap_or(param_name)
        .to_string()
}

/// Complete post-training quantization pipeline
pub fn ptq_pipeline(
    module: &mut dyn Module,
    dataset: &CalibrationDataset,
    config: QuantConfig,
) -> TorshResult<Box<dyn Module>> {
    // Initialize PTQ state
    let mut ptq_state = PTQState::new(config);

    // Attach observers to quantizable layers
    attach_observers(module, &mut ptq_state)?;

    // Calibrate with representative data
    calibrate_model(module, dataset, &mut ptq_state)?;

    // Convert to quantized model
    convert_to_quantized(module, &ptq_state)
}

/// Get quantization statistics for analysis
pub fn get_quantization_stats(ptq_state: &PTQState) -> HashMap<String, (f32, i32, usize)> {
    let mut stats = HashMap::new();

    for (layer_name, observer) in &ptq_state.observers {
        if let Ok((scale, zero_point)) = observer.calculate_qparams(ptq_state.config.dtype) {
            stats.insert(
                layer_name.clone(),
                (scale, zero_point, observer.num_batches()),
            );
        }
    }

    stats
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{ObserverType, QScheme, QuantBackend};
    use torsh_tensor::creation::tensor_1d;

    // Mock module for testing
    struct MockModule {
        weight: Tensor,
        bias: Tensor,
    }

    impl MockModule {
        fn new() -> Self {
            Self {
                weight: tensor_1d(&[1.0, 2.0, 3.0]).unwrap(),
                bias: tensor_1d(&[0.1, 0.2]).unwrap(),
            }
        }
    }

    impl Module for MockModule {
        fn forward(&self, input: &Tensor) -> TorshResult<Tensor> {
            Ok(input.clone())
        }

        fn forward_with_activations(
            &self,
            input: &Tensor,
            sink: &mut ActivationSink<'_>,
        ) -> TorshResult<Tensor> {
            // linear1 sees the input; linear2 sees linear1's (scaled) output.
            let first = input.clone();
            sink("linear1", &first)?;

            let scaled: Vec<f32> = first.data()?.iter().map(|&x| x * 2.0).collect();
            let second = Tensor::from_data(scaled, first.shape().dims().to_vec(), first.device())?;
            sink("linear2", &second)?;

            Ok(second)
        }

        fn parameters(&self) -> Vec<&Tensor> {
            vec![&self.weight, &self.bias]
        }

        fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
            vec![&mut self.weight, &mut self.bias]
        }

        fn named_parameters(&self) -> Vec<(String, &Tensor)> {
            vec![
                ("linear1.weight".to_string(), &self.weight),
                ("linear2.bias".to_string(), &self.bias),
            ]
        }

        fn train(&mut self, _mode: bool) {}
    }

    #[test]
    fn test_calibration_dataset() {
        let tensors = vec![
            tensor_1d(&[1.0, 2.0]).unwrap(),
            tensor_1d(&[3.0, 4.0]).unwrap(),
            tensor_1d(&[5.0, 6.0]).unwrap(),
            tensor_1d(&[7.0, 8.0]).unwrap(),
        ];

        let dataset = CalibrationDataset::new(tensors)
            .with_batch_size(2)
            .with_num_batches(2);

        assert_eq!(dataset.batch_size, 2);
        assert_eq!(dataset.num_batches, Some(2));

        let batches: Vec<_> = dataset.iter().collect();
        assert_eq!(batches.len(), 2);
        assert_eq!(batches[0].len(), 2);
        assert_eq!(batches[1].len(), 2);
    }

    #[test]
    fn test_ptq_state() {
        let config = QuantConfig::int8();
        let mut ptq_state = PTQState::new(config);

        assert!(!ptq_state.is_calibrated());
        assert_eq!(ptq_state.observers.len(), 0);
        assert_eq!(ptq_state.layer_params.len(), 0);

        // Add an observer
        let observer = Observer::new(ObserverType::MinMax);
        ptq_state.add_observer("layer1".to_string(), observer);

        assert_eq!(ptq_state.observers.len(), 1);
        assert!(ptq_state.get_observer("layer1").is_some());
        assert!(ptq_state.get_observer("nonexistent").is_none());
    }

    #[test]
    fn test_layer_quant_params() {
        let params = LayerQuantParams {
            scales: vec![0.1, 0.2],
            zero_points: vec![0, 1],
            scheme: QScheme::PerChannelAffine,
            dtype: DType::I8,
            ch_axis: Some(0),
        };

        assert_eq!(params.scales.len(), 2);
        assert_eq!(params.zero_points.len(), 2);
        assert_eq!(params.scheme, QScheme::PerChannelAffine);
    }

    #[test]
    fn test_dynamic_quant_config() {
        let config = DynamicQuantConfig::default();
        assert!(config.quantize_weights);
        assert!(config.quantize_activations);
        assert_eq!(config.dtype, DType::I8);
        assert!(config.layer_types.contains(&"linear".to_string()));
        assert!(config.layer_types.contains(&"conv".to_string()));
    }

    #[test]
    fn test_conversion_plan() {
        let mut plan = ConversionPlan::new(QuantBackend::Native);
        plan.add_layer("layer1".to_string());
        plan.add_quant_dequant("input".to_string(), "output".to_string());

        assert_eq!(plan.layers_to_quantize.len(), 1);
        assert_eq!(plan.insert_quant_dequant.len(), 1);
        assert_eq!(plan.backend, QuantBackend::Native);
    }

    #[test]
    fn test_extract_layer_name() {
        assert_eq!(extract_layer_name("layer1.weight"), "layer1");
        assert_eq!(extract_layer_name("model.layer2.bias"), "model.layer2");
        assert_eq!(extract_layer_name("simple_name"), "simple_name");
    }

    #[test]
    fn test_quantize_post_training() {
        let mut module = MockModule::new();
        let mut ptq_state = PTQState::new(QuantConfig::int8());

        // Test observer attachment
        assert!(attach_observers(&module, &mut ptq_state).is_ok());
        assert!(!ptq_state.observers.is_empty());

        // Test calibration dataset
        let tensors = vec![tensor_1d(&[1.0, 2.0, 3.0]).unwrap()];
        let dataset = CalibrationDataset::new(tensors);

        // Test calibration
        assert!(calibrate_model(&mut module, &dataset, &mut ptq_state).is_ok());
        assert!(ptq_state.is_calibrated());
    }

    #[test]
    fn test_quantization_stats() {
        let config = QuantConfig::int8();
        let mut ptq_state = PTQState::new(config);

        // Add a mock observer with some data
        let mut observer = Observer::new(ObserverType::MinMax);
        let tensor = tensor_1d(&[1.0, 2.0, 3.0]).unwrap();
        observer.update(&tensor).unwrap();

        ptq_state.add_observer("test_layer".to_string(), observer);

        let stats = get_quantization_stats(&ptq_state);
        assert!(stats.contains_key("test_layer"));

        let (scale, zero_point, num_batches) = stats["test_layer"];
        assert!(scale > 0.0);
        assert!((-128..=127).contains(&zero_point));
        assert_eq!(num_batches, 1);
    }

    #[test]
    fn test_validate_conversion_plan() {
        let config = QuantConfig::int8();
        let mut ptq_state = PTQState::new(config);

        // Add layer parameters
        let layer_params = LayerQuantParams {
            scales: vec![0.1],
            zero_points: vec![0],
            scheme: QScheme::PerTensorAffine,
            dtype: DType::I8,
            ch_axis: None,
        };
        ptq_state
            .layer_params
            .insert("layer1".to_string(), layer_params);
        ptq_state.calibrated = true;

        let mut plan = ConversionPlan::new(QuantBackend::Native);
        plan.add_layer("layer1".to_string());

        // Should validate successfully
        assert!(validate_conversion_plan(&plan, &ptq_state).is_ok());

        // Should fail for non-existent layer
        plan.add_layer("nonexistent".to_string());
        assert!(validate_conversion_plan(&plan, &ptq_state).is_err());
    }
}

/// Quantization conversion plan
#[derive(Debug)]
pub struct ConversionPlan {
    pub layers_to_quantize: Vec<String>,
    pub insert_quant_dequant: Vec<(String, String)>, // (before_layer, after_layer)
    pub backend: QuantBackend,
}

impl ConversionPlan {
    pub fn new(backend: QuantBackend) -> Self {
        Self {
            layers_to_quantize: Vec::new(),
            insert_quant_dequant: Vec::new(),
            backend,
        }
    }

    pub fn add_layer(&mut self, layer_name: String) {
        self.layers_to_quantize.push(layer_name);
    }

    pub fn add_quant_dequant(&mut self, before: String, after: String) {
        self.insert_quant_dequant.push((before, after));
    }
}

/// Convert calibrated model to quantized model
pub fn convert_to_quantized(
    _module: &dyn Module,
    ptq_state: &PTQState,
) -> TorshResult<Box<dyn Module>> {
    if !ptq_state.is_calibrated() {
        return Err(TorshError::Other(
            "Model must be calibrated before conversion".to_string(),
        ));
    }

    // Create conversion plan
    let plan = create_conversion_plan(ptq_state)?;

    // In a real implementation, this would:
    // 1. Create a new quantized model
    // 2. Replace floating-point layers with quantized equivalents
    // 3. Set quantization parameters from calibration
    // 4. Insert quantize/dequantize ops where needed

    validate_conversion_plan(&plan, ptq_state)?;

    Err(TorshError::Other(
        "PTQ conversion not yet implemented - plan validated".to_string(),
    ))
}

/// Create a conversion plan based on PTQ state
fn create_conversion_plan(ptq_state: &PTQState) -> TorshResult<ConversionPlan> {
    let mut plan = ConversionPlan::new(ptq_state.config.backend);

    // Add all calibrated layers to quantization plan
    for layer_name in ptq_state.layer_params.keys() {
        plan.add_layer(layer_name.clone());

        // Add quant/dequant around quantized layers
        plan.add_quant_dequant(
            format!("{layer_name}_input"),
            format!("{layer_name}_output"),
        );
    }

    Ok(plan)
}

/// Validate that the conversion plan is feasible
fn validate_conversion_plan(plan: &ConversionPlan, ptq_state: &PTQState) -> TorshResult<()> {
    // Check that all layers in plan have quantization parameters
    for layer_name in &plan.layers_to_quantize {
        if !ptq_state.layer_params.contains_key(layer_name) {
            return Err(TorshError::InvalidArgument(format!(
                "No quantization parameters for layer: {layer_name}"
            )));
        }
    }

    // Check backend compatibility
    match plan.backend {
        QuantBackend::Fbgemm | QuantBackend::Qnnpack => {
            // These backends have specific requirements
            if ptq_state.config.dtype != DType::I8 {
                return Err(TorshError::InvalidArgument(
                    "FBGEMM/QNNPACK backends require INT8 quantization".to_string(),
                ));
            }
        }
        QuantBackend::Native | QuantBackend::Xnnpack => {
            // More flexible backends
        }
    }

    Ok(())
}

/// Dynamic quantization configuration
#[derive(Debug)]
pub struct DynamicQuantConfig {
    pub quantize_weights: bool,
    pub quantize_activations: bool,
    pub dtype: DType,
    pub layer_types: Vec<String>,
}

impl Default for DynamicQuantConfig {
    fn default() -> Self {
        Self {
            quantize_weights: true,
            quantize_activations: true,
            dtype: DType::I8,
            layer_types: vec!["linear".to_string(), "conv".to_string()],
        }
    }
}

/// Quantize a pre-trained model with dynamic quantization
pub fn quantize_dynamic(module: &mut dyn Module) -> TorshResult<()> {
    quantize_dynamic_with_config(module, &DynamicQuantConfig::default())
}

/// Quantize with specific dynamic quantization configuration
pub fn quantize_dynamic_with_config(
    module: &mut dyn Module,
    config: &DynamicQuantConfig,
) -> TorshResult<()> {
    // Dynamic quantization doesn't require calibration
    // It computes quantization parameters on-the-fly

    let params = module.named_parameters();
    let mut quantized_layers = 0;

    for (name, _param) in params {
        let should_quantize = config
            .layer_types
            .iter()
            .any(|layer_type| name.to_lowercase().contains(&layer_type.to_lowercase()));

        if should_quantize {
            if config.quantize_weights {
                // For linear/conv layers, quantize weights statically
                // This would replace the parameter with a quantized version
                quantized_layers += 1;
            }

            if config.quantize_activations {
                // Insert dynamic quantization for activations
                // This computes quantization parameters on-the-fly during forward pass
            }
        }
    }

    if quantized_layers == 0 {
        return Err(TorshError::Other(
            "No layers found for dynamic quantization".to_string(),
        ));
    }

    Ok(())
}