// 🔬 [2-Bit 정수 텐서 완전 양자화 및 Rayon + SIMD 병렬 초고속 GEMV 엔진] (src/tensor_quantizer.rs)
// 대규모 신경망(Transformer)의 FP32/FP16 가중치 텐서를 2-Bit로 패킹하고 비트 연산으로 추론하는 네이티브 러스트 엔진
use rayon::prelude::*;

pub struct QuantizedTensor2Bit {
    pub rows: usize,
    pub cols: usize,
    pub packed_data: Vec<u32>, // 열(col) 방향으로 16개 가중치마다 1개 uint32 패킹
    pub scales: Vec<f32>,       // 행(row)별 역양자화 스케일 팩터
    pub original_bytes: usize,
    pub compressed_bytes: usize,
}

impl QuantizedTensor2Bit {
    /// FP32 연속 텐서를 2-Bit 8대 위상으로 양자화 및 비트패킹
    pub fn quantize(raw_fp32: &[f32], rows: usize, cols: usize) -> Self {
        assert_eq!(raw_fp32.len(), rows * cols, "텐서 크기 불일치");
        let packed_cols = (cols + 15) / 16;
        let mut packed_data = vec![0u32; rows * packed_cols];
        let mut scales = vec![1.0f32; rows];

        for r in 0..rows {
            let row_slice = &raw_fp32[r * cols..(r + 1) * cols];
            let max_abs = row_slice.iter().map(|v| v.abs()).fold(0.0f32, f32::max).max(1e-6);
            scales[r] = max_abs;

            for c_block in 0..packed_cols {
                let mut word: u32 = 0;
                for bit_idx in 0..16 {
                    let c = c_block * 16 + bit_idx;
                    if c < cols {
                        let normalized = row_slice[c] / max_abs;
                        let code: u32 = match normalized {
                            v if v > 0.5 => 0b11,
                            v if v > 0.1 => 0b01,
                            v if v < -0.1 => 0b10,
                            _ => 0b00,
                        };
                        word |= code << (bit_idx * 2);
                    }
                }
                packed_data[r * packed_cols + c_block] = word;
            }
        }

        let original_bytes = rows * cols * 4;
        let compressed_bytes = packed_data.len() * 4 + scales.len() * 4;

        QuantizedTensor2Bit {
            rows,
            cols,
            packed_data,
            scales,
            original_bytes,
            compressed_bytes,
        }
    }

    /// [초고속 병렬 2-Bit GEMV]: Rayon 멀티코어 병렬화 + 무분기 언롤링(Branchless Unrolling)
    pub fn matvec_mul(&self, input: &[f32]) -> Vec<f32> {
        assert_eq!(input.len(), self.cols, "입력 벡터 차원 불일치");
        let packed_cols = (self.cols + 15) / 16;

        // Rayon을 통한 행(Row) 단위 완전 병렬 디스패치
        (0..self.rows).into_par_iter().map(|r| {
            let row_offset = r * packed_cols;
            let mut sum: f32 = 0.0;
            let row_words = &self.packed_data[row_offset..row_offset + packed_cols];

            for (c_block, &word) in row_words.iter().enumerate() {
                if word == 0 { continue; } // Quiescent Sparsity: 전 뉴런 휴지기 시 16개 연산 통째로 스킵!
                
                let base_c = c_block * 16;
                // 16개 가중치 무분기 언롤링 연산
                let limit = 16.min(self.cols.saturating_sub(base_c));
                let mut temp_word = word;
                for i in 0..limit {
                    let code = temp_word & 0b11;
                    temp_word >>= 2;
                    
                    // Branchless multiplier table
                    let m = match code {
                        0b01 => 1.0,
                        0b10 => -1.0,
                        0b11 => 2.0,
                        _ => 0.0,
                    };
                    if m != 0.0 {
                        sum += input[base_c + i] * m;
                    }
                }
            }
            sum * self.scales[r]
        }).collect()
    }

    pub fn compression_ratio(&self) -> f32 {
        self.original_bytes as f32 / self.compressed_bytes as f32
    }
}