import torch

from mapgs.model import MapGS
from mapgs.ttt import token_tuning, map_guided_densify
from mapgs.ttt.densify import anchors_along_trajectory
from .conftest import requires_cuda


@requires_cuda
def test_token_tuning_runs(cfg, dataset):
    model = MapGS(cfg).cuda()
    g = token_tuning(model, dataset[0], cfg, steps=3)
    M = (cfg.model.tokens.n_map + cfg.model.tokens.n_free) * cfg.model.tokens.gaussians_per_token
    assert g.means.shape == (1, M, 3)
    # network weights must be unchanged by TT (only tokens are tuned)
    before = {n: p.clone() for n, p in model.named_parameters()}
    token_tuning(model, dataset[0], cfg, steps=2)
    for n, p in model.named_parameters():
        assert torch.allclose(before[n], p)


@requires_cuda
def test_densify_increases_gaussians(cfg, dataset):
    model = MapGS(cfg).cuda()
    s = dataset[0]
    traj = torch.stack([torch.zeros(15, device="cuda"), torch.linspace(2, 25, 15, device="cuda")], -1)
    npos, ntyp, nnrm = anchors_along_trajectory(s.ground.to("cuda"), traj, spacing=1.0)
    n_new = min(64, npos.shape[0])
    g = map_guided_densify(model, s, npos[:n_new], ntyp[:n_new], nnrm[:n_new], cfg)
    base = (cfg.model.tokens.n_map + cfg.model.tokens.n_free) * cfg.model.tokens.gaussians_per_token
    assert g.means.shape[1] == base + n_new * cfg.model.tokens.gaussians_per_token