import json
from types import SimpleNamespace

import numpy as np
import pytest
import torch
from PIL import Image

from image3d_student.data import CachedDataset, cache_sample, cell_indices, splat_targets
from image3d_student.engine import split_rows
from image3d_student.model import Student, distillation_loss
from image3d_student.splats import Splats, read_ply, write_ply
from image3d_student.teacher_idol import patch_attention


@pytest.fixture
def splats():
    return Splats(np.array([[-.2, .1, .0], [.2, -.1, .3]]),
                  np.array([[.9,.1,.2], [.1,.8,.4]]), np.array([[.01,.02,.03],[.04,.01,.02]]),
                  np.array([[1,0,0,0], [.5,.5,.5,.5]]), np.array([[.8],[.9]]))


def test_ply_roundtrip_preserves_frame_and_attributes(tmp_path, splats):
    path = tmp_path / "splats.ply"
    write_ply(path, splats)
    reread = read_ply(path)
    for name in splats.__dataclass_fields__:
        np.testing.assert_allclose(getattr(reread, name), getattr(splats, name), atol=1e-6)


def test_non_gaussian_ply_rejected(tmp_path):
    import trimesh
    path = tmp_path / "mesh.ply"
    trimesh.creation.box().export(path)
    with pytest.raises(ValueError, match="Gaussian-splat"):
        read_ply(path)


def test_coordinates_not_silently_clipped():
    with pytest.raises(ValueError, match="aligned"):
        cell_indices(np.array([[1., 0, 0]]), 8)
    assert cell_indices(np.array([[.5,.5,.5]]), 8)[0] == 511


def test_voxel_assignment_selects_highest_opacity(splats):
    splats.xyz[1] = splats.xyz[0] + .001
    target = splat_targets(splats, 8)
    assert target["occ"].sum() == 1
    np.testing.assert_allclose(target["rgb"][target["occ"][:,0] > 0][0], splats.rgb[1])


def test_cache_integrity_and_duplicate_protection(tmp_path, splats):
    image = tmp_path / "input.png"
    Image.new("RGB", (32,32), "red").save(image)
    root = tmp_path / "dataset"
    cache_sample(root, "a", image, splats, 8)
    dataset = CachedDataset(root / "manifest.jsonl")
    tensor, target = dataset[0]
    assert tensor.shape == (3,128,128)
    assert target["occ"].sum() == 2
    with pytest.raises(ValueError, match="already exists"):
        cache_sample(root, "a", image, splats, 8)
    (root / "images/a.png").write_bytes(b"changed")
    with pytest.raises(ValueError, match="hash mismatch"):
        CachedDataset(root / "manifest.jsonl")


def test_geometry_import_requires_frame_and_marks_teachers(tmp_path, splats):
    import trimesh
    image = tmp_path / "input.png"
    Image.new("RGB", (32,32), "green").save(image)
    mesh = tmp_path / "shape.obj"
    trimesh.creation.icosphere(subdivisions=1, radius=.3).export(mesh)
    with pytest.raises(ValueError, match="explicit"):
        cache_sample(tmp_path / "bad", "a", image, splats, 8, coarse=mesh)
    transform = tmp_path / "identity.json"
    transform.write_text(json.dumps(np.eye(4).tolist()))
    root = tmp_path / "good"
    row = cache_sample(root, "a", image, splats, 8, coarse=mesh, coarse_transform=transform,
                       fine=mesh, fine_transform=transform)
    assert row["teachers"] == {"appearance": "idol", "coarse": "InstantMesh", "fine": "Hi3D"}
    with np.load(root / row["target"]) as target:
        assert target["geometry_valid"].tolist() == [1,1]
        assert target["geometry"].sum() > 0


def test_missing_teachers_are_masked_and_quaternions_sign_invariant(splats):
    torch.manual_seed(1)
    model = Student(8,32)
    target = {k: torch.from_numpy(v).float()[None] for k,v in splat_targets(splats,8).items() if k != "grid"}
    prediction = model(torch.zeros(1,3,128,128))
    loss1, parts1 = distillation_loss(prediction, target, 8)
    target["geometry"][:] = 1
    target["quat"] *= -1
    loss2, parts2 = distillation_loss(prediction, target, 8)
    torch.testing.assert_close(loss1, loss2)
    assert parts1["geometry"] == 0
    assert parts2["geometry"] == 0
    target["geometry_valid"][:] = 1
    loss3, parts3 = distillation_loss(prediction, target, 8)
    assert parts3["geometry"] > 0
    assert loss3 > loss2


def test_query_chunk_attention_matches_full_attention():
    from torch.nn import functional as F
    torch.manual_seed(7)
    layers = SimpleNamespace(RePo3DRotaryEmbedding=type("RoPE", (), {}))
    chunked = patch_attention(layers, 3)
    q = torch.randn(2,11,2,8)
    k = torch.randn(2,17,2,8)
    v = torch.randn(2,17,2,8)
    expected = F.scaled_dot_product_attention(q.permute(0,2,1,3), k.permute(0,2,1,3), v.permute(0,2,1,3)).permute(0,2,1,3)
    torch.testing.assert_close(chunked(q,k,v), expected, atol=1e-6, rtol=1e-5)
    qkv = torch.randn(1,7,3,2,8)
    q,k,v = qkv.unbind(2)
    torch.testing.assert_close(chunked(qkv), chunked(q,k,v))
    assert layers.RePo3DRotaryEmbedding.real_valued


def test_split_keeps_object_views_together():
    rows = [{"id": str(i), "image_sha256": str(i), "metadata": {"object_id": str(i//2)}} for i in range(8)]
    train, val = split_rows(rows, .25, 42)
    assert not {r["metadata"]["object_id"] for r in train} & {r["metadata"]["object_id"] for r in val}
    assert len(train) == 6 and len(val) == 2


def test_actual_training_reduces_distillation_loss(splats):
    torch.set_num_threads(2)
    torch.manual_seed(42)
    model = Student(8,32)
    target = {k: torch.from_numpy(v).float()[None] for k,v in splat_targets(splats,8).items() if k != "grid"}
    image = torch.rand(1,3,128,128)
    optimizer = torch.optim.Adam(model.parameters(), lr=.003)
    initial = float(distillation_loss(model(image), target, 8)[0].detach())
    for _ in range(25):
        optimizer.zero_grad()
        loss, _ = distillation_loss(model(image), target, 8)
        loss.backward()
        optimizer.step()
    final = float(distillation_loss(model(image), target, 8)[0].detach())
    assert final < initial * .65, (initial, final)