import collections
import io
import tempfile
import zipfile

import random

has_debugpy = True
try:
    import debugpy
except ImportError:
    has_debugpy = False
import pycolmap
import torch
import numpy as np


def parse_colmap_reconstruction(colmap_data: bytes) -> pycolmap.Reconstruction:
    """Parses a COLMAP reconstruction from a zip file.

    Args:
        colmap_data (bytes): The COLMAP reconstruction data as a zip file.
    Returns:
        pycolmap.Reconstruction: The parsed COLMAP reconstruction.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        with zipfile.ZipFile(io.BytesIO(colmap_data), "r") as zf:
            zf.extractall(tmpdir)
        return pycolmap.Reconstruction(tmpdir)



def map_tensor(input: any, func: callable) -> any:
    if isinstance(input, str):
        return input
    elif isinstance(input, collections.abc.Mapping):
        return {k: map_tensor(sample, func) for k, sample in input.items()}
    elif isinstance(input, collections.abc.Sequence):
        return [map_tensor(sample, func) for sample in input]
    else:
        return func(input)


def batch_to_device(batch: any, device: str, non_blocking: bool = True):
    return map_tensor(batch, lambda x: x.to(device=device, non_blocking=non_blocking))



def start_debug():
    if not has_debugpy:
        raise ImportError("debugpy library is required for debugging.")
    debugpy.listen(5678)
    print("Wait for debugger!")
    debugpy.wait_for_client()
    print("Attached!")


def set_random_seed(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    if hasattr(torch.backends, "cudnn"):
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True