"""Image-reference CNN VAE encoder (single-frame causal convolutions in MLX).

For a single frame, causal temporal zero padding selects the last temporal
kernel plane. This is the exact 2D reduction of the released 3D CNN, including
its spatial reflection padding, group norms and sampled posterior.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path

import mlx.core as mx
import numpy as np

from .layout import patchify
from .ops import silu
from .vae import assemble_tiles, split_tiles
from .vision import bound_reference
from .weights import WeightStore

log = logging.getLogger(__name__)


def reflect_pad(x, left, right, top, bottom):
    # NHWC; explicit slices avoid NumPy/CPU roundtrips and support Metal.
    if top or bottom:
        parts = []
        if top:
            parts.append(x[:, 1 : top + 1][:, ::-1])
        parts.append(x)
        if bottom:
            parts.append(x[:, -bottom - 1 : -1][:, ::-1])
        x = mx.concatenate(parts, axis=1)
    if left or right:
        parts = []
        if left:
            parts.append(x[:, :, 1 : left + 1][:, :, ::-1])
        parts.append(x)
        if right:
            parts.append(x[:, :, -right - 1 : -1][:, :, ::-1])
        x = mx.concatenate(parts, axis=2)
    return x


class ReferenceEncoder:
    def __init__(self, path):
        self.store = WeightStore(path, cache_gb=2)
        self.config = json.loads((Path(path) / "config.json").read_text())

    def conv(self, x, prefix, *, stride=1, pad=1, down=False):
        weight = self.store.get(prefix + ".weight")
        if weight.ndim == 5:
            weight = weight[:, :, -1]
        weight = weight.transpose(0, 2, 3, 1)
        if down:
            x = reflect_pad(x, 0, 1, 0, 1)
        elif pad:
            x = reflect_pad(x, pad, pad, pad, pad)
        out = mx.conv2d(x, weight, stride=stride)
        return out + self.store.get(prefix + ".bias")

    def norm(self, x, prefix):
        b, h, w, c = x.shape
        # GroupNorm reduces spatial dimensions and channels within each group.
        f = x.astype(mx.float32).reshape(b, h, w, 32, c // 32)
        mean = mx.mean(f, axis=(1, 2, 4), keepdims=True)
        var = mx.mean((f - mean) ** 2, axis=(1, 2, 4), keepdims=True)
        f = ((f - mean) * mx.rsqrt(var + 1e-6)).reshape(x.shape)
        return f * self.store.get(prefix + ".weight") + self.store.get(prefix + ".bias")

    def resblock(self, x, prefix):
        h = self.conv(silu(self.norm(x, prefix + ".norm1")), prefix + ".conv1")
        h = self.conv(silu(self.norm(h, prefix + ".norm2")), prefix + ".conv2")
        if prefix + ".nin_shortcut.weight" in self.store:
            x = self.conv(x, prefix + ".nin_shortcut", pad=0)
        return h + x

    def encode_tile(self, x):
        x = self.conv(x, "encoder.conv_in")
        for i in range(6):
            for j in range(2):
                x = self.resblock(x, f"encoder.down.{i}.block.{j}")
                mx.eval(x)
            if i < 4:
                x = self.conv(
                    x, f"encoder.down.{i}.downsample.conv", stride=2, pad=0, down=True
                )
                mx.eval(x)
        x = self.conv(silu(self.norm(x, "encoder.norm_out")), "encoder.conv_out")
        x = self.conv(x, "quant_conv", pad=0)
        mx.eval(x)
        return x.transpose(0, 3, 1, 2)

    def encode(self, image):
        # Native conditioning CNN runs FP32, independently from decoder BF16.
        image = bound_reference(image, 1024, 2048)
        w, h = image.size
        w, h = w // 16 * 16, h // 16 * 16
        array = np.asarray(image)[:h, :w].astype(np.float32) / 255
        x = mx.array(
            (array - np.array([0.485, 0.456, 0.406], np.float32))
            / np.array([0.229, 0.224, 0.225], np.float32)
        )[None]
        ys, yl, yo = split_tiles(h)
        xs, xl, xo = split_tiles(w)
        rows = []
        for y, hh in zip(ys, yl, strict=False):
            row = []
            for xx, ww in zip(xs, xl, strict=False):
                row.append(self.encode_tile(x[:, y : y + hh, xx : xx + ww]))
            rows.append(row)
        moments = assemble_tiles(rows, [n // 16 for n in yo], [n // 16 for n in xo])
        mx.eval(moments)
        return moments


def encode_references(path, images, *, generator, seed):
    encoder = ReferenceEncoder(path)
    bounded = [bound_reference(im, 1024, 2048) for im in images]
    grids = [(im.height // 32 * 2, im.width // 32 * 2) for im in bounded]
    # Consume every anchor draw before any posterior draw, as in H3 native.
    if generator is not None:
        import torch

        anchor_noise = [
            mx.array(
                torch.randn(
                    1, 24, h, w, generator=generator, dtype=torch.float32
                ).numpy()
            )
            for h, w in grids
        ]
    else:
        anchor_noise = [
            mx.random.normal((1, 24, h, w), key=mx.random.key(seed + 1 + i))
            for i, (h, w) in enumerate(grids)
        ]
    anchors = []
    cfg = encoder.config
    mean = mx.array(cfg["latents_mean"], mx.float32)[None, :, None, None]
    std = mx.array(cfg["latents_std"], mx.float32)[None, :, None, None]
    for i, (image, (h, w), noise) in enumerate(
        zip(bounded, grids, anchor_noise, strict=False)
    ):
        moments = encoder.encode(image)
        mu, logvar = mx.split(moments, 2, axis=1)
        logvar = mx.clip(logvar, -30, 20)
        if generator is not None:
            # Native posterior uses CPU torch noise then moves to execution device.
            draw = mx.array(
                torch.randn(
                    tuple(mu.shape), generator=generator, dtype=torch.float32
                ).numpy()
            )
        else:
            draw = mx.random.normal(mu.shape, key=mx.random.key(seed + 100 + i))
        z = mu + mx.exp(0.5 * logvar) * draw
        # The upstream encode_images contract rounds its sampled latent to FP16.
        z = z.astype(mx.float16).astype(mx.float32)[:, :, :h, :w]
        clean = patchify((z - mean) / std)
        t = mx.array(0.999, mx.float32)
        anchors.append(t * clean + (1 - t) * patchify(noise))
        mx.eval(anchors[-1])
        log.info("Encoded edit reference %d/%d", i + 1, len(images))
    encoder.store.clear()
    return tuple(grids), mx.concatenate(anchors)