import os

import torch

from mmdet3d.models.data_preprocessors import voxelize as voxelize_mod
from mmdet3d.models.data_preprocessors.voxelize import voxelization


def run_mmcv(points, voxel_size, coors_range, max_points, max_voxels):
    voxels = points.new_zeros((max_voxels, max_points, points.size(1)))
    coors = points.new_zeros((max_voxels, 3), dtype=torch.int)
    counts = points.new_zeros((max_voxels, ), dtype=torch.int)
    voxel_num = torch.zeros((), dtype=torch.long)
    voxelize_mod.ext_module.hard_voxelize_forward(
        points,
        torch.tensor(voxel_size, dtype=torch.float),
        torch.tensor(coors_range, dtype=torch.float),
        voxels,
        coors,
        counts,
        voxel_num,
        max_points=max_points,
        max_voxels=max_voxels,
        NDim=3,
        deterministic=True)
    n = int(voxel_num.item())
    return voxels[:n].clone(), coors[:n].clone(), counts[:n].clone()


def run_wrapped(points, voxel_size, coors_range, max_points, max_voxels):
    old = os.environ.pop('MMDET3D_DISABLE_VOXELIZE_OPT', None)
    try:
        return tuple(
            x.clone()
            for x in voxelization(points, voxel_size, coors_range, max_points,
                                  max_voxels, True))
    finally:
        if old is not None:
            os.environ['MMDET3D_DISABLE_VOXELIZE_OPT'] = old


def check_case(name, points, voxel_size, coors_range, max_points, max_voxels):
    ref_voxels, ref_coors, ref_counts = run_mmcv(points, voxel_size,
                                                 coors_range, max_points,
                                                 max_voxels)
    opt_voxels, opt_coors, opt_counts = run_wrapped(points, voxel_size,
                                                   coors_range, max_points,
                                                   max_voxels)
    ok = (torch.equal(ref_coors, opt_coors) and torch.equal(
        ref_counts, opt_counts) and torch.equal(ref_voxels, opt_voxels))
    status = 'PASS' if ok else 'FAIL'
    print(f'{name}: {status} voxels={ref_coors.size(0)}')
    if not ok:
        print('  coors_equal:', torch.equal(ref_coors, opt_coors))
        print('  counts_equal:', torch.equal(ref_counts, opt_counts))
        print('  voxels_equal:', torch.equal(ref_voxels, opt_voxels))
        raise SystemExit(1)


def main():
    if not torch.cuda.is_available():
        raise SystemExit('CUDA/HIP device is required')
    device = torch.device('cuda')
    voxel_size = [0.25, 0.25, 8.0]
    coors_range = [-50.0, -50.0, -5.0, 50.0, 50.0, 3.0]
    gen = torch.Generator(device=device)
    gen.manual_seed(20260521)

    cases = []
    pts = torch.empty((30000, 5), device=device).uniform_(
        -1.0, 1.0, generator=gen)
    pts[:, 0] = torch.empty((30000, ), device=device).uniform_(
        -49.9, 49.9, generator=gen)
    pts[:, 1] = torch.empty((30000, ), device=device).uniform_(
        -49.9, 49.9, generator=gen)
    pts[:, 2] = torch.empty((30000, ), device=device).uniform_(
        -4.9, 2.9, generator=gen)
    cases.append(('normal_30k', pts, 64, 30000))

    pts = torch.empty((100000, 5), device=device).uniform_(
        -1.0, 1.0, generator=gen)
    pts[:, 0] = torch.empty((100000, ), device=device).uniform_(
        -55.0, 55.0, generator=gen)
    pts[:, 1] = torch.empty((100000, ), device=device).uniform_(
        -55.0, 55.0, generator=gen)
    pts[:, 2] = torch.empty((100000, ), device=device).uniform_(
        -8.0, 5.0, generator=gen)
    cases.append(('normal_100k', pts, 64, 30000))

    pts = torch.zeros((4096, 5), device=device)
    pts[:, 0] = torch.linspace(-50.0, 49.999, 4096, device=device)
    pts[:, 1] = torch.linspace(49.999, -50.0, 4096, device=device)
    pts[:, 2] = torch.linspace(-5.0, 2.999, 4096, device=device)
    cases.append(('near_boundary', pts, 64, 30000))

    pts = torch.zeros((8192, 5), device=device)
    pts[:, 0] = 0.01
    pts[:, 1] = 0.01
    pts[:, 2] = 0.01
    pts[:, 3:] = torch.randn((8192, 2), device=device, generator=gen)
    cases.append(('dense_overflow', pts, 64, 30000))

    pts = torch.empty((50000, 5), device=device).uniform_(
        -1.0, 1.0, generator=gen)
    pts[:, 0] = torch.empty((50000, ), device=device).uniform_(
        -49.9, 49.9, generator=gen)
    pts[:, 1] = torch.empty((50000, ), device=device).uniform_(
        -49.9, 49.9, generator=gen)
    pts[:, 2] = torch.empty((50000, ), device=device).uniform_(
        -4.9, 2.9, generator=gen)
    cases.append(('max_voxels_clip', pts, 64, 1024))

    for name, points, max_points, max_voxels in cases:
        check_case(name, points.contiguous(), voxel_size, coors_range,
                   max_points, max_voxels)
    print('all cases passed')


if __name__ == '__main__':
    main()
