from __future__ import annotations

import mlx.core as mx
import numpy as np


def patchify(latent):
    """BCHW latent -> [B*H/2*W/2, C*2*2], channel then patch coordinates."""
    b, c, h, w = latent.shape
    if h % 2 or w % 2:
        raise ValueError("Latent height/width must be even")
    return (
        latent.reshape(b, c, h // 2, 2, w // 2, 2)
        .transpose(0, 2, 4, 1, 3, 5)
        .reshape(-1, c * 4)
    )


def unpatchify(rows, h, w, channels=24):
    return (
        rows.reshape(1, h // 2, w // 2, channels, 2, 2)
        .transpose(0, 3, 1, 4, 2, 5)
        .reshape(1, channels, h, w)
    )


def image_positions(text_len, h, w, references=()):
    """H3 ref2va coordinates, generated in float64 on CPU then cast once.

    Padding is omitted: upstream packs it as a separate attention document.
    Real rows cannot attend to it, so it has no effect on these outputs.
    """
    text = np.zeros((text_len, 3), np.float64)
    text[:, 0] = np.arange(text_len)
    result = [text]
    for i, (hh, ww) in enumerate((*references, (h, w))):
        area = np.sqrt(hh * ww)
        axes = []
        for dim in (hh, ww):
            ratio = dim / area
            left = (1 - ratio) / 2
            axes.append(np.linspace(left, left + ratio, dim // 2, endpoint=False) * 32)
        yy, xx = np.meshgrid(*axes, indexing="ij")
        coords = np.stack((np.full_like(yy, text_len + i), yy, xx), axis=-1).reshape(
            -1, 3
        )
        result.append(coords)
    return mx.array(np.concatenate(result).astype(np.float32))


def rope_cache(positions, inv_freq, dtype=mx.bfloat16):
    half = (positions[..., None] * inv_freq).reshape(positions.shape[0], -1)
    angles = mx.concatenate((half, half), axis=-1)
    return mx.cos(angles).astype(dtype)[None, :, None], mx.sin(angles).astype(dtype)[
        None, :, None
    ]