"""Real Metal tests; teacher uses random tiny modules, never pretrained weights."""
import pytest
import torch

from image3d_student.model import Student
from image3d_student.teacher_idol import patch_attention

pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="Metal GPU not visible")


def test_default_student_backward_on_metal():
    model = Student().to("mps")
    result = model(torch.rand(1,3,128,128,device="mps"))
    loss = sum(t.square().mean() for t in result.values())
    loss.backward()
    assert torch.isfinite(loss)
    assert result["xyz"].shape == (1,4096,3)
    assert all(p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters())


@torch.no_grad()
def test_tiny_real_idol_architecture_on_metal():
    pytest.importorskip("idol")
    from tensordict import TensorDict
    from idol.models import LatentSeqMMFlowModel, OctreeProbabilityFixedlenDecoder, GaussianFixedlenDecoder, layers
    from idol.sampler import sample_flow
    patch_attention(layers, query_chunk=4)
    model = LatentSeqMMFlowModel(q_token_length=8,in_channels=4,cam_channels=6,
        model_channels=32,cond_channels=12,cond2_channels=8,num_blocks=1,
        num_refiner_blocks=2,num_heads=2,num_head_channels=16).eval().to("mps")
    noise = TensorDict({"latent": torch.randn(1,8,4,device="mps"),
                        "camera": torch.randn(1,1,6,device="mps")}, batch_size=1, device="mps")
    cond = TensorDict({"feature1": torch.randn(1,6,12,device="mps"),
                       "feature2": torch.randn(1,6,8,device="mps")}, batch_size=1, device="mps")
    result = sample_flow(model,noise,cond,cond.apply(torch.zeros_like),steps=2,verbose=False,batch_cfg=False)
    latent = result["latent"]
    assert torch.isfinite(latent).all()
    octree = OctreeProbabilityFixedlenDecoder(32,4,1,2).eval().to("mps")
    points = octree.sample(octree,latent,num_points=8,level=3)
    assert points.shape == (1,8,3)
    config = {"num_gaussians": 4, "lr": {k: 1. for k in ["_xyz","_features_dc","_scaling","_rotation","_opacity"]},
              "perturbe_size": .1, "scaling_bias": .01, "opacity_bias": .8, "3d_filter_kernel_size": .001}
    decoder = GaussianFixedlenDecoder(3,32,4,1,2,config).eval().to("mps")
    gaussian = decoder.to_representation(points,decoder(points,latent))[0]
    assert gaussian.xyz.shape == (32,3)
    for value in (gaussian.xyz,gaussian.features_dc,gaussian.scaling,gaussian.rotation,gaussian.opacity):
        assert torch.isfinite(value).all()