"""From-scratch MLX reimplementation of microsoft/table-transformer-detection
(ResNet18 backbone + DETR-style encoder-decoder with learned object queries).

No mlx-vlm registry covers DETR-family detection models, so every module is
reimplemented directly against the HF `transformers` source
(`modeling_table_transformer.py`).

Simplification: this implementation assumes a full (unpadded) pixel_mask --
i.e. a single square image resized to exactly (IMAGE_SIZE, IMAGE_SIZE), never
a batch with letterboxed padding. Under that assumption every attention mask
in the original model is a no-op (all-valid), so encoder self-attention,
decoder cross-attention, and the sine position embedding's mask-dependent
cumsum all simplify to their unmasked form. This does NOT support batched
inference with mixed aspect ratios/padding.
"""

import math

import mlx.core as mx
import mlx.nn as nn
import numpy as np

D_MODEL = 256
NUM_HEADS = 8
FFN_DIM = 2048
NUM_ENCODER_LAYERS = 6
NUM_DECODER_LAYERS = 6
NUM_QUERIES = 15
NUM_LABELS_PLUS_ONE = 3  # "table", "table rotated", "no object"
IMAGE_SIZE = 800
BN_EPS = 1e-5


def frozen_bn(x, weight, bias, mean, var):
    scale = weight / mx.sqrt(var + BN_EPS)
    shift = bias - mean * scale
    return x * scale + shift


class BasicBlock(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1, has_downsample=False):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1_weight = mx.ones((out_ch,))
        self.bn1_bias = mx.zeros((out_ch,))
        self.bn1_mean = mx.zeros((out_ch,))
        self.bn1_var = mx.ones((out_ch,))
        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2_weight = mx.ones((out_ch,))
        self.bn2_bias = mx.zeros((out_ch,))
        self.bn2_mean = mx.zeros((out_ch,))
        self.bn2_var = mx.ones((out_ch,))
        self.has_downsample = has_downsample
        if has_downsample:
            self.down_conv = nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=stride, bias=False)
            self.down_bn_weight = mx.ones((out_ch,))
            self.down_bn_bias = mx.zeros((out_ch,))
            self.down_bn_mean = mx.zeros((out_ch,))
            self.down_bn_var = mx.ones((out_ch,))

    def __call__(self, x):
        residual = x
        h = self.conv1(x)
        h = frozen_bn(h, self.bn1_weight, self.bn1_bias, self.bn1_mean, self.bn1_var)
        h = nn.relu(h)
        h = self.conv2(h)
        h = frozen_bn(h, self.bn2_weight, self.bn2_bias, self.bn2_mean, self.bn2_var)
        if self.has_downsample:
            residual = self.down_conv(x)
            residual = frozen_bn(
                residual, self.down_bn_weight, self.down_bn_bias, self.down_bn_mean, self.down_bn_var
            )
        return nn.relu(h + residual)


class ResNet18Backbone(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1_weight = mx.ones((64,))
        self.bn1_bias = mx.zeros((64,))
        self.bn1_mean = mx.zeros((64,))
        self.bn1_var = mx.ones((64,))
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        self.layer1 = [BasicBlock(64, 64), BasicBlock(64, 64)]
        self.layer2 = [BasicBlock(64, 128, stride=2, has_downsample=True), BasicBlock(128, 128)]
        self.layer3 = [BasicBlock(128, 256, stride=2, has_downsample=True), BasicBlock(256, 256)]
        self.layer4 = [BasicBlock(256, 512, stride=2, has_downsample=True), BasicBlock(512, 512)]

    def __call__(self, x):
        x = self.conv1(x)
        x = frozen_bn(x, self.bn1_weight, self.bn1_bias, self.bn1_mean, self.bn1_var)
        x = nn.relu(x)
        x = self.maxpool(x)
        for block in self.layer1:
            x = block(x)
        for block in self.layer2:
            x = block(x)
        for block in self.layer3:
            x = block(x)
        for block in self.layer4:
            x = block(x)
        return x  # (B, H/32, W/32, 512)


def sine_position_embedding(height, width, embedding_dim=128, temperature=10000.0, scale=2 * math.pi):
    """Unmasked (all-ones pixel_mask) version: y_embed/x_embed are just the
    cumulative row/column index, 1-indexed, normalized to [~0, scale]."""
    y_embed = np.arange(1, height + 1, dtype=np.float32)
    x_embed = np.arange(1, width + 1, dtype=np.float32)
    y_embed = y_embed / (height + 1e-6) * scale
    x_embed = x_embed / (width + 1e-6) * scale

    dim_t = np.arange(embedding_dim, dtype=np.float32)
    dim_t = temperature ** (2 * (dim_t // 2) / embedding_dim)

    pos_x = x_embed[None, :, None] / dim_t  # (1, W, D)
    pos_y = y_embed[:, None, None] / dim_t  # (H, 1, D)
    pos_x = np.broadcast_to(pos_x, (height, width, embedding_dim))
    pos_y = np.broadcast_to(pos_y, (height, width, embedding_dim))

    def interleave_sin_cos(p):
        s = np.sin(p[..., 0::2])
        c = np.cos(p[..., 1::2])
        out = np.stack([s, c], axis=-1).reshape(p.shape[:-1] + (p.shape[-1],))
        return out

    pos_x = interleave_sin_cos(pos_x)
    pos_y = interleave_sin_cos(pos_y)
    pos = np.concatenate([pos_y, pos_x], axis=-1)  # (H, W, 2*D) = (H, W, d_model)
    return mx.array(pos[None])  # (1, H, W, d_model)


class MHA(nn.Module):
    """DETR-style attention: object_queries added to Q/K (not V) before projection."""

    def __init__(self, embed_dim=D_MODEL, num_heads=NUM_HEADS):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.scale = self.head_dim**-0.5
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def __call__(self, hidden_states, object_queries=None, key_value_states=None, spatial_position_embeddings=None):
        b, tgt_len, c = hidden_states.shape
        q_in = hidden_states if object_queries is None else hidden_states + object_queries

        if key_value_states is not None:
            kv_in = key_value_states if spatial_position_embeddings is None else key_value_states + spatial_position_embeddings
            v_in = key_value_states
        else:
            kv_in = q_in
            v_in = hidden_states

        src_len = kv_in.shape[1]
        q = self.q_proj(q_in).reshape(b, tgt_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        k = self.k_proj(kv_in).reshape(b, src_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        v = self.v_proj(v_in).reshape(b, src_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)

        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale)
        out = out.transpose(0, 2, 1, 3).reshape(b, tgt_len, c)
        return self.out_proj(out)


class EncoderLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.self_attn = MHA()
        self.self_attn_layer_norm = nn.LayerNorm(D_MODEL)
        self.fc1 = nn.Linear(D_MODEL, FFN_DIM)
        self.fc2 = nn.Linear(FFN_DIM, D_MODEL)
        self.final_layer_norm = nn.LayerNorm(D_MODEL)

    def __call__(self, x, object_queries):
        residual = x
        h = self.self_attn_layer_norm(x)
        h = self.self_attn(h, object_queries=object_queries)
        x = residual + h

        residual = x
        h = self.final_layer_norm(x)
        h = self.fc2(nn.relu(self.fc1(h)))
        x = residual + h
        return x


class DecoderLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.self_attn = MHA()
        self.self_attn_layer_norm = nn.LayerNorm(D_MODEL)
        self.encoder_attn = MHA()
        self.encoder_attn_layer_norm = nn.LayerNorm(D_MODEL)
        self.fc1 = nn.Linear(D_MODEL, FFN_DIM)
        self.fc2 = nn.Linear(FFN_DIM, D_MODEL)
        self.final_layer_norm = nn.LayerNorm(D_MODEL)

    def __call__(self, x, query_position_embeddings, encoder_hidden_states, object_queries):
        residual = x
        h = self.self_attn_layer_norm(x)
        h = self.self_attn(h, object_queries=query_position_embeddings)
        x = residual + h

        residual = x
        h = self.encoder_attn_layer_norm(x)
        h = self.encoder_attn(
            h,
            object_queries=query_position_embeddings,
            key_value_states=encoder_hidden_states,
            spatial_position_embeddings=object_queries,
        )
        x = residual + h

        residual = x
        h = self.final_layer_norm(x)
        h = self.fc2(nn.relu(self.fc1(h)))
        x = residual + h
        return x


class TableTransformerMLX(nn.Module):
    def __init__(self, num_queries=NUM_QUERIES, image_size=IMAGE_SIZE):
        super().__init__()
        self.backbone = ResNet18Backbone()
        self.input_projection = nn.Conv2d(512, D_MODEL, kernel_size=1)
        self.query_position_embeddings = mx.zeros((num_queries, D_MODEL))

        self.encoder_layers = [EncoderLayer() for _ in range(NUM_ENCODER_LAYERS)]
        self.encoder_layernorm = nn.LayerNorm(D_MODEL)
        self.decoder_layers = [DecoderLayer() for _ in range(NUM_DECODER_LAYERS)]
        self.decoder_layernorm = nn.LayerNorm(D_MODEL)

        self.class_labels_classifier = nn.Linear(D_MODEL, NUM_LABELS_PLUS_ONE)
        self.bbox_fc0 = nn.Linear(D_MODEL, D_MODEL)
        self.bbox_fc1 = nn.Linear(D_MODEL, D_MODEL)
        self.bbox_fc2 = nn.Linear(D_MODEL, 4)

        feat_size = image_size // 32
        self._pos_embed = sine_position_embedding(feat_size, feat_size, embedding_dim=D_MODEL // 2)

    def __call__(self, pixel_values_nhwc):
        b = pixel_values_nhwc.shape[0]
        feature_map = self.backbone(pixel_values_nhwc)  # (B, H', W', 512)
        h, w = feature_map.shape[1], feature_map.shape[2]

        projected = self.input_projection(feature_map)  # (B, H', W', 256)
        flattened = projected.reshape(b, h * w, D_MODEL)

        object_queries = mx.broadcast_to(self._pos_embed.reshape(1, h * w, D_MODEL), (b, h * w, D_MODEL))

        x = flattened
        for layer in self.encoder_layers:
            x = layer(x, object_queries)
        encoder_out = self.encoder_layernorm(x)

        query_pos = mx.broadcast_to(self.query_position_embeddings[None], (b, NUM_QUERIES, D_MODEL))
        queries = mx.zeros((b, NUM_QUERIES, D_MODEL))

        y = queries
        for layer in self.decoder_layers:
            y = layer(y, query_pos, encoder_out, object_queries)
        decoder_out = self.decoder_layernorm(y)

        logits = self.class_labels_classifier(decoder_out)
        boxes = nn.relu(self.bbox_fc0(decoder_out))
        boxes = nn.relu(self.bbox_fc1(boxes))
        boxes = self.bbox_fc2(boxes)
        boxes = mx.sigmoid(boxes)
        return logits, boxes