import os

import torch

from mapgs.config import load_config
from mapgs.geometry import resize_with_intrinsics
from mapgs.hdmap.ground_field import GridGroundField
from .conftest import requires_cuda, tiny_overrides


def test_ground_field_from_raster_and_scaled():
    gf = GridGroundField.from_raster(torch.ones(10, 12) * 2.0, x0=-3.0, y0=-2.0, dx=0.5)
    z, valid = gf.height_at(torch.tensor([[0.0, 0.0]]))
    assert valid.all() and abs(float(z[0]) - 2.0) < 1e-4
    g2 = gf.scaled(0.5)
    assert abs(g2.x0 - (-1.5)) < 1e-6 and abs(g2.dx - 0.25) < 1e-6
    # scaling halves heights too (spatial)
    z2, _ = g2.height_at(torch.tensor([[0.0, 0.0]]))
    assert abs(float(z2[0]) - 1.0) < 1e-4


def test_resize_with_intrinsics():
    img = torch.rand(3, 100, 200)
    K = torch.tensor([[200.0, 0, 100], [0, 200, 50], [0, 0, 1]])
    out, Ks = resize_with_intrinsics(img, K, 50, 50)
    assert out.shape == (3, 50, 50)
    assert abs(float(Ks[0, 0]) - 50.0) < 1e-3   # fx * (50/200)
    assert abs(float(Ks[1, 1]) - 100.0) < 1e-3  # fy * (50/100)
    assert abs(float(Ks[0, 2]) - 25.0) < 1e-3   # cx * (50/200)


@requires_cuda
def test_unified_convert_load_and_mixed_roots(tmp_path):
    cfg = load_config(overrides=tiny_overrides(str(tmp_path / "src")) + ["data.name=unified"])
    from mapgs.data.convert import convert_synthetic
    from mapgs.data import UnifiedClipDataset, collate_samples

    ra = str(tmp_path / "ua"); rb = str(tmp_path / "ub")
    convert_synthetic(cfg, ra, "train", n_clips=2, device="cuda")
    convert_synthetic(cfg, rb, "train", n_clips=2, device="cuda")

    # on-disk layout
    clip = sorted(os.listdir(os.path.join(ra, "train")))[0]
    files = os.listdir(os.path.join(ra, "train", clip))
    assert "meta.pt" in files and "images" in files

    ds = UnifiedClipDataset(cfg, roots=[ra, rb], split="train", n_sup_views=3)
    assert len(ds) == 4  # mixed roots concatenated
    s = ds[0]
    assert s.ctx_images.shape[-2:] == (cfg.data.height, cfg.data.width)
    assert s.anchor_pos.shape[0] == cfg.model.tokens.n_map
    batch = collate_samples([ds[0], ds[1]])
    assert batch["ctx_images"].shape[0] == 2


@requires_cuda
def test_unified_train_step(tmp_path):
    cfg = load_config(overrides=tiny_overrides(str(tmp_path / "src")) + ["data.name=unified"])
    from mapgs.data.convert import convert_synthetic
    from mapgs.data import UnifiedClipDataset, collate_samples
    from mapgs.train import Trainer

    root = str(tmp_path / "u")
    convert_synthetic(cfg, root, "train", n_clips=2, device="cuda")
    ds = UnifiedClipDataset(cfg, roots=root, split="train", n_sup_views=3)
    tr = Trainer(cfg)
    log = tr.train_step(collate_samples([ds[0], ds[1]]))
    assert log["total"] == log["total"]  # finite