import hashlib
import json

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

from h3_mlx.compare import make_comparison
from h3_mlx.reference import reflect_pad
from h3_mlx.vae import assemble_tiles, split_tiles, to_image


def test_tiling_native_overlap_and_coverage():
    for length in (32, 256, 512, 1024, 2048, 4096):
        starts, lens, overlaps = split_tiles(length)
        assert starts[0] == 0 and starts[-1] + lens[-1] == length
        assert sum(lens) - sum(overlaps) == length
        assert all(v % 16 == 0 for v in overlaps)
    assert split_tiles(512) == ([0, 128, 256], [256] * 3, [128, 128])


def test_native_vertical_then_horizontal_overlap_blending():
    rows = [
        [mx.full((1, 1, 4, 4), i * 2 + j, mx.float32) for j in range(2)]
        for i in range(2)
    ]
    out = np.asarray(assemble_tiles(rows, [2], [2]))[0, 0]
    # Hand-computed native blend: row 1/col 1 at its first overlap pixel
    # inherits the original tile above, then the original tile to its left.
    assert out.shape == (6, 6)
    assert out[0, 0] == 0 and out[-1, -1] == 3
    assert out[2, 2] == 2  # horizontal blend takes the unblended left source
    assert out[3, 3] == 2


def test_reflection_pad_matches_pytorch():
    x = torch.arange(2 * 3 * 4 * 2, dtype=torch.float32).reshape(2, 3, 4, 2)
    out = reflect_pad(mx.array(x.numpy()), 1, 1, 1, 1)
    ref = torch.nn.functional.pad(
        x.permute(0, 3, 1, 2), (1, 1, 1, 1), mode="reflect"
    ).permute(0, 2, 3, 1)
    np.testing.assert_array_equal(np.asarray(out), ref.numpy())


def test_comparison_rejects_mismatched_seed_and_retains_native_crops(tmp_path):
    a = tmp_path / "ref.png"
    b = tmp_path / "candidate.png"
    Image.new("RGB", (64, 64), (50, 70, 90)).save(a)
    Image.new("RGB", (64, 64), (51, 70, 90)).save(b)
    metadata = {
        "prompt": "test",
        "seed": 7,
        "size": [64, 64],
        "canvas": [64, 64],
        "rng": "torch",
        "reference_images": [],
    }
    a.with_suffix(".json").write_text(json.dumps(metadata))
    b.with_suffix(".json").write_text(json.dumps({**metadata, "seed": 8}))
    with pytest.raises(ValueError, match="seed"):
        make_comparison(a, [b], tmp_path / "bad")
    b.with_suffix(".json").write_text(json.dumps(metadata))
    result = make_comparison(a, [b], tmp_path / "report", crop=(1, 2, 33, 34))
    assert result["crop_xyxy"] == [1, 2, 33, 34]
    assert result["rows"][1]["mae_8bit"] == pytest.approx(1 / 3)
    assert (tmp_path / "report/stack.png").exists()
    assert (tmp_path / "report/grid.png").exists()
    a.with_suffix(".trace.safetensors").write_bytes(b"reference latent")
    decoder = {
        **metadata,
        "validation_scope": "Same final latent; decoder only",
        "source_trace_sha256": "different latent",
    }
    b.with_suffix(".json").write_text(json.dumps(decoder))
    with pytest.raises(ValueError, match="final latent"):
        make_comparison(a, [b], tmp_path / "wrong-latent")
    decoder["source_trace_sha256"] = hashlib.sha256(b"reference latent").hexdigest()
    b.with_suffix(".json").write_text(json.dumps(decoder))
    verified = make_comparison(a, [b], tmp_path / "decoder")
    assert "decoder-only" in verified["validation_scope"]


def test_output_quantization_truncates_like_native_writer():
    x = mx.array([[[[0.5, 0.501]], [[0.25, 0.251]], [[0.75, 0.751]]]])
    image = to_image(x, 2, 1)
    np.testing.assert_array_equal(np.asarray(image), [[[127, 63, 191], [127, 64, 191]]])