import numpy as np
import torch

from mapgs.config import MapConfig
from mapgs.hdmap import (
    grid_field_from_points, HDMap, sample_anchors, rasterize_map_depth, project_polylines,
)
from mapgs.hdmap.anchors import resample_polyline, weighted_fps
from .conftest import requires_cuda


def _toy_map():
    pts = np.random.RandomState(0).uniform([-20, -5], [20, 50], size=(2000, 2))
    z = 0.02 * pts[:, 0]
    gp = np.concatenate([pts, z[:, None]], 1)
    gf = grid_field_from_points(gp, 0.5)
    lanes = [torch.tensor([[x, y, 0.02 * x] for y in np.arange(0, 45, 1.0)], dtype=torch.float32)
             for x in [-3.5, 0, 3.5]]
    bnds = [torch.tensor([[x, y, 0.02 * x] for y in np.arange(0, 45, 2.0)], dtype=torch.float32)
            for x in [-7, 7]]
    return HDMap(ground=gf, lanes=lanes, boundaries=bnds)


def test_ground_field_bilinear_and_grad():
    hd = _toy_map()
    xy = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
    z, valid = hd.height_at(xy)
    assert valid.all()
    assert torch.allclose(z, 0.02 * xy[:, 0].detach(), atol=0.05)
    z.sum().backward()
    assert xy.grad.abs().sum() > 0


def test_anchor_ratios_and_count():
    hd = _toy_map()
    cfg = MapConfig(n_anchors=1000)
    A = sample_anchors(hd, cfg, seed=0)
    assert len(A) == 1000
    counts = torch.bincount(A.types, minlength=3).float() / 1000
    assert abs(counts[0] - 0.6) < 0.05
    assert abs(counts[1] - 0.3) < 0.05
    assert abs(counts[2] - 0.1) < 0.05


def test_resample_polyline_spacing():
    pl = torch.tensor([[0., 0, 0], [0, 10, 0]])
    out = resample_polyline(pl, 1.0)
    assert out.shape[0] >= 9


def test_weighted_fps_prefers_high_weight():
    pts = torch.randn(200, 3)
    w = torch.ones(200)
    w[0] = 100.0
    sel = weighted_fps(pts, w, 10)
    assert 0 in sel.tolist()  # highest-weight point is picked first


@requires_cuda
def test_map_depth_coverage():
    from mapgs.geometry import look_at_pose
    hd = _toy_map().to("cuda")
    K = torch.tensor([[60., 0, 32], [0, 60, 24], [0, 0, 1]], device="cuda")[None]
    c2w = look_at_pose(torch.tensor([0., 0, 1.5]), torch.tensor([0., 20, 0.5]))[None].cuda()
    depth, mask = rasterize_map_depth(hd.ground, K, c2w, 48, 64)
    assert mask.float().mean() > 0.2  # lower half sees ground
    assert (depth[mask] > 0).all()
    uv = project_polylines(hd.lanes, K, c2w, 48, 64)
    assert uv[0].shape[0] > 0