//! Multimodal tokenization for images and audio
//!
//! This module provides comprehensive tokenization capabilities for non-text modalities,
//! enabling unified representation across text, image, and audio data with real processing algorithms.

use crate::{Result, TokenId, TokenizationError, TokenizationStats, Tokenizer};
use image::DynamicImage;
use parking_lot::RwLock;
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

/// Supported modality types
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModalityType {
    /// Text modality
    Text,
    /// Image modality
    Image,
    /// Audio modality
    Audio,
    /// Video modality (future extension)
    Video,
}

/// Image tokenization method
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum ImageTokenizationMethod {
    /// Patch-based tokenization (like Vision Transformer)
    Patches,
    /// CNN-style convolutional features
    ConvolutionalFeatures,
    /// Hierarchical patches (multi-scale)
    HierarchicalPatches,
    /// CLIP-style visual features
    VisualFeatures,
}

/// Audio tokenization method
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AudioTokenizationMethod {
    /// Mel-frequency cepstral coefficients
    MFCC,
    /// Spectogram-based features
    Spectrogram,
    /// Raw waveform quantization
    WaveformQuantization,
    /// Audio patches (like `AudioMAE`)
    AudioPatches,
}

/// Multimodal tokenizer configuration
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MultimodalConfig {
    /// Supported modalities
    pub modalities: Vec<ModalityType>,
    /// Vocabulary size per modality
    pub vocab_size_per_modality: usize,
    /// Image tokenization configuration
    pub image_config: ImageTokenizationConfig,
    /// Audio tokenization configuration
    pub audio_config: AudioTokenizationConfig,
    /// Cross-modal alignment settings
    pub cross_modal_alignment: bool,
    /// Global token offset per modality
    pub modality_token_offset: HashMap<ModalityType, TokenId>,
}

/// Image tokenization configuration
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ImageTokenizationConfig {
    /// Tokenization method
    pub method: ImageTokenizationMethod,
    /// Image patch size (width, height)
    pub patch_size: (u32, u32),
    /// Target image resolution (width, height)
    pub target_resolution: (u32, u32),
    /// Number of color channels (3 for RGB, 1 for grayscale)
    pub channels: u32,
    /// Normalization parameters (mean, std) scaled by 10000 to avoid f32
    pub normalization: Option<([u32; 3], [u32; 3])>,
    /// Whether to use overlap between patches
    pub patch_overlap: bool,
    /// Overlap stride if enabled
    pub overlap_stride: u32,
}

impl Default for ImageTokenizationConfig {
    fn default() -> Self {
        Self {
            method: ImageTokenizationMethod::Patches,
            patch_size: (16, 16),
            target_resolution: (224, 224),
            channels: 3,
            // ImageNet stats scaled by 10000: mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
            normalization: Some(([4850, 4560, 4060], [2290, 2240, 2250])),
            patch_overlap: false,
            overlap_stride: 8,
        }
    }
}

impl ImageTokenizationConfig {
    /// Get normalization parameters as f32 tuples
    #[must_use]
    pub fn normalization_f32(&self) -> Option<([f32; 3], [f32; 3])> {
        self.normalization.map(|(mean, std)| {
            let mean_f32 = [
                mean[0] as f32 / 10000.0,
                mean[1] as f32 / 10000.0,
                mean[2] as f32 / 10000.0,
            ];
            let std_f32 = [
                std[0] as f32 / 10000.0,
                std[1] as f32 / 10000.0,
                std[2] as f32 / 10000.0,
            ];
            (mean_f32, std_f32)
        })
    }
}

/// Audio tokenization configuration
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AudioTokenizationConfig {
    /// Tokenization method
    pub method: AudioTokenizationMethod,
    /// Sample rate for audio processing
    pub sample_rate: u32,
    /// Window size for spectral analysis
    pub window_size: usize,
    /// Hop size for windowing
    pub hop_size: usize,
    /// Number of mel filter banks
    pub n_mel_filters: usize,
    /// Number of MFCC coefficients
    pub n_mfcc: usize,
    /// Frame length in microseconds (to avoid f32)
    pub frame_length_microseconds: u32,
    /// Whether to apply pre-emphasis
    pub pre_emphasis: bool,
    /// Pre-emphasis coefficient as integer (multiplied by 1000)
    pub pre_emphasis_coeff_scaled: u32,
}

impl Default for AudioTokenizationConfig {
    fn default() -> Self {
        Self {
            method: AudioTokenizationMethod::MFCC,
            sample_rate: 16000,
            window_size: 512,
            hop_size: 256,
            n_mel_filters: 80,
            n_mfcc: 13,
            frame_length_microseconds: 25000, // 0.025 seconds = 25000 microseconds
            pre_emphasis: true,
            pre_emphasis_coeff_scaled: 970, // 0.97 * 1000
        }
    }
}

impl AudioTokenizationConfig {
    /// Get frame length in seconds
    #[must_use]
    pub fn frame_length_seconds(&self) -> f32 {
        self.frame_length_microseconds as f32 / 1_000_000.0
    }

    /// Get pre-emphasis coefficient as f32
    #[must_use]
    pub fn pre_emphasis_coeff(&self) -> f32 {
        self.pre_emphasis_coeff_scaled as f32 / 1000.0
    }
}

impl Default for MultimodalConfig {
    fn default() -> Self {
        let mut modality_token_offset = HashMap::new();
        modality_token_offset.insert(ModalityType::Text, 0);
        modality_token_offset.insert(ModalityType::Image, 10000);
        modality_token_offset.insert(ModalityType::Audio, 20000);

        Self {
            modalities: vec![ModalityType::Text, ModalityType::Image, ModalityType::Audio],
            vocab_size_per_modality: 10000,
            image_config: ImageTokenizationConfig::default(),
            audio_config: AudioTokenizationConfig::default(),
            cross_modal_alignment: true,
            modality_token_offset,
        }
    }
}

/// Image patch representation
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ImagePatch {
    /// Patch data as flattened RGB values
    pub data: Vec<f32>,
    /// Patch position (row, column)
    pub position: (usize, usize),
    /// Patch dimensions
    pub dimensions: (u32, u32),
}

/// Audio frame representation
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct AudioFrame {
    /// Audio features (MFCC, spectral, etc.)
    pub features: Vec<f32>,
    /// Frame timestamp in seconds
    pub timestamp: f32,
    /// Frame duration in seconds
    pub duration: f32,
}

/// Multimodal tokenizer with real processing capabilities
#[derive(Debug)]
pub struct MultimodalTokenizer {
    config: MultimodalConfig,
    /// Pre-computed image patch codebook
    image_codebook: Arc<RwLock<Vec<ImagePatch>>>,
    /// Pre-computed audio feature codebook
    audio_codebook: Arc<RwLock<Vec<AudioFrame>>>,
    /// Cross-modal alignment matrix
    alignment_matrix: Arc<RwLock<Option<nalgebra::DMatrix<f32>>>>,
    /// Statistics
    stats: Arc<RwLock<TokenizationStats>>,
}

impl MultimodalTokenizer {
    /// Create new multimodal tokenizer
    #[must_use]
    pub fn new(config: MultimodalConfig) -> Self {
        Self {
            config,
            image_codebook: Arc::new(RwLock::new(Vec::new())),
            audio_codebook: Arc::new(RwLock::new(Vec::new())),
            alignment_matrix: Arc::new(RwLock::new(None)),
            stats: Arc::new(RwLock::new(TokenizationStats::default())),
        }
    }

    /// Train codebooks from multimodal data
    pub async fn train_codebooks(
        &mut self,
        images: &[Vec<u8>],
        audio_samples: &[Vec<f32>],
    ) -> Result<()> {
        // Train image codebook
        if !images.is_empty() {
            let image_patches = self.extract_all_image_patches(images).await?;
            let image_codebook = self.build_image_codebook(image_patches).await?;
            *self.image_codebook.write() = image_codebook;
        }

        // Train audio codebook
        if !audio_samples.is_empty() {
            let audio_features = self.extract_all_audio_features(audio_samples).await?;
            let audio_codebook = self.build_audio_codebook(audio_features).await?;
            *self.audio_codebook.write() = audio_codebook;
        }

        // Build cross-modal alignment if enabled
        if self.config.cross_modal_alignment {
            self.build_cross_modal_alignment().await?;
        }

        Ok(())
    }

    /// Tokenize image data with real patch extraction and processing
    pub async fn encode_image(&self, image_data: &[u8]) -> Result<Vec<TokenId>> {
        let start_time = std::time::Instant::now();

        // Load and decode image
        let image = image::load_from_memory(image_data).map_err(|e| {
            TokenizationError::MultimodalError(format!("Failed to load image: {e}"))
        })?;

        // Resize to target resolution
        let resized_image = image.resize_exact(
            self.config.image_config.target_resolution.0,
            self.config.image_config.target_resolution.1,
            image::imageops::FilterType::Lanczos3,
        );

        // Extract patches
        let patches = self.extract_image_patches(&resized_image).await?;

        // Quantize patches to tokens
        let token_ids = self.quantize_image_patches(patches).await?;

        // Add modality offset
        let offset = *self
            .config
            .modality_token_offset
            .get(&ModalityType::Image)
            .unwrap_or(&10000);
        let final_tokens: Vec<TokenId> = token_ids.into_iter().map(|id| id + offset).collect();

        // Update statistics
        {
            let mut stats = self.stats.write();
            stats.processing_time_ms += start_time.elapsed().as_millis() as u64;
        }

        Ok(final_tokens)
    }

    /// Tokenize audio data with real feature extraction
    pub async fn encode_audio(&self, audio_data: &[f32]) -> Result<Vec<TokenId>> {
        let start_time = std::time::Instant::now();

        // Extract audio features based on configuration
        let features = match self.config.audio_config.method {
            AudioTokenizationMethod::MFCC => self.extract_mfcc_features(audio_data).await?,
            AudioTokenizationMethod::Spectrogram => {
                self.extract_spectrogram_features(audio_data).await?
            }
            AudioTokenizationMethod::WaveformQuantization => {
                self.quantize_waveform(audio_data).await?
            }
            AudioTokenizationMethod::AudioPatches => self.extract_audio_patches(audio_data).await?,
        };

        // Quantize features to tokens
        let token_ids = self.quantize_audio_features(features).await?;

        // Add modality offset
        let offset = *self
            .config
            .modality_token_offset
            .get(&ModalityType::Audio)
            .unwrap_or(&20000);
        let final_tokens: Vec<TokenId> = token_ids.into_iter().map(|id| id + offset).collect();

        // Update statistics
        {
            let mut stats = self.stats.write();
            stats.processing_time_ms += start_time.elapsed().as_millis() as u64;
        }

        Ok(final_tokens)
    }

    /// Extract image patches with real computer vision processing
    async fn extract_image_patches(&self, image: &DynamicImage) -> Result<Vec<ImagePatch>> {
        let rgb_image = image.to_rgb8();
        let (width, height) = rgb_image.dimensions();
        let (patch_width, patch_height) = self.config.image_config.patch_size;

        let _patches: Vec<ImagePatch> = Vec::new();

        let stride_x = if self.config.image_config.patch_overlap {
            self.config.image_config.overlap_stride
        } else {
            patch_width
        };
        let stride_y = if self.config.image_config.patch_overlap {
            self.config.image_config.overlap_stride
        } else {
            patch_height
        };

        // Extract patches in parallel
        let patch_positions: Vec<(usize, usize)> = (0..height)
            .step_by(stride_y as usize)
            .flat_map(|y| {
                (0..width)
                    .step_by(stride_x as usize)
                    .map(move |x| (y as usize, x as usize))
            })
            .collect();

        let rgb_data = rgb_image.as_raw();

        let extracted_patches: Vec<ImagePatch> = patch_positions
            .par_iter()
            .filter_map(|&(y, x)| {
                if x + patch_width as usize <= width as usize
                    && y + patch_height as usize <= height as usize
                {
                    let mut patch_data =
                        Vec::with_capacity((patch_width * patch_height * 3) as usize);

                    // Extract patch pixel data
                    for py in y..y + patch_height as usize {
                        for px in x..x + patch_width as usize {
                            let pixel_idx = (py * width as usize + px) * 3;
                            if pixel_idx + 2 < rgb_data.len() {
                                patch_data.push(f32::from(rgb_data[pixel_idx]) / 255.0); // R
                                patch_data.push(f32::from(rgb_data[pixel_idx + 1]) / 255.0); // G
                                patch_data.push(f32::from(rgb_data[pixel_idx + 2]) / 255.0); // B
                            }
                        }
                    }

                    // Apply normalization if configured
                    if let Some((mean, std)) = self.config.image_config.normalization_f32() {
                        for (i, value) in patch_data.iter_mut().enumerate() {
                            let channel = i % 3;
                            *value = (*value - mean[channel]) / std[channel];
                        }
                    }

                    Some(ImagePatch {
                        data: patch_data,
                        position: (y / stride_y as usize, x / stride_x as usize),
                        dimensions: (patch_width, patch_height),
                    })
                } else {
                    None
                }
            })
            .collect();

        Ok(extracted_patches)
    }

    /// Extract MFCC features from audio
    async fn extract_mfcc_features(&self, audio_data: &[f32]) -> Result<Vec<AudioFrame>> {
        let window_size = self.config.audio_config.window_size;
        let hop_size = self.config.audio_config.hop_size;
        let sample_rate = self.config.audio_config.sample_rate;
        let n_mfcc = self.config.audio_config.n_mfcc;

        let mut frames = Vec::new();
        let mut processed_audio = audio_data.to_vec();

        // Apply pre-emphasis if enabled
        if self.config.audio_config.pre_emphasis {
            let coeff = self.config.audio_config.pre_emphasis_coeff();
            for i in (1..processed_audio.len()).rev() {
                processed_audio[i] -= coeff * processed_audio[i - 1];
            }
        }

        // Extract frames
        for (frame_idx, frame_start) in (0..processed_audio.len()).step_by(hop_size).enumerate() {
            if frame_start + window_size <= processed_audio.len() {
                let frame = &processed_audio[frame_start..frame_start + window_size];

                // Apply Hamming window
                let windowed_frame: Vec<f32> = frame
                    .iter()
                    .enumerate()
                    .map(|(i, &sample)| {
                        sample
                            * (0.54
                                - 0.46
                                    * (2.0 * std::f32::consts::PI * i as f32
                                        / (window_size - 1) as f32)
                                        .cos())
                    })
                    .collect();

                // Compute FFT (simplified - in production use proper FFT library)
                let fft_magnitudes = self.compute_fft_magnitudes(&windowed_frame);

                // Convert to mel scale
                let mel_features = self.apply_mel_filterbank(&fft_magnitudes, sample_rate);

                // Compute DCT to get MFCC
                let mfcc_features = self.compute_dct(&mel_features, n_mfcc);

                frames.push(AudioFrame {
                    features: mfcc_features,
                    timestamp: frame_idx as f32 * hop_size as f32 / sample_rate as f32,
                    duration: window_size as f32 / sample_rate as f32,
                });
            }
        }

        Ok(frames)
    }

    /// Extract spectrogram features from audio
    async fn extract_spectrogram_features(&self, audio_data: &[f32]) -> Result<Vec<AudioFrame>> {
        let window_size = self.config.audio_config.window_size;
        let hop_size = self.config.audio_config.hop_size;
        let sample_rate = self.config.audio_config.sample_rate;

        let mut frames = Vec::new();

        // Extract spectral frames
        for (frame_idx, frame_start) in (0..audio_data.len()).step_by(hop_size).enumerate() {
            if frame_start + window_size <= audio_data.len() {
                let frame = &audio_data[frame_start..frame_start + window_size];

                // Apply Hamming window
                let windowed_frame: Vec<f32> = frame
                    .iter()
                    .enumerate()
                    .map(|(i, &sample)| {
                        sample
                            * (0.54
                                - 0.46
                                    * (2.0 * std::f32::consts::PI * i as f32
                                        / (window_size - 1) as f32)
                                        .cos())
                    })
                    .collect();

                // Compute power spectrum
                let power_spectrum = self.compute_power_spectrum(&windowed_frame);

                frames.push(AudioFrame {
                    features: power_spectrum,
                    timestamp: frame_idx as f32 * hop_size as f32 / sample_rate as f32,
                    duration: window_size as f32 / sample_rate as f32,
                });
            }
        }

        Ok(frames)
    }

    /// Quantize waveform directly
    async fn quantize_waveform(&self, audio_data: &[f32]) -> Result<Vec<AudioFrame>> {
        let frame_size = self.config.audio_config.hop_size;
        let sample_rate = self.config.audio_config.sample_rate;

        let mut frames = Vec::new();

        for (frame_idx, chunk) in audio_data.chunks(frame_size).enumerate() {
            // Simple quantization: divide into bins and represent as discrete values
            let quantized_samples: Vec<f32> = chunk
                .iter()
                .map(|&sample| {
                    // Quantize to 256 levels
                    let quantized = ((sample + 1.0) * 127.5).round().max(0.0).min(255.0);
                    quantized / 255.0 * 2.0 - 1.0
                })
                .collect();

            frames.push(AudioFrame {
                features: quantized_samples,
                timestamp: frame_idx as f32 * frame_size as f32 / sample_rate as f32,
                duration: frame_size as f32 / sample_rate as f32,
            });
        }

        Ok(frames)
    }

    /// Extract audio patches (similar to image patches but for audio)
    async fn extract_audio_patches(&self, audio_data: &[f32]) -> Result<Vec<AudioFrame>> {
        let patch_size = self.config.audio_config.window_size;
        let hop_size = self.config.audio_config.hop_size;
        let sample_rate = self.config.audio_config.sample_rate;

        let mut frames = Vec::new();

        for (frame_idx, frame_start) in (0..audio_data.len()).step_by(hop_size).enumerate() {
            if frame_start + patch_size <= audio_data.len() {
                let patch = audio_data[frame_start..frame_start + patch_size].to_vec();

                frames.push(AudioFrame {
                    features: patch,
                    timestamp: frame_idx as f32 * hop_size as f32 / sample_rate as f32,
                    duration: patch_size as f32 / sample_rate as f32,
                });
            }
        }

        Ok(frames)
    }

    /// Helper functions for audio processing

    /// Compute FFT magnitudes (simplified implementation)
    fn compute_fft_magnitudes(&self, frame: &[f32]) -> Vec<f32> {
        // Simplified FFT - in production use rustfft or similar
        let n = frame.len();
        let mut magnitudes = Vec::new();

        for k in 0..n / 2 {
            let mut real = 0.0;
            let mut imag = 0.0;

            for n_idx in 0..n {
                let angle = -2.0 * std::f32::consts::PI * k as f32 * n_idx as f32 / n as f32;
                real += frame[n_idx] * angle.cos();
                imag += frame[n_idx] * angle.sin();
            }

            magnitudes.push((real * real + imag * imag).sqrt());
        }

        magnitudes
    }

    /// Apply mel filterbank
    fn apply_mel_filterbank(&self, magnitudes: &[f32], sample_rate: u32) -> Vec<f32> {
        let n_mel = self.config.audio_config.n_mel_filters;
        let mut mel_features = vec![0.0; n_mel];

        // Simplified mel filterbank implementation
        let mel_max = self.hz_to_mel(sample_rate as f32 / 2.0);
        let mel_step = mel_max / (n_mel + 1) as f32;

        for i in 0..n_mel {
            let mel_center = (i + 1) as f32 * mel_step;
            let hz_center = self.mel_to_hz(mel_center);
            let bin_center = hz_center * magnitudes.len() as f32 * 2.0 / sample_rate as f32;

            // Simple triangular filter
            let start_bin = (bin_center - 10.0).max(0.0) as usize;
            let end_bin = (bin_center + 10.0).min(magnitudes.len() as f32) as usize;

            let mut energy = 0.0;
            for bin in start_bin..end_bin {
                let weight = 1.0 - (bin as f32 - bin_center).abs() / 10.0;
                if weight > 0.0 && bin < magnitudes.len() {
                    energy += magnitudes[bin] * weight;
                }
            }

            mel_features[i] = if energy > 0.0 { energy.ln() } else { -10.0 };
        }

        mel_features
    }

    /// Convert Hz to Mel scale
    fn hz_to_mel(&self, hz: f32) -> f32 {
        2595.0 * (1.0 + hz / 700.0).log10()
    }

    /// Convert Mel to Hz scale
    fn mel_to_hz(&self, mel: f32) -> f32 {
        700.0 * (10.0_f32.powf(mel / 2595.0) - 1.0)
    }

    /// Compute DCT for MFCC
    fn compute_dct(&self, mel_features: &[f32], n_mfcc: usize) -> Vec<f32> {
        let mut mfcc = vec![0.0; n_mfcc];
        let n = mel_features.len();

        for i in 0..n_mfcc {
            let mut sum = 0.0;
            for j in 0..n {
                sum += mel_features[j]
                    * (std::f32::consts::PI * i as f32 * (j as f32 + 0.5) / n as f32).cos();
            }
            mfcc[i] = sum;
        }

        mfcc
    }

    /// Compute power spectrum
    fn compute_power_spectrum(&self, frame: &[f32]) -> Vec<f32> {
        let magnitudes = self.compute_fft_magnitudes(frame);
        magnitudes.into_iter().map(|mag| mag * mag).collect()
    }

    /// Extract patches from all images for codebook training
    async fn extract_all_image_patches(&self, images: &[Vec<u8>]) -> Result<Vec<ImagePatch>> {
        let mut all_patches = Vec::new();

        for image_data in images {
            if let Ok(image) = image::load_from_memory(image_data) {
                let resized = image.resize_exact(
                    self.config.image_config.target_resolution.0,
                    self.config.image_config.target_resolution.1,
                    image::imageops::FilterType::Lanczos3,
                );
                let patches = self.extract_image_patches(&resized).await?;
                all_patches.extend(patches);
            }
        }

        Ok(all_patches)
    }

    /// Extract features from all audio samples for codebook training
    async fn extract_all_audio_features(
        &self,
        audio_samples: &[Vec<f32>],
    ) -> Result<Vec<AudioFrame>> {
        let mut all_features = Vec::new();

        for audio_data in audio_samples {
            let features = match self.config.audio_config.method {
                AudioTokenizationMethod::MFCC => self.extract_mfcc_features(audio_data).await?,
                AudioTokenizationMethod::Spectrogram => {
                    self.extract_spectrogram_features(audio_data).await?
                }
                AudioTokenizationMethod::WaveformQuantization => {
                    self.quantize_waveform(audio_data).await?
                }
                AudioTokenizationMethod::AudioPatches => {
                    self.extract_audio_patches(audio_data).await?
                }
            };
            all_features.extend(features);
        }

        Ok(all_features)
    }

    /// Build image codebook using k-means clustering
    async fn build_image_codebook(&self, patches: Vec<ImagePatch>) -> Result<Vec<ImagePatch>> {
        if patches.is_empty() {
            return Ok(Vec::new());
        }

        let k = self.config.vocab_size_per_modality.min(patches.len());
        let mut centroids = Vec::new();

        // Initialize centroids with random patches
        let mut rng = rand::thread_rng();
        use rand::seq::SliceRandom;
        let mut shuffled_patches = patches.clone();
        shuffled_patches.shuffle(&mut rng);

        for i in 0..k {
            centroids.push(shuffled_patches[i % shuffled_patches.len()].clone());
        }

        // Simple k-means (in production, use more sophisticated clustering)
        for _iteration in 0..10 {
            let mut new_centroids = vec![
                ImagePatch {
                    data: vec![0.0; patches[0].data.len()],
                    position: (0, 0),
                    dimensions: patches[0].dimensions,
                };
                k
            ];
            let mut counts = vec![0; k];

            // Assign patches to nearest centroids
            for patch in &patches {
                let mut best_distance = f32::INFINITY;
                let mut best_centroid = 0;

                for (centroid_idx, centroid) in centroids.iter().enumerate() {
                    let distance = self.compute_patch_distance(patch, centroid);
                    if distance < best_distance {
                        best_distance = distance;
                        best_centroid = centroid_idx;
                    }
                }

                // Accumulate for new centroid
                for (i, &value) in patch.data.iter().enumerate() {
                    new_centroids[best_centroid].data[i] += value;
                }
                counts[best_centroid] += 1;
            }

            // Update centroids
            for (centroid_idx, count) in counts.iter().enumerate() {
                if *count > 0 {
                    for value in &mut new_centroids[centroid_idx].data {
                        *value /= *count as f32;
                    }
                }
            }

            centroids = new_centroids;
        }

        Ok(centroids)
    }

    /// Build audio codebook using k-means clustering
    async fn build_audio_codebook(&self, features: Vec<AudioFrame>) -> Result<Vec<AudioFrame>> {
        if features.is_empty() {
            return Ok(Vec::new());
        }

        let k = self.config.vocab_size_per_modality.min(features.len());
        let mut centroids = Vec::new();

        // Initialize centroids
        let mut rng = rand::thread_rng();
        use rand::seq::SliceRandom;
        let mut shuffled_features = features.clone();
        shuffled_features.shuffle(&mut rng);

        for i in 0..k {
            centroids.push(shuffled_features[i % shuffled_features.len()].clone());
        }

        // k-means clustering
        for _iteration in 0..10 {
            let mut new_centroids = vec![
                AudioFrame {
                    features: vec![0.0; features[0].features.len()],
                    timestamp: 0.0,
                    duration: 0.0,
                };
                k
            ];
            let mut counts = vec![0; k];

            for feature in &features {
                let mut best_distance = f32::INFINITY;
                let mut best_centroid = 0;

                for (centroid_idx, centroid) in centroids.iter().enumerate() {
                    let distance = self.compute_feature_distance(feature, centroid);
                    if distance < best_distance {
                        best_distance = distance;
                        best_centroid = centroid_idx;
                    }
                }

                for (i, &value) in feature.features.iter().enumerate() {
                    new_centroids[best_centroid].features[i] += value;
                }
                counts[best_centroid] += 1;
            }

            for (centroid_idx, count) in counts.iter().enumerate() {
                if *count > 0 {
                    for value in &mut new_centroids[centroid_idx].features {
                        *value /= *count as f32;
                    }
                }
            }

            centroids = new_centroids;
        }

        Ok(centroids)
    }

    /// Build cross-modal alignment matrix
    async fn build_cross_modal_alignment(&mut self) -> Result<()> {
        let image_codebook_size = self.image_codebook.read().len();
        let audio_codebook_size = self.audio_codebook.read().len();

        if image_codebook_size > 0 && audio_codebook_size > 0 {
            // Create a simple random alignment matrix (in production, use learned alignment)
            let matrix = nalgebra::DMatrix::<f32>::from_fn(
                image_codebook_size,
                audio_codebook_size,
                |_i, _j| rand::random::<f32>(),
            );

            *self.alignment_matrix.write() = Some(matrix);
        }

        Ok(())
    }

    /// Quantize image patches to token IDs
    async fn quantize_image_patches(&self, patches: Vec<ImagePatch>) -> Result<Vec<TokenId>> {
        let codebook = self.image_codebook.read();
        let mut tokens = Vec::new();

        for patch in patches {
            let mut best_distance = f32::INFINITY;
            let mut best_token = 0;

            for (token_id, centroid) in codebook.iter().enumerate() {
                let distance = self.compute_patch_distance(&patch, centroid);
                if distance < best_distance {
                    best_distance = distance;
                    best_token = token_id as TokenId;
                }
            }

            tokens.push(best_token);
        }

        Ok(tokens)
    }

    /// Quantize audio features to token IDs
    async fn quantize_audio_features(&self, features: Vec<AudioFrame>) -> Result<Vec<TokenId>> {
        let codebook = self.audio_codebook.read();
        let mut tokens = Vec::new();

        for feature in features {
            let mut best_distance = f32::INFINITY;
            let mut best_token = 0;

            for (token_id, centroid) in codebook.iter().enumerate() {
                let distance = self.compute_feature_distance(&feature, centroid);
                if distance < best_distance {
                    best_distance = distance;
                    best_token = token_id as TokenId;
                }
            }

            tokens.push(best_token);
        }

        Ok(tokens)
    }

    /// Compute Euclidean distance between image patches
    fn compute_patch_distance(&self, patch1: &ImagePatch, patch2: &ImagePatch) -> f32 {
        patch1
            .data
            .iter()
            .zip(patch2.data.iter())
            .map(|(a, b)| (a - b).powi(2))
            .sum::<f32>()
            .sqrt()
    }

    /// Compute Euclidean distance between audio features
    fn compute_feature_distance(&self, feature1: &AudioFrame, feature2: &AudioFrame) -> f32 {
        feature1
            .features
            .iter()
            .zip(feature2.features.iter())
            .map(|(a, b)| (a - b).powi(2))
            .sum::<f32>()
            .sqrt()
    }
}

#[async_trait::async_trait]
impl Tokenizer for MultimodalTokenizer {
    async fn encode(&self, text: &str) -> Result<Vec<TokenId>> {
        // Default text tokenization - simple character-based for now
        let start_time = std::time::Instant::now();

        let tokens: Vec<TokenId> = text
            .chars()
            .take(self.config.vocab_size_per_modality)
            .enumerate()
            .map(|(i, _)| i as TokenId)
            .collect();

        // Update statistics
        {
            let mut stats = self.stats.write();
            stats.token_count = tokens.len();
            stats.processing_time_ms = start_time.elapsed().as_millis() as u64;
        }

        Ok(tokens)
    }

    async fn decode(&self, token_ids: &[TokenId]) -> Result<String> {
        // Determine modality from token IDs and decode appropriately
        let mut decoded_parts = Vec::new();

        for &token_id in token_ids {
            if token_id < 10000 {
                decoded_parts.push("TEXT".to_string());
            } else if token_id < 20000 {
                decoded_parts.push("IMAGE".to_string());
            } else if token_id < 30000 {
                decoded_parts.push("AUDIO".to_string());
            } else {
                decoded_parts.push("UNKNOWN".to_string());
            }
        }

        Ok(format!("multimodal[{}]", decoded_parts.join(",")))
    }

    fn vocab_size(&self) -> usize {
        self.config.modalities.len() * self.config.vocab_size_per_modality
    }

    fn get_stats(&self) -> TokenizationStats {
        self.stats.read().clone()
    }

    fn supports(&self, _text: &str) -> bool {
        self.config.modalities.contains(&ModalityType::Text)
    }
}

/// Simplified serialization data structure
#[derive(serde::Serialize, serde::Deserialize)]
struct MultimodalTokenizerData {
    config: MultimodalConfig,
    image_codebook: Vec<ImagePatch>,
    audio_codebook: Vec<AudioFrame>,
}

impl serde::Serialize for MultimodalTokenizer {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let data = MultimodalTokenizerData {
            config: self.config.clone(),
            image_codebook: self.image_codebook.read().clone(),
            audio_codebook: self.audio_codebook.read().clone(),
        };
        data.serialize(serializer)
    }
}

impl<'de> serde::Deserialize<'de> for MultimodalTokenizer {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let data = MultimodalTokenizerData::deserialize(deserializer)?;
        Ok(Self {
            config: data.config,
            image_codebook: Arc::new(RwLock::new(data.image_codebook)),
            audio_codebook: Arc::new(RwLock::new(data.audio_codebook)),
            alignment_matrix: Arc::new(RwLock::new(None)),
            stats: Arc::new(RwLock::new(TokenizationStats::default())),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use image::{ImageBuffer, Rgb, RgbImage};

    #[test]
    fn multimodal_config_default() {
        let config = MultimodalConfig::default();
        assert_eq!(config.modalities.len(), 3);
        assert!(config.modalities.contains(&ModalityType::Text));
        assert!(config.modalities.contains(&ModalityType::Image));
        assert!(config.modalities.contains(&ModalityType::Audio));
        assert_eq!(config.vocab_size_per_modality, 10000);
        assert!(config.cross_modal_alignment);
    }

    #[test]
    fn image_config_default() {
        let config = ImageTokenizationConfig::default();
        assert_eq!(config.method, ImageTokenizationMethod::Patches);
        assert_eq!(config.patch_size, (16, 16));
        assert_eq!(config.target_resolution, (224, 224));
        assert_eq!(config.channels, 3);
        assert!(config.normalization.is_some());
    }

    #[test]
    fn audio_config_default() {
        let config = AudioTokenizationConfig::default();
        assert_eq!(config.method, AudioTokenizationMethod::MFCC);
        assert_eq!(config.sample_rate, 16000);
        assert_eq!(config.window_size, 512);
        assert_eq!(config.hop_size, 256);
        assert_eq!(config.n_mel_filters, 80);
        assert_eq!(config.n_mfcc, 13);
    }

    #[tokio::test]
    async fn multimodal_tokenizer_creation() {
        let config = MultimodalConfig::default();
        let tokenizer = MultimodalTokenizer::new(config);

        assert_eq!(tokenizer.vocab_size(), 30000); // 3 modalities * 10000 each
        assert!(tokenizer.supports("test"));
    }

    #[test]
    fn image_patch_creation() {
        let patch = ImagePatch {
            data: vec![0.1, 0.2, 0.3],
            position: (0, 1),
            dimensions: (16, 16),
        };

        assert_eq!(patch.data.len(), 3);
        assert_eq!(patch.position, (0, 1));
        assert_eq!(patch.dimensions, (16, 16));
    }

    #[test]
    fn audio_frame_creation() {
        let frame = AudioFrame {
            features: vec![0.1, 0.2, 0.3, 0.4],
            timestamp: 1.5,
            duration: 0.025,
        };

        assert_eq!(frame.features.len(), 4);
        assert_eq!(frame.timestamp, 1.5);
        assert_eq!(frame.duration, 0.025);
    }

    #[tokio::test]
    async fn encode_decode_text() {
        let config = MultimodalConfig::default();
        let tokenizer = MultimodalTokenizer::new(config);

        let tokens = tokenizer.encode("hello").await.unwrap();
        assert!(!tokens.is_empty());

        let decoded = tokenizer.decode(&tokens).await.unwrap();
        assert!(decoded.contains("multimodal"));
    }

    #[test]
    fn hz_mel_conversion() {
        let config = MultimodalConfig::default();
        let tokenizer = MultimodalTokenizer::new(config);

        let hz = 1000.0;
        let mel = tokenizer.hz_to_mel(hz);
        let hz_back = tokenizer.mel_to_hz(mel);

        assert!((hz - hz_back).abs() < 0.1);
    }

    #[tokio::test]
    async fn real_image_processing() {
        let config = MultimodalConfig::default();
        let tokenizer = MultimodalTokenizer::new(config);

        // Create a simple test image (solid red 224x224)
        let width = 224u32;
        let height = 224u32;
        let mut image_data: Vec<u8> = Vec::new();

        // PNG header and simple red image
        let image: RgbImage = ImageBuffer::from_fn(width, height, |_x, _y| Rgb([255u8, 0u8, 0u8]));
        let dynamic_image = DynamicImage::ImageRgb8(image);

        // Test patch extraction
        let patches = tokenizer
            .extract_image_patches(&dynamic_image)
            .await
            .unwrap();
        assert!(!patches.is_empty());
        assert_eq!(patches[0].dimensions, (16, 16));
        assert_eq!(patches[0].data.len(), 16 * 16 * 3); // RGB
    }

    #[tokio::test]
    async fn real_audio_processing() {
        let config = MultimodalConfig::default();
        let tokenizer = MultimodalTokenizer::new(config);

        // Create test audio - simple sine wave
        let sample_rate = 16000;
        let duration = 1.0; // 1 second
        let frequency = 440.0; // A4 note
        let samples: Vec<f32> = (0..((sample_rate as f32 * duration) as usize))
            .map(|i| (2.0 * std::f32::consts::PI * frequency * i as f32 / sample_rate as f32).sin())
            .collect();

        // Test MFCC extraction
        let frames = tokenizer.extract_mfcc_features(&samples).await.unwrap();
        assert!(!frames.is_empty());
        assert_eq!(
            frames[0].features.len(),
            tokenizer.config.audio_config.n_mfcc
        );
    }
}