from __future__ import annotations

import gc
import json
import logging
import math
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image

from .ops import attention, layer_norm, rms_norm, rotate, swiglu
from .weights import Linear, WeightStore

log = logging.getLogger(__name__)


def split_tiles(length, tile=256, overlap=64, ratio=16):
    if tile <= overlap or tile % ratio or overlap % ratio:
        raise ValueError(
            "Tile size/overlap must be multiples of 16, with tile > overlap"
        )
    if length <= tile:
        return [0], [length], []
    n = math.ceil(length / tile)
    while tile * n - overlap * (n - 1) < length:
        n += 1
    overlaps = [overlap] * (n - 1)
    for i in range((tile * n - sum(overlaps) - length) // ratio):
        overlaps[i % (n - 1)] += ratio
    starts = [0]
    for ov in overlaps:
        starts.append(starts[-1] + tile - ov)
    return starts, [tile] * n, overlaps


def blend(a, b, extent, axis):
    extent = min(a.shape[axis], b.shape[axis], extent)
    if extent == 0:
        return b
    shape = [1] * b.ndim
    shape[axis] = extent
    weight = mx.arange(extent, dtype=b.dtype).reshape(shape) / extent
    aa, bb, rest = (
        [slice(None)] * b.ndim,
        [slice(None)] * b.ndim,
        [slice(None)] * b.ndim,
    )
    aa[axis], bb[axis], rest[axis] = (
        slice(-extent, None),
        slice(0, extent),
        slice(extent, None),
    )
    mixed = a[tuple(aa)] * (1 - weight) + b[tuple(bb)] * weight
    return mx.concatenate((mixed, b[tuple(rest)]), axis=axis)


def assemble_tiles(rows, y_overlap, x_overlap):
    assembled = []
    for i, row in enumerate(rows):
        pieces = []
        for j, tile in enumerate(row):
            if i:
                tile = blend(rows[i - 1][j], tile, y_overlap[i - 1], -2)
            if j:
                tile = blend(row[j - 1], tile, x_overlap[j - 1], -1)
            if i < len(rows) - 1:
                tile = tile[..., : -y_overlap[i], :]
            if j < len(row) - 1:
                tile = tile[..., : -x_overlap[j]]
            pieces.append(tile)
        assembled.append(mx.concatenate(pieces, axis=-1))
    return mx.concatenate(assembled, axis=-2)


class DecoderBlock:
    def __init__(self, store, index, dtype):
        p = f"decoder.transformer_blocks.{index}"
        self.n1, self.n2 = (store.get(p + f".norm{i}.weight") for i in (1, 2))
        self.s1, self.s2 = (store.get(p + f".scale{i}") for i in (1, 2))
        self.qkv = Linear.load(store, p + ".attn.to_qkv")
        self.out = Linear.load(store, p + ".attn.to_out")
        self.w1 = Linear.load(store, p + ".ff.w1")
        self.w2 = Linear.load(store, p + ".ff.w2")

    def __call__(self, x, rope):
        b, n, _ = x.shape
        h = rms_norm(x, self.n1)
        # VAE QKV is head-interleaved too, but keep its native layout here.
        qkv = self.qkv(h).reshape(b, n, 16, 3, 64)
        q, k, v = (qkv[:, :, :, i] for i in range(3))
        q, k = rms_norm(q), rms_norm(k)
        q, k = rotate(q, *rope), rotate(k, *rope)
        a = self.out(attention(q, k, v).reshape(b, n, -1))
        x = x + a * self.s1
        h = rms_norm(x, self.n2)
        return x + self.w2(swiglu(self.w1(h))) * self.s2


class VAE:
    def __init__(
        self,
        path,
        *,
        tile_batch=1,
        dtype="bfloat16",
        tile_size=256,
        compile_blocks=False,
    ):
        self.path = Path(path)
        self.config = json.loads((self.path / "config.json").read_text())
        self.dtype = {"bfloat16": mx.bfloat16, "float32": mx.float32}[dtype]
        self.store = WeightStore(self.path, cast=self.dtype)
        self.tile_batch, self.tile_size = tile_batch, tile_size
        self.blocks = None
        self.compile_blocks = compile_blocks
        self.runners = []

    def _load(self):
        if self.blocks is not None:
            return
        s = self.store
        self.post = Linear(
            s.get("post_quant_conv.weight").reshape(24, 24),
            s.get("post_quant_conv.bias"),
        )
        self.embed = Linear.load(s, "decoder.x_embedder")
        self.register = s.get("decoder.register_tokens")
        self.norm_w, self.norm_b = (
            s.get("decoder.norm_out.weight"),
            s.get("decoder.norm_out.bias"),
        )
        self.out = Linear.load(s, "decoder.proj_out")
        self.blocks = [DecoderBlock(s, i, self.dtype) for i in range(24)]
        if self.compile_blocks:
            for block in self.blocks:

                def forward(x, rope, block=block):
                    return block(x, rope)

                self.runners.append(mx.compile(forward))

    def decode_tiles(self, raw):
        self._load()
        b, c, h, w = raw.shape
        # Native image decoder is BF16, including its pointwise convolution.
        x = raw.transpose(0, 2, 3, 1).astype(self.dtype)
        x = self.embed(self.post(x)).reshape(b, h * w, 1024)
        x = mx.concatenate(
            (
                x,
                mx.broadcast_to(self.register, (b, 4, 1024)),
                mx.zeros((b, 1, 1024), self.dtype),
            ),
            axis=1,
        )
        # Coordinates start in latent dtype (BF16), as in released ViT2D.
        yy = 2 * ((mx.arange(h, dtype=self.dtype) + mx.array(0.5, self.dtype)) / h) - 1
        xx = 2 * ((mx.arange(w, dtype=self.dtype) + mx.array(0.5, self.dtype)) / w) - 1
        yy, xx = mx.meshgrid(yy, xx, indexing="ij")
        positions = mx.stack((yy, xx), axis=-1).reshape(-1, 2)
        positions = mx.concatenate((positions, mx.zeros((5, 2), self.dtype)), axis=0)
        inv = 1 / (100 ** mx.arange(0, 1, 0.125, dtype=mx.float32))
        scaled = (2 * math.pi * positions.astype(mx.float32)).astype(self.dtype)
        angles = (scaled.astype(mx.float32)[..., None] * inv).reshape(h * w + 5, 16)
        angles = mx.concatenate((angles, angles), axis=-1)
        rope = (
            mx.cos(angles).astype(self.dtype)[None, :, None],
            mx.sin(angles).astype(self.dtype)[None, :, None],
        )
        for block in self.runners if self.compile_blocks else self.blocks:
            x = block(x, rope)
            if not self.compile_blocks:
                mx.eval(x)
        x = layer_norm(x, self.norm_w, self.norm_b)
        x = self.out(x)[:, : h * w]
        # Output channels are [C, patchH, patchW], matching the VAE checkpoint.
        return (
            x.reshape(b, h, w, 3, 16, 16)
            .transpose(0, 3, 1, 4, 2, 5)
            .reshape(b, 3, h * 16, w * 16)
        )

    def decode(self, latent, *, on_batch=None):
        c = self.config
        mean = mx.array(c["latents_mean"], mx.float32)[None, :, None, None]
        std = mx.array(c["latents_std"], mx.float32)[None, :, None, None]
        raw = latent.astype(mx.float32) * std + mean
        h, w = raw.shape[-2:]
        ys, ylen, yo = split_tiles(h * 16, self.tile_size)
        xs, xlen, xo = split_tiles(w * 16, self.tile_size)
        tasks = [
            (i, j, raw[..., y // 16 : (y + yl) // 16, x // 16 : (x + xl) // 16])
            for i, (y, yl) in enumerate(zip(ys, ylen, strict=False))
            for j, (x, xl) in enumerate(zip(xs, xlen, strict=False))
        ]
        rows = [[None for _ in xs] for _ in ys]
        for start in range(0, len(tasks), self.tile_batch):
            group = tasks[start : start + self.tile_batch]
            tiles = mx.concatenate([t for _, _, t in group], axis=0)
            decoded = self.decode_tiles(tiles)
            mx.eval(decoded)
            for k, (i, j, _) in enumerate(group):
                rows[i][j] = decoded[k : k + 1]
            if on_batch:
                on_batch(start + len(group), len(tasks))
        # Decode tile overlap arithmetic stays in the decoder dtype.
        out = assemble_tiles(rows, yo, xo)
        # torchvision's inverse Normalize subtracts its inverse mean, then divides.
        inv_std = mx.array([1 / 0.229, 1 / 0.224, 1 / 0.225], self.dtype)[
            None, :, None, None
        ]
        inv_mean = mx.array(
            [-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], self.dtype
        )[None, :, None, None]
        out = mx.clip((out - inv_mean) / inv_std, 0, 1).astype(mx.float32)
        mx.eval(out)
        return out

    def clear(self):
        self.runners.clear()
        self.blocks = None
        for name in ("post", "embed", "register", "norm_w", "norm_b", "out"):
            if hasattr(self, name):
                delattr(self, name)
        self.store.clear()
        gc.collect()
        mx.clear_cache()


def to_image(pixels, width, height):
    array = np.asarray(
        mx.clip(pixels[0].transpose(1, 2, 0) * 255, 0, 255).astype(mx.uint8)
    )
    image = Image.fromarray(array)
    if image.size != (width, height):
        image = image.resize((width, height), Image.Resampling.LANCZOS)
    return image