from typing import List, Optional, Tuple

import mlx.core as mx
import mlx.nn as nn

from ..base import LanguageModelOutput, scaled_dot_product_attention
from ..cache import KVCache
from .config import ModelConfig, TextConfig


def precompute_freqs_1d(
    dim: int, end: int, theta: float = 10000.0
) -> Tuple[mx.array, mx.array]:
    freqs = 1.0 / (theta ** (mx.arange(0, dim, 2).astype(mx.float32)[: dim // 2] / dim))
    t = mx.arange(end).astype(mx.float32)
    freqs = t[:, None] * freqs[None, :]
    return mx.cos(freqs), mx.sin(freqs)


def apply_rotary_emb_1d(
    xq: mx.array, xk: mx.array, cos: mx.array, sin: mx.array
) -> Tuple[mx.array, mx.array]:
    dtype = xq.dtype
    *shape_q, d = xq.shape
    *shape_k, _ = xk.shape
    xq_r = xq.astype(mx.float32).reshape(*shape_q, d // 2, 2)
    xk_r = xk.astype(mx.float32).reshape(*shape_k, d // 2, 2)
    xq_0, xq_1 = xq_r[..., 0], xq_r[..., 1]
    xk_0, xk_1 = xk_r[..., 0], xk_r[..., 1]
    c = cos.reshape(1, 1, -1, cos.shape[-1])
    s = sin.reshape(1, 1, -1, sin.shape[-1])
    oq = mx.stack([xq_0 * c - xq_1 * s, xq_0 * s + xq_1 * c], axis=-1)
    ok = mx.stack([xk_0 * c - xk_1 * s, xk_0 * s + xk_1 * c], axis=-1)
    return oq.reshape(*shape_q, d).astype(dtype), ok.reshape(*shape_k, d).astype(dtype)


def compute_golden_freqs(
    freqs_golden: mx.array, pos_hw: mx.array
) -> Tuple[mx.array, mx.array]:
    theta = mx.einsum(
        "bsp,hfp->bshf", pos_hw.astype(mx.float32), freqs_golden.astype(mx.float32)
    )
    return mx.cos(theta), mx.sin(theta)


def apply_golden_rotary_emb(
    x: mx.array, cos_2d: mx.array, sin_2d: mx.array
) -> mx.array:
    dtype = x.dtype
    cos = cos_2d.transpose(0, 2, 1, 3)
    sin = sin_2d.transpose(0, 2, 1, 3)
    x_f = x.astype(mx.float32)
    x_even, x_odd = x_f[..., 0::2], x_f[..., 1::2]
    o_even = x_even * cos - x_odd * sin
    o_odd = x_even * sin + x_odd * cos
    return mx.stack([o_even, o_odd], axis=-1).reshape(x.shape).astype(dtype)


def apply_3d_rotary_emb(
    xq: mx.array,
    xk: mx.array,
    cos_1d: mx.array,
    sin_1d: mx.array,
    cos_2d: Optional[mx.array] = None,
    sin_2d: Optional[mx.array] = None,
) -> Tuple[mx.array, mx.array]:
    half = xq.shape[-1] // 2
    xq_t, xq_hw = xq[..., :half], xq[..., half:]
    xk_t, xk_hw = xk[..., :half], xk[..., half:]
    xq_t, xk_t = apply_rotary_emb_1d(xq_t, xk_t, cos_1d, sin_1d)
    if cos_2d is not None:
        xq_hw = apply_golden_rotary_emb(xq_hw, cos_2d, sin_2d)
        xk_hw = apply_golden_rotary_emb(xk_hw, cos_2d, sin_2d)
    dtype_q, dtype_k = xq.dtype, xk.dtype
    return (
        mx.concatenate([xq_t, xq_hw], axis=-1).astype(dtype_q),
        mx.concatenate([xk_t, xk_hw], axis=-1).astype(dtype_k),
    )


def compute_pos_hw(
    input_ids: mx.array,
    image_token_id: int,
    image_grid_hws: Optional[List[Tuple[int, int]]] = None,
) -> mx.array:
    num_tokens = input_ids.shape[-1]
    ids = input_ids.reshape(-1).tolist()
    img_indices = [i for i, t in enumerate(ids) if t == image_token_id]
    if not img_indices:
        return mx.zeros((1, num_tokens, 2))

    all_coords: List[Tuple[float, float]] = []
    if image_grid_hws:
        for gh, gw in image_grid_hws:
            for hi in range(gh):
                for wi in range(gw):
                    h_val = -((gh / gw) ** 0.5) + 2 * ((gh / gw) ** 0.5) * hi / max(
                        gh - 1, 1
                    )
                    w_val = -((gw / gh) ** 0.5) + 2 * ((gw / gh) ** 0.5) * wi / max(
                        gw - 1, 1
                    )
                    all_coords.append((h_val, w_val))

    hw_vals = [[0.0, 0.0]] * num_tokens
    for idx_i, tok_idx in enumerate(img_indices):
        if idx_i < len(all_coords):
            hw_vals[tok_idx] = list(all_coords[idx_i])
    return mx.array(hw_vals).reshape(1, num_tokens, 2)


def create_falcon_perception_mask(
    input_ids: mx.array,
    image_cls_id: int,
    img_end_id: int,
) -> mx.array:
    ids = input_ids.reshape(-1)
    S = ids.shape[0]

    soi = (ids == image_cls_id).astype(mx.int32)
    eoi = (ids == img_end_id).astype(mx.int32)
    acc_soi = mx.cumsum(soi)
    acc_eoi = mx.cumsum(eoi)
    in_image = (acc_soi - acc_eoi) > 0
    block_id = acc_soi * in_image.astype(mx.int32)

    q = mx.arange(S)
    causal = q[:, None] >= q[None, :]

    q_in = in_image[:, None]
    kv_in = in_image[None, :]
    q_blk = block_id[:, None]
    kv_blk = block_id[None, :]
    same_image = q_in & kv_in & (q_blk == kv_blk)

    attend = causal | same_image
    return attend.reshape(1, 1, S, S)


class Attention(nn.Module):

    def __init__(self, args: TextConfig):
        super().__init__()
        self.n_heads = args.num_attention_heads
        self.n_kv_heads = args.num_key_value_heads
        self.n_rep = self.n_heads // self.n_kv_heads
        self.head_dim = args.head_dim
        self.scale = self.head_dim**-0.5
        self.eps = args.rms_norm_eps
        q_size = self.n_heads * self.head_dim
        kv_size = self.n_kv_heads * self.head_dim
        self.q_size = q_size
        self.kv_size = kv_size
        self.wqkv = nn.Linear(args.hidden_size, q_size + 2 * kv_size, bias=False)
        self.wo = nn.Linear(q_size, args.hidden_size, bias=False)
        self.sinks = mx.zeros((self.n_heads,))
        self._norm_w_in = mx.ones((args.hidden_size,))
        self._norm_w_qk = mx.ones((self.head_dim,))

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array] = None,
        cache: Optional[KVCache] = None,
        cos_1d: Optional[mx.array] = None,
        sin_1d: Optional[mx.array] = None,
        cos_2d: Optional[mx.array] = None,
        sin_2d: Optional[mx.array] = None,
    ) -> mx.array:
        B, L, _ = x.shape
        x_norm = mx.fast.rms_norm(x, self._norm_w_in, eps=self.eps)

        qkv = self.wqkv(x_norm)
        q = qkv[..., : self.q_size]
        k = qkv[..., self.q_size : self.q_size + self.kv_size]
        v = qkv[..., self.q_size + self.kv_size :]

        q = q.reshape(B, L, self.n_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = k.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = v.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)

        q = mx.fast.rms_norm(q, self._norm_w_qk, eps=self.eps)
        k = mx.fast.rms_norm(k, self._norm_w_qk, eps=self.eps)

        if self.n_rep > 1:
            k = mx.repeat(k, self.n_rep, axis=1)
            v = mx.repeat(v, self.n_rep, axis=1)

        if cos_1d is not None:
            q, k = apply_3d_rotary_emb(q, k, cos_1d, sin_1d, cos_2d, sin_2d)

        if cache is not None:
            k, v = cache.update_and_fetch(k, v)

        output = scaled_dot_product_attention(
            q, k, v, cache=cache, scale=self.scale, mask=mask, sinks=self.sinks
        )

        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
        return self.wo(output)


class MLP(nn.Module):

    def __init__(self, args: TextConfig):
        super().__init__()
        self.hidden_dim = args.intermediate_size
        self.eps = args.rms_norm_eps
        self.w13 = nn.Linear(args.hidden_size, 2 * args.intermediate_size, bias=False)
        self.w2 = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
        self._norm_w = mx.ones((args.hidden_size,))

    def __call__(self, x: mx.array) -> mx.array:
        x_norm = mx.fast.rms_norm(x, self._norm_w, eps=self.eps)
        w13_out = self.w13(x_norm)
        gate = w13_out[..., : self.hidden_dim]
        up = w13_out[..., self.hidden_dim :]
        activated = nn.relu(gate) ** 2 * up
        return self.w2(activated)


class DecoderLayer(nn.Module):
    def __init__(self, args: TextConfig):
        super().__init__()
        self.self_attn = Attention(args)
        self.mlp = MLP(args)

    def __call__(self, x, mask=None, cache=None, **kwargs):
        x = x + self.self_attn(x, mask=mask, cache=cache, **kwargs)
        x = x + self.mlp(x)
        return x


class FalconPerceptionTransformerModel(nn.Module):
    def __init__(self, args: TextConfig, config: ModelConfig):
        super().__init__()
        self.args = args
        self.config = config
        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)

        patch_dim = (
            config.vision_config.temporal_patch_size
            * (config.vision_config.spatial_patch_size**2)
            * config.vision_config.channel_size
        )
        self.img_projector = nn.Linear(patch_dim, args.hidden_size, bias=False)
        self.layers = [DecoderLayer(args) for _ in range(args.num_hidden_layers)]
        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)

        rope_dim = args.head_dim // 2
        cos_1d, sin_1d = precompute_freqs_1d(
            rope_dim, args.max_position_embeddings, args.rope_theta
        )
        self.cos_1d = cos_1d
        self.sin_1d = sin_1d
        self.freqs_cis_golden = mx.zeros((args.num_attention_heads, rope_dim // 2, 2))
        self._last_hidden_state = None

    def __call__(
        self,
        inputs: mx.array,
        inputs_embeds: Optional[mx.array] = None,
        mask: Optional[mx.array] = None,
        cache=None,
        position_ids: Optional[mx.array] = None,
        pos_hw: Optional[mx.array] = None,
        **kwargs,
    ):
        if inputs_embeds is None:
            h = self.embed_tokens(inputs)
        else:
            h = inputs_embeds

        if cache is None:
            cache = [None] * len(self.layers)

        B, L, _ = h.shape

        if position_ids is None:
            offset = 0
            if cache[0] is not None:
                offset = cache[0].offset
                if isinstance(offset, mx.array):
                    offset = offset.item()
            position_ids = mx.arange(offset, offset + L)

        if position_ids.ndim > 1:
            pos_t = position_ids.reshape(-1)[:L]
        else:
            pos_t = position_ids[:L]

        cos_1d = self.cos_1d[pos_t]
        sin_1d = self.sin_1d[pos_t]

        cos_2d, sin_2d = None, None
        if pos_hw is not None:
            cos_2d, sin_2d = compute_golden_freqs(self.freqs_cis_golden, pos_hw)

        for layer, c in zip(self.layers, cache):
            h = layer(
                h,
                mask=mask,
                cache=c,
                cos_1d=cos_1d,
                sin_1d=sin_1d,
                cos_2d=cos_2d,
                sin_2d=sin_2d,
            )

        h = self.norm(h)
        self._last_hidden_state = h
        return h


class LanguageModel(nn.Module):
    def __init__(self, args: TextConfig, config: ModelConfig = None):
        super().__init__()
        self.args = args
        self.config = config
        self.model_type = args.model_type
        self.model = FalconPerceptionTransformerModel(args, config)
        if not args.tie_word_embeddings:
            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
        self._rope_delta = None
        self._position_ids = None
        self._pos_hw = None
        self._full_attn_mask = None
        # Set by parent Model via object.__setattr__:
        # _coord_encoder, _coord_decoder, _size_encoder, _size_decoder, _perception_config
        self._pending_coord_xy = None
        self._pending_size_hw = None
        self._detections = []
        self._current_det = {}

    def __call__(
        self,
        inputs: mx.array,
        inputs_embeds: Optional[mx.array] = None,
        mask: Optional[mx.array] = None,
        cache=None,
        **kwargs,
    ):
        kwargs.pop("image_grid_hw", None)
        kwargs.pop("pixel_values", None)
        kwargs.pop("cross_attention_states", None)

        cache_offset = 0
        if cache and cache[0] is not None:
            offset = cache[0].offset
            if isinstance(offset, int):
                cache_offset = offset
            elif isinstance(offset, mx.array):
                cache_offset = (offset if offset.ndim == 0 else offset[0]).item()

        if inputs_embeds is not None:
            L = inputs_embeds.shape[1]
        elif inputs.ndim > 1:
            L = inputs.shape[1]
        else:
            L = 1

        # Auto coord/size feedback during decode
        has_perception = hasattr(self, "_perception_config")
        if has_perception and inputs_embeds is None and cache_offset > 0 and L == 1:
            cfg = self._perception_config
            input_ids = inputs if inputs.ndim == 2 else inputs[None]
            token_id = input_ids.reshape(-1)[0].item()

            # Accumulate detection from the token just generated
            if token_id == cfg.coord_token_id and self._pending_coord_xy is not None:
                xy = self._pending_coord_xy
                self._current_det["xy"] = {"x": xy[0, 0].item(), "y": xy[0, 1].item()}
            elif token_id == cfg.size_token_id and self._pending_size_hw is not None:
                hw = self._pending_size_hw
                self._current_det["hw"] = {"h": hw[0, 0].item(), "w": hw[0, 1].item()}
            elif token_id == cfg.seg_token_id:
                # Decode segmentation mask via parent model's segm pipeline
                if hasattr(self, "_proj_segm"):
                    # Get parent Model to lazily compute segm features
                    # We find it by going: self (LM) is model.language_model
                    # The segm features are on the parent Model
                    parent = getattr(self, "_parent_model_ref", None)
                    if parent is not None:
                        parent._ensure_segm_features()
                        if parent._segm_features is not None:
                            h_state = self.model._last_hidden_state
                            if h_state is not None:
                                seg_h = h_state[:, -1, :].reshape(-1)
                                seg_mask = parent.decode_segm_mask(
                                    seg_h,
                                    parent._segm_features,
                                    parent._orig_hw[0],
                                    parent._orig_hw[1],
                                )
                                mx.eval(seg_mask)
                                self._current_det["mask"] = seg_mask
                if "xy" in self._current_det and "hw" in self._current_det:
                    self._detections.append(self._current_det)
                self._current_det = {}

            # Apply coord/size encoding to embeddings
            embeds = self.model.embed_tokens(input_ids)
            if self._pending_coord_xy is not None:
                coord_mask = input_ids == cfg.coord_token_id
                if mx.any(coord_mask).item():
                    encoded = self._coord_encoder(self._pending_coord_xy.reshape(-1, 2))
                    encoded = encoded.reshape(1, 1, -1)
                    embeds = mx.where(mx.expand_dims(coord_mask, -1), encoded, embeds)
            if self._pending_size_hw is not None:
                size_mask = input_ids == cfg.size_token_id
                if mx.any(size_mask).item():
                    encoded = self._size_encoder(self._pending_size_hw.reshape(-1, 2))
                    encoded = encoded.reshape(1, 1, -1)
                    embeds = mx.where(mx.expand_dims(size_mask, -1), encoded, embeds)
            self._pending_coord_xy = None
            self._pending_size_hw = None
            inputs_embeds = embeds

        position_ids = None
        pos_hw = None

        is_prefill = inputs_embeds is not None and L > 1
        if is_prefill and self._position_ids is not None:
            position_ids = self._position_ids[cache_offset : cache_offset + L]
            if self._pos_hw is not None:
                pos_hw = self._pos_hw[:, cache_offset : cache_offset + L, :]
        elif self._rope_delta is not None and cache_offset > 0:
            start = cache_offset + self._rope_delta
            position_ids = mx.arange(start, start + L, dtype=mx.int32)

        if mask is None and self._full_attn_mask is not None and L > 1:
            end = cache_offset + L
            mask = self._full_attn_mask[:, :, cache_offset:end, :end]

        out = self.model(
            inputs,
            cache=cache,
            inputs_embeds=inputs_embeds,
            mask=mask,
            position_ids=position_ids,
            pos_hw=pos_hw,
        )
        if self.args.tie_word_embeddings:
            out = self.model.embed_tokens.as_linear(out)
        else:
            out = self.lm_head(out)

        # After prefill, save hidden state for lazy segm feature computation
        if has_perception and L > 1:
            parent = getattr(self, "_parent_model_ref", None)
            if parent is not None and not parent._segm_features_computed:
                parent._prefill_hidden_state = self.model._last_hidden_state
                mx.eval(parent._prefill_hidden_state)

        # After forward, decode coord/size for next step's feedback
        if has_perception:
            import math

            h = self.model._last_hidden_state
            if h is not None:
                h_last = h[:, -1, :]

                # Coord decode with deduplication
                coord_logits = self._coord_decoder(h_last)
                half_c = coord_logits.shape[-1] // 2
                coord_logits = coord_logits.reshape(-1, 2, half_c)
                cl = coord_logits[0]  # (2, half_c)

                # Suppress already-detected coordinates
                existing = [d["xy"] for d in self._detections if "xy" in d]
                threshold = 0.01
                for _ in range(100):
                    bins = mx.argmax(cl, axis=-1)
                    px = bins[0].item() / (half_c - 1)
                    py = bins[1].item() / (half_c - 1)
                    is_dup = any(
                        abs(e["x"] - px) < threshold and abs(e["y"] - py) < threshold
                        for e in existing
                    )
                    if not is_dup:
                        break
                    cl = cl.at[0, bins[0]].add(-1e9)
                    cl = cl.at[1, bins[1]].add(-1e9)
                self._pending_coord_xy = mx.array([[px, py]])

                size_logits = self._size_decoder(h_last)
                half_s = size_logits.shape[-1] // 2
                size_logits = size_logits.reshape(-1, 2, half_s)
                pred = mx.argmax(size_logits, axis=-1).astype(mx.float32) / (half_s - 1)
                min_sz = math.log2(1.0 / half_s)
                pred = pred * (0.0 - min_sz) + min_sz
                self._pending_size_hw = 2.0**pred

        return LanguageModelOutput(
            logits=out.astype(self.model.embed_tokens.weight.dtype),
        )

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

    @property
    def head_dim(self):
        return self.args.head_dim

    @property
    def n_kv_heads(self):
        return self.args.num_key_value_heads
