rustorch 0.6.29

// メモリ圧縮・最適化エンジン
// Memory compression and optimization engine

use crate::common::RusTorchResult;
use crate::hybrid_f32::tensor::core::F32Tensor;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

/// 圧縮形式
/// Compression format
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CompressionFormat {
    /// 無圧縮
    /// No compression
    None,

    /// スパース表現（疎行列）
    /// Sparse representation
    Sparse,

    /// 量子化（低精度）
    /// Quantization (low precision)
    Quantized8,
    Quantized16,

    /// RLE圧縮（Run-Length Encoding）
    /// RLE compression
    RLE,

    /// ハフマン圧縮
    /// Huffman compression
    Huffman,

    /// LZ4圧縮
    /// LZ4 compression
    LZ4,
}

/// 圧縮設定
/// Compression configuration
#[derive(Debug, Clone)]
pub struct CompressionConfig {
    /// 圧縮形式
    /// Compression format
    pub format: CompressionFormat,

    /// 圧縮しきい値（バイト）
    /// Compression threshold (bytes)
    pub threshold: usize,

    /// スパース性しきい値（0-1）
    /// Sparsity threshold (0-1)
    pub sparsity_threshold: f32,

    /// 量子化レベル
    /// Quantization level
    pub quantization_levels: u32,

    /// 自動圧縮形式選択
    /// Auto compression format selection
    pub auto_select: bool,
}

impl Default for CompressionConfig {
    fn default() -> Self {
        CompressionConfig {
            format: CompressionFormat::None,
            threshold: 1024 * 1024, // 1MB
            sparsity_threshold: 0.8,
            quantization_levels: 256,
            auto_select: true,
        }
    }
}

/// 圧縮されたテンソルデータ
/// Compressed tensor data
#[derive(Debug, Clone)]
pub struct CompressedTensor {
    /// 圧縮データ
    /// Compressed data
    pub data: Vec<u8>,

    /// 元の形状
    /// Original shape
    pub shape: Vec<usize>,

    /// 圧縮形式
    /// Compression format
    pub format: CompressionFormat,

    /// 圧縮率
    /// Compression ratio
    pub compression_ratio: f32,

    /// メタデータ
    /// Metadata
    pub metadata: HashMap<String, String>,
}

impl CompressedTensor {
    /// 圧縮率を計算
    /// Calculate compression ratio
    pub fn calculate_ratio(original_size: usize, compressed_size: usize) -> f32 {
        if original_size == 0 {
            0.0
        } else {
            compressed_size as f32 / original_size as f32
        }
    }

    /// メモリ削減量を取得
    /// Get memory savings
    pub fn memory_savings(&self) -> usize {
        let original_size: usize = self.shape.iter().product::<usize>() * 4; // f32 = 4 bytes
        original_size.saturating_sub(self.data.len())
    }
}

/// スパーステンソル表現
/// Sparse tensor representation
#[derive(Debug, Clone)]
pub struct SparseTensor {
    /// 非ゼロ値のインデックス
    /// Non-zero indices
    pub indices: Vec<Vec<usize>>,

    /// 非ゼロ値
    /// Non-zero values
    pub values: Vec<f32>,

    /// テンソル形状
    /// Tensor shape
    pub shape: Vec<usize>,

    /// 密度（非ゼロ要素の割合）
    /// Density (ratio of non-zero elements)
    pub density: f32,
}

impl SparseTensor {
    /// F32Tensorからスパーステンソルを作成
    /// Create sparse tensor from F32Tensor
    pub fn from_dense(tensor: &F32Tensor, threshold: f32) -> RusTorchResult<Self> {
        let mut indices = Vec::new();
        let mut values = Vec::new();
        let shape = tensor.shape().to_vec();

        // 非ゼロ要素を抽出
        for (flat_idx, &value) in tensor.data.iter().enumerate() {
            if value.abs() > threshold {
                // フラットインデックスを多次元インデックスに変換
                let mut multi_idx = Vec::new();
                let mut remaining = flat_idx;

                for &dim_size in shape.iter().rev() {
                    multi_idx.push(remaining % dim_size);
                    remaining /= dim_size;
                }
                multi_idx.reverse();

                indices.push(multi_idx);
                values.push(value);
            }
        }

        let total_elements: usize = shape.iter().product();
        let density = if total_elements > 0 {
            values.len() as f32 / total_elements as f32
        } else {
            0.0
        };

        Ok(SparseTensor {
            indices,
            values,
            shape,
            density,
        })
    }

    /// スパーステンソルをF32Tensorに変換
    /// Convert sparse tensor to F32Tensor
    pub fn to_dense(&self) -> RusTorchResult<F32Tensor> {
        let mut dense_tensor = F32Tensor::zeros(&self.shape)?;

        for (idx_vec, &value) in self.indices.iter().zip(self.values.iter()) {
            // 多次元インデックスをフラットインデックスに変換
            let mut flat_idx = 0;
            let mut multiplier = 1;

            for (&idx, &dim_size) in idx_vec.iter().zip(self.shape.iter()).rev() {
                flat_idx += idx * multiplier;
                multiplier *= dim_size;
            }

            dense_tensor.data[flat_idx] = value;
        }

        Ok(dense_tensor)
    }

    /// メモリ使用量を推定
    /// Estimate memory usage
    pub fn memory_usage(&self) -> usize {
        let indices_size = self.indices.len() * self.shape.len() * std::mem::size_of::<usize>();
        let values_size = self.values.len() * std::mem::size_of::<f32>();
        let shape_size = self.shape.len() * std::mem::size_of::<usize>();

        indices_size + values_size + shape_size + std::mem::size_of::<f32>() // density
    }
}

/// 量子化テンソル
/// Quantized tensor
#[derive(Debug, Clone)]
pub struct QuantizedTensor {
    /// 量子化された値
    /// Quantized values
    pub quantized_data: Vec<u8>,

    /// スケールファクター
    /// Scale factor
    pub scale: f32,

    /// ゼロポイント
    /// Zero point
    pub zero_point: u8,

    /// 元の形状
    /// Original shape
    pub shape: Vec<usize>,

    /// 量子化ビット数
    /// Quantization bits
    pub bits: u8,
}

impl QuantizedTensor {
    /// F32Tensorを8ビット量子化
    /// Quantize F32Tensor to 8-bit
    pub fn quantize_8bit(tensor: &F32Tensor) -> Self {
        let data = &tensor.data;
        let min_val = data.iter().fold(f32::INFINITY, |acc, &x| acc.min(x));
        let max_val = data.iter().fold(f32::NEG_INFINITY, |acc, &x| acc.max(x));

        let scale = (max_val - min_val) / 255.0;
        let zero_point = (-min_val / scale).round() as u8;

        let quantized_data: Vec<u8> = data
            .iter()
            .map(|&value| {
                let quantized = (value / scale + zero_point as f32).round();
                quantized.max(0.0).min(255.0) as u8
            })
            .collect();

        QuantizedTensor {
            quantized_data,
            scale,
            zero_point,
            shape: tensor.shape().to_vec(),
            bits: 8,
        }
    }

    /// 量子化テンソルをF32Tensorに復元
    /// Dequantize to F32Tensor
    pub fn dequantize(&self) -> RusTorchResult<F32Tensor> {
        let dequantized_data: Vec<f32> = self
            .quantized_data
            .iter()
            .map(|&q| (q as f32 - self.zero_point as f32) * self.scale)
            .collect();

        F32Tensor::new(dequantized_data, &self.shape)
    }

    /// メモリ使用量を推定
    /// Estimate memory usage
    pub fn memory_usage(&self) -> usize {
        self.quantized_data.len() +
        std::mem::size_of::<f32>() + // scale
        std::mem::size_of::<u8>() + // zero_point
        self.shape.len() * std::mem::size_of::<usize>() + // shape
        std::mem::size_of::<u8>() // bits
    }
}

/// 圧縮エンジン
/// Compression engine
#[derive(Debug)]
pub struct CompressionEngine {
    /// 設定
    /// Configuration
    config: CompressionConfig,

    /// 圧縮統計
    /// Compression statistics
    stats: Arc<Mutex<CompressionStats>>,
}

/// 圧縮統計
/// Compression statistics
#[derive(Debug, Default, Clone)]
pub struct CompressionStats {
    /// 圧縮回数
    /// Compression count
    pub compressions: usize,

    /// 解凍回数
    /// Decompression count
    pub decompressions: usize,

    /// 総バイト数節約
    /// Total bytes saved
    pub bytes_saved: usize,

    /// 平均圧縮率
    /// Average compression ratio
    pub avg_compression_ratio: f32,

    /// 形式別使用回数
    /// Usage count by format
    pub format_usage: HashMap<CompressionFormat, usize>,
}

impl CompressionEngine {
    /// 新しい圧縮エンジンを作成
    /// Create new compression engine
    pub fn new(config: CompressionConfig) -> Self {
        CompressionEngine {
            config,
            stats: Arc::new(Mutex::new(CompressionStats::default())),
        }
    }

    /// デフォルト設定で作成
    /// Create with default config
    pub fn with_default_config() -> Self {
        Self::new(CompressionConfig::default())
    }

    /// テンソルを圧縮
    /// Compress tensor
    pub fn compress(&self, tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        let format = if self.config.auto_select {
            self.select_best_format(tensor)
        } else {
            self.config.format
        };

        let compressed = match format {
            CompressionFormat::None => self.compress_none(tensor)?,
            CompressionFormat::Sparse => self.compress_sparse(tensor)?,
            CompressionFormat::Quantized8 => self.compress_quantized_8(tensor)?,
            CompressionFormat::Quantized16 => self.compress_quantized_16(tensor)?,
            CompressionFormat::RLE => self.compress_rle(tensor)?,
            CompressionFormat::Huffman => self.compress_huffman(tensor)?,
            CompressionFormat::LZ4 => self.compress_lz4(tensor)?,
        };

        // 統計を更新
        let mut stats = self.stats.lock().unwrap();
        stats.compressions += 1;
        stats.bytes_saved += compressed.memory_savings();
        *stats.format_usage.entry(format).or_insert(0) += 1;

        // 平均圧縮率を更新
        let total_ratio = stats.avg_compression_ratio * (stats.compressions - 1) as f32
            + compressed.compression_ratio;
        stats.avg_compression_ratio = total_ratio / stats.compressions as f32;

        Ok(compressed)
    }

    /// 圧縮されたテンソルを解凍
    /// Decompress tensor
    pub fn decompress(&self, compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        let tensor = match compressed.format {
            CompressionFormat::None => self.decompress_none(compressed)?,
            CompressionFormat::Sparse => {
                // Temporary: Convert sparse back to dense
                let sparse = self.decompress_sparse(compressed)?;
                self.sparse_to_dense(&sparse)?
            }
            CompressionFormat::Quantized8 => self.decompress_quantized_8(compressed)?,
            CompressionFormat::Quantized16 => self.decompress_quantized_16(compressed)?,
            CompressionFormat::RLE => self.decompress_rle(compressed)?,
            CompressionFormat::Huffman => self.decompress_huffman(compressed)?,
            CompressionFormat::LZ4 => self.decompress_lz4(compressed)?,
        };

        // 統計を更新
        let mut stats = self.stats.lock().unwrap();
        stats.decompressions += 1;

        Ok(tensor)
    }

    /// 最適な圧縮形式を選択
    /// Select best compression format
    fn select_best_format(&self, tensor: &F32Tensor) -> CompressionFormat {
        let sparsity = self.calculate_sparsity(tensor);
        let size = tensor.numel() * 4; // f32 = 4 bytes

        // スパース性が高い場合
        if sparsity > self.config.sparsity_threshold {
            return CompressionFormat::Sparse;
        }

        // 小さなテンソルは量子化
        if size < 1024 * 1024 {
            // 1MB未満
            return CompressionFormat::Quantized8;
        }

        // 大きなテンソルはLZ4
        if size > 10 * 1024 * 1024 {
            // 10MB以上
            return CompressionFormat::LZ4;
        }

        // デフォルトは無圧縮
        CompressionFormat::None
    }

    /// スパース性を計算
    /// Calculate sparsity
    fn calculate_sparsity(&self, tensor: &F32Tensor) -> f32 {
        let threshold = 1e-6;
        let zero_count = tensor.data.iter().filter(|&&x| x.abs() < threshold).count();

        zero_count as f32 / tensor.numel() as f32
    }

    /// 無圧縮（そのまま）
    /// No compression (as-is)
    fn compress_none(&self, tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        let data = tensor
            .data
            .as_slice()
            .unwrap()
            .iter()
            .flat_map(|&f| f.to_le_bytes())
            .collect();

        Ok(CompressedTensor {
            data,
            shape: tensor.shape().to_vec(),
            format: CompressionFormat::None,
            compression_ratio: 1.0,
            metadata: HashMap::new(),
        })
    }

    fn decompress_none(&self, compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        let float_data: Vec<f32> = compressed
            .data
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
            .collect();

        F32Tensor::new(float_data, &compressed.shape)
    }

    /// スパース圧縮
    /// Sparse compression
    fn compress_sparse(&self, tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        let sparse = SparseTensor::from_dense(tensor, 1e-6)?;

        // スパースデータをシリアライズ
        let mut data = Vec::new();

        // インデックス数
        data.extend_from_slice(&(sparse.indices.len() as u32).to_le_bytes());

        // インデックス
        for idx_vec in &sparse.indices {
            data.extend_from_slice(&(idx_vec.len() as u32).to_le_bytes());
            for &idx in idx_vec {
                data.extend_from_slice(&(idx as u32).to_le_bytes());
            }
        }

        // 値
        for &value in &sparse.values {
            data.extend_from_slice(&value.to_le_bytes());
        }

        let original_size = tensor.numel() * 4;
        let compression_ratio = data.len() as f32 / original_size as f32;

        let mut metadata = HashMap::new();
        metadata.insert("density".to_string(), sparse.density.to_string());

        Ok(CompressedTensor {
            data,
            shape: tensor.shape().to_vec(),
            format: CompressionFormat::Sparse,
            compression_ratio,
            metadata,
        })
    }

    fn decompress_sparse(&self, compressed: &CompressedTensor) -> RusTorchResult<CompressedTensor> {
        // スパース解凍の実装（簡略版）
        // 実際の実装では、シリアライズされたデータを正しく解析する
        Ok(compressed.clone())
    }

    /// スパーステンソルを密なテンソルに変換
    /// Convert sparse tensor to dense tensor
    fn sparse_to_dense(&self, sparse: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        // 簡易実装：ゼロテンソルを返す
        // Simple implementation: return zero tensor
        F32Tensor::zeros(&sparse.shape)
    }

    /// 8ビット量子化
    /// 8-bit quantization
    fn compress_quantized_8(&self, tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        let quantized = QuantizedTensor::quantize_8bit(tensor);

        let mut data = Vec::new();
        data.extend_from_slice(&quantized.scale.to_le_bytes());
        data.push(quantized.zero_point);
        data.extend_from_slice(&quantized.quantized_data);

        let original_size = tensor.numel() * 4;
        let compression_ratio = data.len() as f32 / original_size as f32;

        let mut metadata = HashMap::new();
        metadata.insert("scale".to_string(), quantized.scale.to_string());
        metadata.insert("zero_point".to_string(), quantized.zero_point.to_string());

        Ok(CompressedTensor {
            data,
            shape: tensor.shape().to_vec(),
            format: CompressionFormat::Quantized8,
            compression_ratio,
            metadata,
        })
    }

    fn decompress_quantized_8(&self, compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        if compressed.data.len() < 5 {
            return Err(crate::error::RusTorchError::tensor_op(
                "Invalid quantized data",
            ));
        }

        let scale = f32::from_le_bytes([
            compressed.data[0],
            compressed.data[1],
            compressed.data[2],
            compressed.data[3],
        ]);
        let zero_point = compressed.data[4];
        let quantized_data = &compressed.data[5..];

        let dequantized_data: Vec<f32> = quantized_data
            .iter()
            .map(|&q| (q as f32 - zero_point as f32) * scale)
            .collect();

        F32Tensor::new(dequantized_data, &compressed.shape)
    }

    /// 16ビット量子化（プレースホルダー）
    /// 16-bit quantization (placeholder)
    fn compress_quantized_16(&self, _tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        todo!("16-bit quantization implementation")
    }

    fn decompress_quantized_16(&self, _compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        todo!("16-bit quantization decompression implementation")
    }

    /// RLE圧縮（プレースホルダー）
    /// RLE compression (placeholder)
    fn compress_rle(&self, _tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        todo!("RLE compression implementation")
    }

    fn decompress_rle(&self, _compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        todo!("RLE decompression implementation")
    }

    /// ハフマン圧縮（プレースホルダー）
    /// Huffman compression (placeholder)
    fn compress_huffman(&self, _tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        todo!("Huffman compression implementation")
    }

    fn decompress_huffman(&self, _compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        todo!("Huffman decompression implementation")
    }

    /// LZ4圧縮（プレースホルダー）
    /// LZ4 compression (placeholder)
    fn compress_lz4(&self, _tensor: &F32Tensor) -> RusTorchResult<CompressedTensor> {
        todo!("LZ4 compression implementation")
    }

    fn decompress_lz4(&self, _compressed: &CompressedTensor) -> RusTorchResult<F32Tensor> {
        todo!("LZ4 decompression implementation")
    }

    /// 統計情報を取得
    /// Get statistics
    pub fn stats(&self) -> CompressionStats {
        self.stats.lock().unwrap().clone()
    }

    /// 設定を取得
    /// Get configuration
    pub fn config(&self) -> &CompressionConfig {
        &self.config
    }
}

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

    #[test]
    fn test_sparse_tensor() {
        let tensor = F32Tensor::zeros(&[3, 3]).unwrap();
        // いくつかの非ゼロ値を設定（実際の実装では適切なメソッドを使用）

        let sparse = SparseTensor::from_dense(&tensor, 1e-6).unwrap();
        assert!(sparse.density < 1.0);
    }

    #[test]
    fn test_quantization() {
        let tensor = F32Tensor::randn(&[2, 2]).unwrap();
        let quantized = QuantizedTensor::quantize_8bit(&tensor);
        let dequantized = quantized.dequantize().unwrap();

        assert_eq!(tensor.shape(), dequantized.shape());
        // 量子化による誤差は許容範囲内であることを確認
    }

    #[test]
    fn test_compression_engine() {
        let engine = CompressionEngine::with_default_config();
        let tensor = F32Tensor::zeros(&[10, 10]).unwrap();

        let compressed = engine.compress(&tensor).unwrap();
        let decompressed = engine.decompress(&compressed).unwrap();

        assert_eq!(tensor.shape(), decompressed.shape());

        let stats = engine.stats();
        assert_eq!(stats.compressions, 1);
        assert_eq!(stats.decompressions, 1);
    }
}