import json

import mlx.core as mx
import numpy as np

from h3_mlx.config import DiTConfig
from h3_mlx.dit import DiT


def test_refined_conditioning_cache_is_exact_and_input_specific(tmp_path):
    model_path = tmp_path / "model"
    transformer = model_path / "transformer"
    transformer.mkdir(parents=True)
    cfg = DiTConfig(hidden_size=32, text_dim=24, token_refiner_num_layers=0)
    (transformer / "config.json").write_text(json.dumps(vars(cfg)))
    mx.random.seed(732)
    weights = {
        "condition_proj.weight": (mx.random.normal((32, 24)) * 0.1).astype(mx.bfloat16),
        "condition_proj.bias": mx.zeros(32, mx.bfloat16),
        "token_refiner.final_norm.weight": mx.ones(32, mx.bfloat16),
    }
    mx.save_safetensors(str(transformer / "model.safetensors"), weights)
    mx.save_safetensors(
        str(model_path / "pdd_fused_heads.safetensors"), {"unused": mx.zeros(1)}
    )
    model = DiT(model_path, cache_dir=tmp_path / "cache")
    text = mx.random.normal((17, 24)).astype(mx.bfloat16)
    first = model.refine(text)
    assert not model.refiner_cache_hit
    second = model.refine(text)
    assert model.refiner_cache_hit
    np.testing.assert_array_equal(
        np.asarray(first.view(mx.uint16)), np.asarray(second.view(mx.uint16))
    )
    model.refine(text + mx.array(1, mx.bfloat16))
    assert not model.refiner_cache_hit
    assert len(list((tmp_path / "cache").glob("refined-*.safetensors"))) == 2
    model.cache_refined_text = False
    direct = model.refine(text)
    assert not model.refiner_cache_hit
    np.testing.assert_array_equal(
        np.asarray(first.view(mx.uint16)), np.asarray(direct.view(mx.uint16))
    )
    model.clear()