use crate::error::{MultimodalError, Result};
use crate::vision::vit::{LayerNorm, MultiHeadAttention, TransformerBlock, TransformerConfig};
use rtx_tensor::Device;
use rtx_tensor::Tensor;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeSformerConfig {
    pub image_size: usize,
    pub patch_size: usize,
    pub in_channels: usize,
    pub embed_dim: usize,
    pub depth: usize,
    pub num_heads: usize,
    pub mlp_ratio: f32,
    pub num_frames: usize,
    pub num_classes: usize,
    pub dropout: f32,
    pub attention_dropout: f32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DividedSpaceTimeAttentionConfig {
    pub embed_dim: usize,
    pub num_heads: usize,
    pub dropout: f32,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimeSformerBlockConfig {
    pub embed_dim: usize,
    pub num_heads: usize,
    pub mlp_ratio: f32,
    pub dropout: f32,
    pub attention_dropout: f32,
}

pub struct VideoPatchEmbedding {
    conv: Tensor, // 3D convolution weights
    bias: Tensor,
    num_patches: usize,
    embed_dim: usize,
    num_frames: usize,
    device: Device,
}

impl VideoPatchEmbedding {
    pub fn new(
        image_size: usize,
        patch_size: usize,
        in_channels: usize,
        embed_dim: usize,
        num_frames: usize,
        device: &Device,
    ) -> Result<Self> {
        if !image_size.is_multiple_of(patch_size) {
            return Err(MultimodalError::InvalidPatchSize {
                patch_size,
                image_size,
            });
        }

        let num_patches = (image_size / patch_size).pow(2);

        // 2D convolution for spatial patches (we handle temporal dimension separately)
        let fan_out = embed_dim * patch_size * patch_size;
        let bound = (6.0 / (in_channels * patch_size * patch_size + fan_out) as f32).sqrt();

        // Use randn and scale to [-bound, bound] range
        let conv = Tensor::randn(&[embed_dim, in_channels, patch_size, patch_size], device)?
            .mul_scalar(bound / 3.0)?; // Scale by bound/3 for similar distribution

        let bias = Tensor::zeros([embed_dim], device)?;

        Ok(Self {
            conv,
            bias,
            num_patches,
            embed_dim,
            num_frames,
            device: device.clone(),
        })
    }

    pub fn num_patches(&self) -> usize {
        self.num_patches
    }

    pub fn embed_dim(&self) -> usize {
        self.embed_dim
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // x shape: [batch_size, channels, frames, height, width]
        let _batch_size = x.shape()[0];
        let _channels = x.shape()[1];
        let frames = x.shape()[2];
        let _height = x.shape()[3];
        let _width = x.shape()[4];

        // Process each frame separately
        let mut frame_patches = Vec::new();
        for t in 0..frames {
            // Extract frame: [batch, channels, height, width]
            let frame = x.slice(2, t, t + 1)?.squeeze(Some(2))?;

            // Apply 2D convolution to create patches
            let patches = frame.conv2d(&self.conv, Some(&self.bias), 1, 0, 1, 1)?;

            // Flatten spatial dimensions
            let batch_size = patches.shape()[0];
            let embed_dim = patches.shape()[1];
            let h_patches = patches.shape()[2];
            let w_patches = patches.shape()[3];

            let flattened = patches.reshape([batch_size, embed_dim, h_patches * w_patches])?;
            let transposed = flattened.transpose(1, 2)?; // [batch, num_patches, embed_dim]

            frame_patches.push(transposed);
        }

        // Stack frame patches: [batch, frames, num_patches, embed_dim]
        let video_patches = Tensor::stack(&frame_patches, 1)?;

        Ok(video_patches)
    }
}

pub struct TemporalPositionalEncoding {
    embeddings: Tensor,
    max_frames: usize,
    device: Device,
}

impl TemporalPositionalEncoding {
    pub fn new(max_frames: usize, embed_dim: usize, device: &Device) -> Result<Self> {
        let embeddings = Tensor::randn(&[1, max_frames, embed_dim], device)?;

        Ok(Self {
            embeddings,
            max_frames,
            device: device.clone(),
        })
    }

    pub fn forward(&self, num_frames: usize) -> Result<Tensor> {
        if num_frames > self.max_frames {
            return Err(MultimodalError::InvalidSequenceLength {
                seq_len: num_frames,
                max_len: self.max_frames,
            });
        }

        let indices = Tensor::arange(0, num_frames as i64, &self.device)?
            .unsqueeze(0)?
            .unsqueeze(2)?;

        Ok(self.embeddings.gather(1, &indices)?)
    }
}

pub struct SpatialPositionalEncoding {
    embeddings: Tensor,
    max_patches: usize,
    device: Device,
}

impl SpatialPositionalEncoding {
    pub fn new(max_patches: usize, embed_dim: usize, device: &Device) -> Result<Self> {
        let embeddings = Tensor::randn(&[1, max_patches, embed_dim], device)?;

        Ok(Self {
            embeddings,
            max_patches,
            device: device.clone(),
        })
    }

    pub fn forward(&self, num_patches: usize) -> Result<Tensor> {
        if num_patches > self.max_patches {
            return Err(MultimodalError::InvalidSequenceLength {
                seq_len: num_patches,
                max_len: self.max_patches,
            });
        }

        let indices = Tensor::arange(0, num_patches as i64, &self.device)?
            .unsqueeze(0)?
            .unsqueeze(2)?;

        Ok(self.embeddings.gather(1, &indices)?)
    }
}

pub struct SpaceOnlyAttention {
    attention: MultiHeadAttention,
    device: Device,
}

impl SpaceOnlyAttention {
    pub fn new(embed_dim: usize, num_heads: usize, _dropout: f32, device: &Device) -> Result<Self> {
        let attention = MultiHeadAttention::new(embed_dim, num_heads, device)?;

        Ok(Self {
            attention,
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // x shape: [batch, frames, patches, embed_dim]
        let _batch_size = x.shape()[0];
        let frames = x.shape()[1];
        let _patches = x.shape()[2];
        let _embed_dim = x.shape()[3];

        // Process each frame independently
        let mut frame_outputs = Vec::new();
        for t in 0..frames {
            let frame = x.slice(1, t, t + 1)?.squeeze(Some(1))?; // [batch, patches, embed_dim]
            let frame_out = self.attention.forward(&frame)?;
            frame_outputs.push(frame_out);
        }

        // Stack outputs: [batch, frames, patches, embed_dim]
        let output = Tensor::stack(&frame_outputs, 1)?;

        Ok(output)
    }
}

pub struct TimeOnlyAttention {
    attention: MultiHeadAttention,
    device: Device,
}

impl TimeOnlyAttention {
    pub fn new(embed_dim: usize, num_heads: usize, _dropout: f32, device: &Device) -> Result<Self> {
        let attention = MultiHeadAttention::new(embed_dim, num_heads, device)?;

        Ok(Self {
            attention,
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // x shape: [batch, frames, patches, embed_dim]
        let _batch_size = x.shape()[0];
        let _frames = x.shape()[1];
        let patches = x.shape()[2];
        let _embed_dim = x.shape()[3];

        // Process each spatial location independently across time
        let mut patch_outputs = Vec::new();
        for p in 0..patches {
            let patch_temporal = x.slice(2, p, p + 1)?.squeeze(Some(2))?; // [batch, frames, embed_dim]
            let patch_out = self.attention.forward(&patch_temporal)?;
            patch_outputs.push(patch_out);
        }

        // Stack outputs and permute back: [batch, patches, frames, embed_dim] -> [batch, frames, patches, embed_dim]
        let stacked = Tensor::stack(&patch_outputs, 1)?; // [batch, patches, frames, embed_dim]
        let output = stacked.transpose(1, 2)?; // [batch, frames, patches, embed_dim]

        Ok(output)
    }
}

pub struct JointSpaceTimeAttention {
    attention: MultiHeadAttention,
    device: Device,
}

impl JointSpaceTimeAttention {
    pub fn new(embed_dim: usize, num_heads: usize, _dropout: f32, device: &Device) -> Result<Self> {
        let attention = MultiHeadAttention::new(embed_dim, num_heads, device)?;

        Ok(Self {
            attention,
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // x shape: [batch, frames, patches, embed_dim]
        let batch_size = x.shape()[0];
        let frames = x.shape()[1];
        let patches = x.shape()[2];
        let embed_dim = x.shape()[3];

        // Flatten spatial-temporal dimensions for joint attention
        let x_flat = x.reshape([batch_size, frames * patches, embed_dim])?;
        let out_flat = self.attention.forward(&x_flat)?;

        // Reshape back to original dimensions
        let output = out_flat.reshape([batch_size, frames, patches, embed_dim])?;

        Ok(output)
    }
}

pub struct DividedSpaceTimeAttention {
    space_attention: SpaceOnlyAttention,
    time_attention: TimeOnlyAttention,
    device: Device,
}

impl DividedSpaceTimeAttention {
    pub fn new(config: &DividedSpaceTimeAttentionConfig, device: &Device) -> Result<Self> {
        let space_attention =
            SpaceOnlyAttention::new(config.embed_dim, config.num_heads, config.dropout, device)?;

        let time_attention =
            TimeOnlyAttention::new(config.embed_dim, config.num_heads, config.dropout, device)?;

        Ok(Self {
            space_attention,
            time_attention,
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // Apply spatial attention first
        let x = self.space_attention.forward(x)?;

        // Then apply temporal attention
        let x = self.time_attention.forward(&x)?;

        Ok(x)
    }
}

pub struct TimeSformerBlock {
    temporal_attn: TimeOnlyAttention,
    spatial_attn: SpaceOnlyAttention,
    mlp: TransformerBlock, // Reuse the MLP from TransformerBlock
    ln1: LayerNorm,
    ln2: LayerNorm,
    ln3: LayerNorm,
    device: Device,
}

impl TimeSformerBlock {
    pub fn new(config: &TimeSformerBlockConfig, device: &Device) -> Result<Self> {
        let temporal_attn = TimeOnlyAttention::new(
            config.embed_dim,
            config.num_heads,
            config.attention_dropout,
            device,
        )?;

        let spatial_attn = SpaceOnlyAttention::new(
            config.embed_dim,
            config.num_heads,
            config.attention_dropout,
            device,
        )?;

        let mlp_config = TransformerConfig {
            embed_dim: config.embed_dim,
            num_heads: config.num_heads,
            mlp_ratio: config.mlp_ratio,
            dropout: config.dropout,
            attention_dropout: config.attention_dropout,
        };
        let mlp = TransformerBlock::new(&mlp_config, device)?;

        let ln1 = LayerNorm::new(config.embed_dim, device)?;
        let ln2 = LayerNorm::new(config.embed_dim, device)?;
        let ln3 = LayerNorm::new(config.embed_dim, device)?;

        Ok(Self {
            temporal_attn,
            spatial_attn,
            mlp,
            ln1,
            ln2,
            ln3,
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // Temporal attention
        let temp_out = self.temporal_attn.forward(&self.ln1.forward(x)?)?;
        let x = (x + &temp_out)?;

        // Spatial attention
        let spat_out = self.spatial_attn.forward(&self.ln2.forward(&x)?)?;
        let x = (&x + &spat_out)?;

        // We need to reshape for the MLP which expects [batch, seq, embed]
        let batch_size = x.shape()[0];
        let frames = x.shape()[1];
        let patches = x.shape()[2];
        let embed_dim = x.shape()[3];

        let x_flat = x.reshape([batch_size, frames * patches, embed_dim])?;
        let mlp_out = self.mlp.forward(&x_flat)?;
        let mlp_out = mlp_out.reshape([batch_size, frames, patches, embed_dim])?;

        Ok(mlp_out)
    }
}

pub struct TimeSformerModel {
    patch_embed: VideoPatchEmbedding,
    cls_token: Tensor,
    temporal_pos_embed: TemporalPositionalEncoding,
    spatial_pos_embed: SpatialPositionalEncoding,
    blocks: Vec<TimeSformerBlock>,
    ln_f: LayerNorm,
    head: Tensor,
    bias: Tensor,
    config: TimeSformerConfig,
    device: Device,
}

impl TimeSformerModel {
    pub fn new(config: &TimeSformerConfig, device: &Device) -> Result<Self> {
        let patch_embed = VideoPatchEmbedding::new(
            config.image_size,
            config.patch_size,
            config.in_channels,
            config.embed_dim,
            config.num_frames,
            device,
        )?;

        // Class token for each frame
        let cls_token = Tensor::randn(&[1, config.num_frames, 1, config.embed_dim], device)?;

        let temporal_pos_embed =
            TemporalPositionalEncoding::new(config.num_frames, config.embed_dim, device)?;

        let spatial_pos_embed = SpatialPositionalEncoding::new(
            patch_embed.num_patches() + 1, // +1 for cls token
            config.embed_dim,
            device,
        )?;

        let block_config = TimeSformerBlockConfig {
            embed_dim: config.embed_dim,
            num_heads: config.num_heads,
            mlp_ratio: config.mlp_ratio,
            dropout: config.dropout,
            attention_dropout: config.attention_dropout,
        };

        let mut blocks = Vec::new();
        for _ in 0..config.depth {
            blocks.push(TimeSformerBlock::new(&block_config, device)?);
        }

        let ln_f = LayerNorm::new(config.embed_dim, device)?;
        let head = Tensor::randn(&[config.num_classes, config.embed_dim], device)?;
        let bias = Tensor::zeros([config.num_classes], device)?;

        Ok(Self {
            patch_embed,
            cls_token,
            temporal_pos_embed,
            spatial_pos_embed,
            blocks,
            ln_f,
            head,
            bias,
            config: config.clone(),
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let batch_size = x.shape()[0];

        // Video patch embedding: [batch, frames, patches, embed_dim]
        let patches = self.patch_embed.forward(x)?;

        // Add cls tokens for each frame
        let cls_tokens = self.cls_token.clone().expand(&[
            batch_size,
            self.config.num_frames,
            1,
            self.config.embed_dim,
        ])?;
        let x = Tensor::cat(&[cls_tokens, patches], 2)?; // Concat along patch dimension

        // Add positional encodings
        let spatial_pos = self.spatial_pos_embed.forward(x.shape()[2])?; // [1, patches+1, embed_dim]
        let temporal_pos = self.temporal_pos_embed.forward(self.config.num_frames)?; // [1, frames, embed_dim]

        // Broadcast and add positional encodings
        let spatial_pos_expanded = spatial_pos.unsqueeze(1)?.expand(&[
            batch_size,
            self.config.num_frames,
            x.shape()[2],
            self.config.embed_dim,
        ])?;
        let temporal_pos_expanded = temporal_pos.unsqueeze(2)?.expand(&[
            batch_size,
            self.config.num_frames,
            x.shape()[2],
            self.config.embed_dim,
        ])?;

        let x = ((&x + &spatial_pos_expanded)? + temporal_pos_expanded)?;

        // Pass through TimeSformer blocks
        let mut x = x;
        for block in &self.blocks {
            x = block.forward(&x)?;
        }

        // Final layer norm
        let x = self.ln_f.forward(&x)?;

        // Classification head (use cls tokens)
        let cls_outputs = x.slice(2, 0, 1)?.squeeze(Some(2))?; // [batch, frames, embed_dim]
        let cls_pooled = cls_outputs.mean(&[1], false)?; // Pool across time: [batch, embed_dim]
        let logits = cls_pooled
            .matmul(&self.head.transpose(0, 1)?)?
            .add(&self.bias)?;

        Ok(logits)
    }

    pub fn extract_features(&self, x: &Tensor) -> Result<Tensor> {
        let batch_size = x.shape()[0];

        // Video patch embedding
        let patches = self.patch_embed.forward(x)?;

        // Add cls tokens for each frame
        let cls_tokens = self.cls_token.clone().expand(&[
            batch_size,
            self.config.num_frames,
            1,
            self.config.embed_dim,
        ])?;
        let x = Tensor::cat(&[cls_tokens, patches], 2)?;

        // Add positional encodings
        let spatial_pos = self.spatial_pos_embed.forward(x.shape()[2])?;
        let temporal_pos = self.temporal_pos_embed.forward(self.config.num_frames)?;

        let spatial_pos_expanded = spatial_pos.unsqueeze(1)?.expand(&[
            batch_size,
            self.config.num_frames,
            x.shape()[2],
            self.config.embed_dim,
        ])?;
        let temporal_pos_expanded = temporal_pos.unsqueeze(2)?.expand(&[
            batch_size,
            self.config.num_frames,
            x.shape()[2],
            self.config.embed_dim,
        ])?;

        let x = ((&x + &spatial_pos_expanded)? + temporal_pos_expanded)?;

        // Pass through blocks
        let mut x = x;
        for block in &self.blocks {
            x = block.forward(&x)?;
        }

        // Return cls token features across time
        let cls_features = x.slice(2, 0, 1)?.squeeze(Some(2))?; // [batch, frames, embed_dim]

        Ok(cls_features)
    }
}

pub struct VideoFrameSampler {
    num_frames: usize,
    uniform_sampling: bool,
    device: Device,
}

impl VideoFrameSampler {
    pub fn new(num_frames: usize, uniform_sampling: bool, device: &Device) -> Result<Self> {
        Ok(Self {
            num_frames,
            uniform_sampling,
            device: device.clone(),
        })
    }

    pub fn sample_frames(&self, video: &Tensor) -> Result<Tensor> {
        let total_frames = video.shape()[2];

        if total_frames <= self.num_frames {
            // If video has fewer frames than requested, repeat last frame
            return Ok(video.clone());
        }

        if self.uniform_sampling {
            // Uniform sampling
            let step = total_frames as f32 / self.num_frames as f32;
            let mut indices = Vec::new();
            for i in 0..self.num_frames {
                let idx = ((i as f32 + 0.5) * step) as usize;
                indices.push(idx.min(total_frames - 1));
            }

            let _index_tensor = Tensor::from_vec(
                indices.into_iter().map(|i| i as f32).collect(),
                &[self.num_frames],
                &self.device,
            )?;

            // For now, return the video as-is (would need proper frame selection)
            Ok(video.clone())
        } else {
            // Random sampling (simplified)
            let indices: Vec<f32> = (0..self.num_frames)
                .map(|_| (rand::random::<f32>() * total_frames as f32) as f32)
                .collect();

            let _index_tensor = Tensor::from_vec(indices, &[self.num_frames], &self.device)?;

            // For now, return the video as-is (would need proper frame selection)
            Ok(video.clone())
        }
    }
}

pub struct VideoAugmentor {
    temporal_crop_ratio: f32,
    spatial_crop_ratio: f32,
    random_horizontal_flip: bool,
    device: Device,
}

impl VideoAugmentor {
    pub fn new(
        temporal_crop_ratio: f32,
        spatial_crop_ratio: f32,
        random_horizontal_flip: bool,
        device: &Device,
    ) -> Result<Self> {
        Ok(Self {
            temporal_crop_ratio,
            spatial_crop_ratio,
            random_horizontal_flip,
            device: device.clone(),
        })
    }

    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // Simplified augmentation - in practice you'd implement proper video augmentation
        // This is just a placeholder that returns the input
        Ok(x.clone())
    }
}

pub struct VideoPreprocessor {
    target_size: usize,
    num_frames: usize,
    fps: f32,
    device: Device,
}

impl VideoPreprocessor {
    pub fn new(target_size: usize, num_frames: usize, fps: f32, device: &Device) -> Result<Self> {
        Ok(Self {
            target_size,
            num_frames,
            fps,
            device: device.clone(),
        })
    }

    pub fn preprocess(&self, raw_frames: &[Vec<Vec<Vec<u8>>>]) -> Result<Tensor> {
        // Convert raw frames to tensor format
        let total_frames = raw_frames.len();
        let height = raw_frames[0].len();
        let width = raw_frames[0][0].len();
        let channels = 3;

        // Sample frames if needed
        let frame_indices = if total_frames > self.num_frames {
            let step = total_frames as f32 / self.num_frames as f32;
            (0..self.num_frames)
                .map(|i| ((i as f32 + 0.5) * step) as usize)
                .collect::<Vec<_>>()
        } else {
            (0..total_frames).collect()
        };

        // Convert to f32 and normalize
        let mut video_data = Vec::new();
        for &frame_idx in &frame_indices {
            let frame_idx = frame_idx.min(total_frames - 1);
            let frame = &raw_frames[frame_idx];

            for c in 0..channels {
                for h in 0..height {
                    for w in 0..width {
                        let pixel_value =
                            if h < frame.len() && w < frame[h].len() && c < frame[h][w].len() {
                                frame[h][w][c] as f32 / 255.0
                            } else {
                                0.0
                            };
                        video_data.push(pixel_value);
                    }
                }
            }
        }

        let actual_frames = frame_indices.len();
        let tensor = Tensor::from_vec(
            video_data,
            &[1, channels, actual_frames, height, width],
            &self.device,
        )?;

        Ok(tensor)
    }
}