import math
from typing import Optional

import mlx.core as mx
import mlx.nn as nn
import numpy as np
from transformers import AutoProcessor

from mlx_vlm.models.base import InputEmbeddingsFeatures

from .config import ModelConfig, ProjectorConfig, SAMViTConfig
from .language import LanguageModel
from .processing_deepseekocr import DeepseekOCRProcessor
from .sam import SAMEncoder
from .vision import VisionModel

AutoProcessor.register("deepseekocr", DeepseekOCRProcessor)


class MlpProjector(nn.Module):
    def __init__(self, config: ProjectorConfig):
        super().__init__()
        self.config = config

        if config.projector_config.projector_type == "linear":
            modules = nn.Linear(
                config.projector_config.input_dim, config.projector_config.n_embed
            )

        elif config.projector_config.projector_type == "downsample_mlp_gelu":
            mlp_depth = config.projector_config.depth
            mlp_ratio = config.projector_config.mlp_ratio
            modules = [
                nn.Linear(
                    config.projector_config.input_dim
                    * config.projector_config.downsample_ratio
                    * config.projector_config.downsample_ratio,
                    config.projector_config.n_embed * mlp_ratio,
                )
            ]
            for _ in range(1, mlp_depth - 1):
                modules.append(nn.GELU())
                modules.append(
                    nn.Linear(
                        config.projector_config.n_embed * mlp_ratio,
                        config.projector_config.n_embed * mlp_ratio,
                    )
                )
            modules.append(nn.GELU())
            modules.append(
                nn.Linear(
                    config.projector_config.n_embed * mlp_ratio,
                    config.projector_config.n_embed,
                )
            )
        else:
            raise ValueError(
                f"Unknown projector type: {config.projector_config.projector_type}"
            )

        self.layers = modules

    def __call__(self, x):
        if self.config.projector_config.projector_type == "downsample_mlp_gelu":
            bs, hw, input_dim = x.shape
            h = w = int(math.sqrt(hw))

            # Compute padding
            pad = (
                0
                if h % self.config.projector_config.downsample_ratio == 0
                else self.config.projector_config.downsample_ratio
                - h % self.config.projector_config.downsample_ratio
            )

            x = mx.reshape(x, (bs, h, w, input_dim))
            if pad > 0:
                x = mx.pad(x, [(0, 0), (0, pad), (0, pad), (0, 0)], constant_values=0)

            x = mx.transpose(x, (0, 3, 1, 2))  # B, C, H, W

            # Manual implementation of unfold for downsampling
            h_pad, w_pad = x.shape[2], x.shape[3]
            ds = self.config.projector_config.downsample_ratio
            patches = []

            for i in range(0, h_pad - ds + 1, ds):
                for j in range(0, w_pad - ds + 1, ds):
                    patch = x[:, :, i : i + ds, j : j + ds]
                    patches.append(mx.reshape(patch, (bs, -1)))

            x = mx.stack(patches, axis=1)  # B, N_patches, C*ds*ds

        if self.config.projector_config.projector_type == "linear":
            x = self.layers(x)
        else:
            for layer in self.layers:
                x = layer(x)
        return x


class Model(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.vision_model = VisionModel(config.vision_config)
        sam_config = SAMViTConfig()
        self.sam_model = SAMEncoder(
            img_size=sam_config.image_size,
            patch_size=sam_config.patch_size,
            embed_dim=sam_config.width,
            depth=sam_config.layers,
            num_heads=sam_config.heads,
            window_size=sam_config.window_size,
            global_attn_indexes=sam_config.global_attn_indexes,
        )
        self.language_model = LanguageModel(config.text_config)
        self.projector = MlpProjector(config)

        self.tile_tag = config.tile_tag
        self.global_view_pos = config.global_view_pos
        # 用于format image token sequence的特殊token
        embed_std = 1 / mx.sqrt(
            mx.array(config.projector_config.n_embed, dtype=mx.float32)
        )

        if self.tile_tag == "2D":

            # <|view_separator|>, <|\n|>
            self.image_newline = mx.array(
                mx.random.normal((config.projector_config.n_embed,)) * embed_std
            )
            # fix the typo: view_seperater
            self.view_separator = mx.array(
                mx.random.normal((config.projector_config.n_embed,)) * embed_std
            )
        else:
            raise ValueError(
                f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
            )

    def get_input_embeddings(
        self,
        input_ids: Optional[mx.array] = None,
        pixel_values: Optional[mx.array] = None,
        images_spatial_crop: Optional[mx.array] = None,
        images_seq_mask: Optional[mx.array] = None,
        **kwargs,
    ):
        input_embeds = self.language_model.model.embed_tokens(input_ids)

        if pixel_values is None:
            return InputEmbeddingsFeatures(inputs_embeds=input_embeds)

        # Only process images on prefill (input_ids.shape[1] != 1), not during autoregressive decoding
        if (
            self.sam_model is not None
            and input_ids.shape[1] != 1
            and mx.sum(pixel_values[1]).item() != 0
        ):

            idx = 0
            patch_idx = 0  # Track patch offset for batch processing
            all_patches = pixel_values[0]
            all_image_ori = pixel_values[1]

            for crop_shape in images_spatial_crop.tolist():
                images_in_this_batch = []
                width_crop_num, height_crop_num = int(crop_shape[0]), int(crop_shape[1])

                # Calculate number of patches for this image
                has_crops = width_crop_num > 1 or height_crop_num > 1
                num_patches = width_crop_num * height_crop_num if has_crops else 0

                # Extract patches for current image
                if has_crops and num_patches > 0:
                    patches = all_patches[patch_idx : patch_idx + num_patches]
                    patch_idx += num_patches
                else:
                    patches = None

                # Extract global image for current image (one per batch item)
                image_ori = all_image_ori[idx : idx + 1]

                if patches is not None and mx.sum(patches).item() != 0:
                    local_features_1 = self.sam_model(patches.transpose(0, 2, 3, 1))

                    local_features_2 = self.vision_model(
                        patches.transpose(0, 2, 3, 1), patch_embeds=local_features_1
                    )

                    local_features = mx.concatenate(
                        (
                            local_features_2[:, 1:],
                            local_features_1.flatten(start_axis=1, end_axis=2),
                        ),
                        axis=-1,
                    )

                    local_features = self.projector(local_features)

                    global_features_1 = self.sam_model(image_ori.transpose(0, 2, 3, 1))
                    global_features_2 = self.vision_model(
                        image_ori.transpose(0, 2, 3, 1), global_features_1
                    )

                    global_features = mx.concatenate(
                        (
                            global_features_2[:, 1:],
                            global_features_1.flatten(start_axis=1, end_axis=2),
                        ),
                        axis=-1,
                    )
                    global_features = self.projector(global_features)

                    # Remove batch dimension for single image processing
                    global_features = global_features[0]  # (hw, n_dim)
                    hw, n_dim = global_features.shape
                    h = w = int(hw**0.5)

                    _, hw2, n_dim2 = local_features.shape
                    h2 = w2 = int(hw2**0.5)

                    global_features = global_features.reshape(h, w, n_dim)

                    global_features = mx.concatenate(
                        [
                            global_features,
                            mx.broadcast_to(
                                self.image_newline[None, None, :], (h, 1, n_dim)
                            ),
                        ],
                        axis=1,
                    )

                    global_features = global_features.reshape(-1, n_dim)

                    local_features = (
                        local_features.reshape(
                            height_crop_num, width_crop_num, h2, w2, n_dim2
                        )
                        .transpose(0, 2, 1, 3, 4)
                        .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2)
                    )
                    local_features = mx.concatenate(
                        [
                            local_features,
                            mx.broadcast_to(
                                self.image_newline[None, None, :],
                                (height_crop_num * h2, 1, n_dim2),
                            ),
                        ],
                        axis=1,
                    )
                    local_features = local_features.reshape(-1, n_dim2)

                    global_local_features = mx.concatenate(
                        [local_features, global_features, self.view_separator[None, :]],
                        axis=0,
                    )

                else:
                    global_features_1 = self.sam_model(image_ori.transpose(0, 2, 3, 1))
                    global_features_2 = self.vision_model(
                        image_ori.transpose(0, 2, 3, 1), global_features_1
                    )
                    global_features = mx.concatenate(
                        (
                            global_features_2[:, 1:],
                            global_features_1.flatten(start_axis=1, end_axis=2),
                        ),
                        axis=-1,
                    )
                    global_features = self.projector(global_features)

                    # Remove batch dimension for single image processing
                    global_features = global_features[0]  # (hw, n_dim)
                    hw, n_dim = global_features.shape
                    h = w = int(hw**0.5)

                    global_features = global_features.reshape(h, w, n_dim)

                    global_features = mx.concatenate(
                        [
                            global_features,
                            mx.broadcast_to(
                                self.image_newline[None, None, :], (h, 1, n_dim)
                            ),
                        ],
                        axis=1,
                    )

                    global_features = global_features.reshape(-1, n_dim)

                    global_local_features = mx.concatenate(
                        [global_features, self.view_separator[None, :]], axis=0
                    )

                images_in_this_batch.append(global_local_features)

                if images_in_this_batch:
                    images_in_this_batch = mx.concatenate(images_in_this_batch, axis=0)
                    # Find positions where images should be placed
                    image_indices = np.where(images_seq_mask[idx])[0].tolist()
                    # Directly assign the image features to those positions
                    input_embeds[idx, image_indices] = images_in_this_batch

                idx += 1

        return InputEmbeddingsFeatures(inputs_embeds=input_embeds)

    @property
    def layers(self):
        return self.language_model.model.layers

    def __call__(
        self,
        input_ids: mx.array,
        pixel_values: Optional[mx.array] = None,
        mask: Optional[mx.array] = None,
        cache=None,
        **kwargs,
    ):

        images_spatial_crop = kwargs.get("images_spatial_crop", None)
        images_seq_mask = kwargs.get("images_seq_mask", None)

        input_embeddings = self.get_input_embeddings(
            input_ids, pixel_values, images_spatial_crop, images_seq_mask
        )

        logits = self.language_model(
            input_ids, cache=cache, inputs_embeds=input_embeddings.inputs_embeds
        )
        return logits

    @staticmethod
    def sanitize(weights):
        def transform_key(key):
            if "model.layers" in key and "language_model" not in key:
                key = key.replace("model.layers", "language_model.model.layers")

            if "model.embed_tokens" in key and "language_model" not in key:
                key = key.replace(
                    "model.embed_tokens", "language_model.model.embed_tokens"
                )

            if "model.norm" in key and "language_model" not in key:
                key = key.replace("model.norm", "language_model.model.norm")

            if "model.vision_model" in key:
                key = key.replace("model.vision_model", "vision_model")

            if "model.sam_model" in key:
                key = key.replace("model.sam_model", "sam_model")

            if "model.projector" in key:
                key = key.replace("model.projector", "projector")

            if "model.view_seperator" in key:
                key = key.replace("model.view_seperator", "view_separator")

            if "model.image_newline" in key:
                key = key.replace("model.image_newline", "image_newline")

            if "lm_head.weight" in key and "language_model" not in key:
                key = key.replace("lm_head.weight", "language_model.lm_head.weight")

            return key

        return {transform_key(k): v for k, v in weights.items()}
