# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import enum
import math
import os
from typing import Union, Optional, Tuple


def _next_power_of_2(n: int) -> int:
    """Round up to the next power of 2. If n is already a power of 2, return n."""
    if n <= 0:
        return 1
    n -= 1
    n |= n >> 1
    n |= n >> 2
    n |= n >> 4
    n |= n >> 8
    n |= n >> 16
    n |= n >> 32
    return n + 1


import numpy as np
import torch
from torch import Tensor

import warpconvnet._C as _C

_EXPAND_STATUS_SUCCESS = 0
_EXPAND_STATUS_VECTOR_OVERFLOW = 1
_EXPAND_STATUS_TABLE_FULL = 2


class HashMethod(enum.Enum):
    """Hash method enumeration for the vector hash table.

    Attributes:
        FNV1A: FNV-1a hash algorithm
        CITY: CityHash algorithm
        MURMUR: MurmurHash algorithm
    """

    FNV1A = 0
    CITY = 1
    MURMUR = 2

    def kernel_suffix(self) -> str:
        """Return the suffix used for templated kernel names."""
        return self.name.lower()  # fnv1a, city, murmur


class TorchHashTable:
    """
    A hash table implementation using CUDA kernels via the _C extension for
    vector key storage and lookup, operating on pytorch Tensors.
    """

    _capacity: int
    _hash_method_enum: HashMethod
    _table_kvs: Tensor = None  # Shape: (capacity, 2), dtype=torch.int32
    _vector_keys: Tensor = None  # Shape: (num_keys, key_dim), dtype=torch.int32
    _key_dim: int = -1
    _device: torch.device = None
    _num_entries: int = 0
    _vector_capacity: int = 0

    def __init__(
        self,
        capacity: int,
        hash_method: HashMethod = HashMethod.CITY,
        device: Union[str, torch.device] = "cuda",
    ):
        """Initialize the hash table using PyTorch tensors.

        Args:
            capacity: Maximum number of entries the table can store (number of slots)
            hash_method: HashMethod enum value (default: CITY)
            device: The torch device to allocate tensors on (e.g., 'cuda', 'cuda:0').
        """
        if not isinstance(hash_method, HashMethod):
            raise TypeError(
                f"hash_method must be a HashMethod enum member, got {type(hash_method)}"
            )
        assert capacity > 0

        # CUDA kernels use bitwise AND (& (capacity-1)) for hash slot
        # computation, so capacity must always be a power of 2.
        self._capacity = _next_power_of_2(capacity)
        self._hash_method_enum = hash_method
        self._device = torch.device(device)
        self._num_entries = 0
        self._vector_capacity = 0

    @property
    def capacity(self) -> int:
        return self._capacity

    @property
    def hash_method(self) -> HashMethod:
        return self._hash_method_enum

    @property
    def device(self) -> torch.device:
        if self._table_kvs is not None:
            return self._table_kvs.device
        return self._device  # Return the device specified during init

    @property
    def key_dim(self) -> int:
        return self._key_dim

    @property
    def num_entries(self) -> int:
        return self._num_entries

    def insert(
        self,
        vec_keys: Tensor,
        threads_per_block: int = 256,
        storage_capacity: Optional[int] = None,
    ):
        """Insert vector keys (PyTorch Tensor) into the hash table.

        Args:
            vec_keys: PyTorch tensor of int32 vector keys, shape (num_keys, key_dim), on CUDA device.

        Raises:
            AssertionError: If capacity is invalid or number of keys exceeds capacity/2.
            TypeError: If input is not a Torch Tensor or has wrong dtype.
            ValueError: If input is not 2D or not on a CUDA device.
        """
        if not isinstance(vec_keys, torch.Tensor):
            raise TypeError(
                f"Input vec_keys must be a PyTorch Tensor, got {type(vec_keys)}"
            )
        if not vec_keys.is_cuda:
            raise ValueError(
                f"Input vec_keys must be on a CUDA device, got {vec_keys.device}"
            )
        if vec_keys.ndim != 2:
            raise ValueError(
                f"Input vec_keys must be 2D, got {vec_keys.ndim} dimensions"
            )

        # Ensure correct device and dtype
        if vec_keys.device != self.device:
            vec_keys = vec_keys.to(self.device)

        if vec_keys.dtype != torch.int32:
            if vec_keys.dtype == torch.int64:
                vec_keys = vec_keys.to(dtype=torch.int32)
            else:
                raise TypeError(
                    f"Input vec_keys must have dtype torch.int32, got {vec_keys.dtype}"
                )

        vec_keys = vec_keys.contiguous()

        num_keys, key_dim = vec_keys.shape
        self._key_dim = key_dim

        if storage_capacity is None:
            storage_capacity = num_keys
        storage_capacity = max(storage_capacity, num_keys)

        assert self._capacity > 0
        assert (
            num_keys <= self._capacity / 2
        ), f"Number of keys {num_keys} exceeds recommended capacity/2 ({self._capacity / 2}) for table size {self._capacity}"

        # Allocate table on the target device
        self._table_kvs = torch.empty(
            (self._capacity, 2), dtype=torch.int32, device=self.device
        )
        # Allocate storage for vector keys to allow future growth
        self._vector_keys = torch.empty(
            (storage_capacity, key_dim), dtype=torch.int32, device=self.device
        )
        self._vector_keys[:num_keys] = vec_keys
        self._vector_capacity = storage_capacity
        self._num_entries = num_keys

        # --- Launch Prepare Kernel ---
        _C.coords.hashmap_prepare(self._table_kvs, self._capacity)
        # No sync needed: CUDA stream ordering guarantees prepare completes
        # before insert starts on the same stream.

        # --- Launch Insert Kernel ---
        _C.coords.hashmap_insert(
            self._table_kvs,
            self._vector_keys,
            num_keys,
            self._key_dim,
            self._capacity,
            self._hash_method_enum.value,
        )
        # No sync needed: downstream kernels on the same stream will see
        # the inserted data. Only sync if CPU needs to read results.

    @classmethod
    def from_keys(
        cls,
        vec_keys: Union[Tensor, np.ndarray],
        hash_method: HashMethod = HashMethod.CITY,
        device: Union[str, torch.device] = "cuda",
        capacity: Optional[int] = None,
        vector_capacity: Optional[int] = None,
    ):
        """Create a hash table from a set of vector keys.

        Args:
            vec_keys: Vector keys (PyTorch Tensor or NumPy array). If NumPy, moved to `device`.
            hash_method: HashMethod enum value to use (default: CITY)
            device: The torch device to use.

        Returns:
            TorchHashTable: New hash table instance containing the keys.
        """
        target_device = torch.device(device)

        if not isinstance(vec_keys, torch.Tensor):
            # If NumPy or other array-like, convert to Tensor on the target device
            try:
                vec_keys = torch.as_tensor(vec_keys, device=target_device)
            except Exception as e:
                raise TypeError(
                    f"Could not convert input vec_keys to a Torch Tensor on device {target_device}: {e}"
                )

        # Ensure correct device and dtype
        if vec_keys.device != target_device:
            vec_keys = vec_keys.to(target_device)

        if vec_keys.dtype != torch.int32:
            if vec_keys.dtype == torch.int64:
                vec_keys = vec_keys.to(dtype=torch.int32)
            else:
                raise TypeError(
                    f"Input vec_keys for from_keys must have dtype torch.int32 or compatible, got {vec_keys.dtype}"
                )

        if vec_keys.ndim != 2:
            raise ValueError(
                f"Input vec_keys for from_keys must be 2D, got {vec_keys.ndim} dimensions"
            )

        vec_keys = vec_keys.contiguous()

        num_keys = len(vec_keys)
        # Constructor enforces power-of-2 for bitwise AND hash slot computation
        chosen_capacity = (
            capacity if capacity is not None else max(16, int(num_keys * 2))
        )
        storage_capacity = vector_capacity if vector_capacity is not None else num_keys
        # Pass the hash_method and device to the constructor
        obj = cls(
            capacity=chosen_capacity, hash_method=hash_method, device=target_device
        )
        obj.insert(vec_keys, storage_capacity=storage_capacity)
        return obj

    def search(
        self,
        search_keys: Union[Tensor, np.ndarray],
        threads_per_block: int = 256,
        cooperative: bool = True,
    ) -> Tensor:
        """Search for keys (PyTorch Tensor) in the hash table.

        Args:
            search_keys: Keys to search for (Torch Tensor or NumPy array). If NumPy, moved to table's device.
                         Shape (num_search, key_dim).

        Returns:
            torch.Tensor: Array of indices (int32) where keys were found in the original
                          `vector_keys` tensor. -1 if not found. On the same device as the table.
        Raises:
            RuntimeError: If insert() has not been called yet.
            ValueError: If search_keys dimensions don't match inserted keys.
            TypeError: If search_keys cannot be converted to a Tensor or have wrong dtype.
        """
        if self._table_kvs is None or self._vector_keys is None:
            raise RuntimeError(
                "Hash table has not been populated. Call insert() or from_keys() first."
            )

        table_device = self.device

        if not isinstance(search_keys, torch.Tensor):
            try:
                # Convert NumPy or other to Tensor on the hash table's device
                search_keys = torch.as_tensor(search_keys, device=table_device)
            except Exception as e:
                raise TypeError(
                    f"Could not convert input search_keys to a Torch Tensor on device {table_device}: {e}"
                )

        # Ensure correct device and dtype
        if search_keys.device != table_device:
            search_keys = search_keys.to(table_device)

        if search_keys.ndim != 2:
            raise ValueError(
                f"Input search_keys must be 2D, got {search_keys.ndim} dimensions"
            )
        if search_keys.shape[1] != self._key_dim:
            raise ValueError(
                f"Search keys dimension ({search_keys.shape[1]}) must match "
                f"inserted keys dimension ({self._key_dim})"
            )

        if search_keys.dtype != torch.int32:
            if search_keys.dtype == torch.int64:
                search_keys = search_keys.to(dtype=torch.int32)
            else:
                raise TypeError(
                    f"Input search_keys must have dtype torch.int32 or compatible, got {search_keys.dtype}"
                )

        num_search_keys = len(search_keys)
        # Allocate results tensor on the correct device
        results = torch.empty(num_search_keys, dtype=torch.int32, device=table_device)

        # --- Launch Search Kernel ---
        if cooperative and hasattr(_C.coords, "hashmap_warp_search"):
            _C.coords.hashmap_warp_search(
                self._table_kvs,
                self._vector_keys,
                search_keys,
                results,
                num_search_keys,
                self._key_dim,
                self._capacity,
                self._hash_method_enum.value,
            )
        else:
            _C.coords.hashmap_search(
                self._table_kvs,
                self._vector_keys,
                search_keys,
                results,
                num_search_keys,
                self._key_dim,
                self._capacity,
                self._hash_method_enum.value,
            )
        # No sync needed: downstream PyTorch ops on the same stream will
        # see the results. Only sync if CPU needs to read the tensor.

        return results

    def expand_with_offsets(
        self,
        base_coords: Tensor,
        offsets: Tensor,
        threads_per_block: int = 256,
    ):
        """Expand the hash table by applying offsets to base coordinates on device."""
        if self._table_kvs is None or self._vector_keys is None:
            raise RuntimeError(
                "Hash table has not been populated. Call insert() or from_keys() first."
            )

        target_device = self.device
        if not isinstance(base_coords, torch.Tensor):
            base_coords = torch.as_tensor(base_coords, device=target_device)
        if not isinstance(offsets, torch.Tensor):
            offsets = torch.as_tensor(offsets, device=target_device)

        if base_coords.device != target_device:
            base_coords = base_coords.to(target_device)
        if offsets.device != target_device:
            offsets = offsets.to(target_device)

        if base_coords.dtype != torch.int32:
            base_coords = base_coords.to(dtype=torch.int32)
        if offsets.dtype != torch.int32:
            offsets = offsets.to(dtype=torch.int32)

        base_coords = base_coords.contiguous()
        offsets = offsets.contiguous()

        if base_coords.ndim != 2 or offsets.ndim != 2:
            raise ValueError("base_coords and offsets must be 2D tensors.")
        if base_coords.shape[1] != self._key_dim or offsets.shape[1] != self._key_dim:
            raise ValueError("Key dimension mismatch between hash table and inputs.")

        num_base = base_coords.shape[0]
        num_offsets = offsets.shape[0]
        if num_base == 0 or num_offsets == 0:
            return

        total_candidates = num_base * num_offsets
        self._ensure_vector_storage(self._num_entries + total_candidates)

        num_entries_tensor = torch.tensor(
            [self._num_entries], dtype=torch.int32, device=target_device
        )
        status_tensor = torch.zeros(1, dtype=torch.int32, device=target_device)

        _C.coords.hashmap_expand(
            self._table_kvs,
            self._vector_keys,
            base_coords,
            offsets,
            num_base,
            num_offsets,
            self._key_dim,
            self._capacity,
            self._vector_capacity,
            num_entries_tensor,
            status_tensor,
            self._hash_method_enum.value,
        )
        # No explicit sync needed: .item() below implicitly synchronizes the
        # current stream before copying the value to CPU.

        status = int(status_tensor.item())
        if status == _EXPAND_STATUS_VECTOR_OVERFLOW:
            raise RuntimeError(
                "TorchHashTable.expand_with_offsets exceeded vector storage capacity."
            )
        if status == _EXPAND_STATUS_TABLE_FULL:
            raise RuntimeError(
                "TorchHashTable.expand_with_offsets ran out of hash table slots."
            )

        self._num_entries = int(num_entries_tensor.item())

    @property
    def unique_index(self) -> Tensor:
        """Get sorted unique indices from the hash table.

        Returns:
            torch.Tensor: Sorted tensor of unique indices (int32) corresponding
                          to the originally inserted keys. On the same device as the table.
        Raises:
            RuntimeError: If insert() has not been called yet.
        """
        if self._vector_keys is None:
            raise RuntimeError(
                "Hash table has not been populated. Call insert() or from_keys() first."
            )

        vec_keys = self.vector_keys
        indices = self.search(vec_keys)
        valid_indices = indices[indices != -1]
        # torch.unique returns sorted unique values
        unique_indices = torch.unique(valid_indices)
        return unique_indices

    @property
    def vector_keys(self) -> Tensor:
        """Return the 2D vector keys tensor used to build the hash table."""
        if self._vector_keys is None:
            raise RuntimeError(
                "Hash table has not been populated. Call insert() or from_keys() first."
            )
        return self._vector_keys[: self._num_entries]

    @property
    def unique_vector_keys(self) -> Tensor:
        """Return the unique 2D vector keys present in the hash table, sorted by index.

        Returns:
            torch.Tensor: Tensor shape (num_unique_keys, key_dim), dtype int32.
                          On the same device as the table.
        Raises:
            RuntimeError: If insert() has not been called yet.
        """
        unique_idx = self.unique_index
        if unique_idx.numel() == 0:  # Use numel() for checking empty tensors
            # Return an empty tensor with the correct shape, dtype, and device
            return torch.empty((0, self.key_dim), dtype=torch.int32, device=self.device)
        vec_keys = self.vector_keys
        return vec_keys[unique_idx]

    def _ensure_vector_storage(self, required_capacity: int):
        if self._vector_keys is None:
            raise RuntimeError(
                "Hash table has not been populated. Call insert() or from_keys() first."
            )
        if required_capacity <= self._vector_capacity:
            return
        growth = max(required_capacity // 2, 1)
        new_capacity = max(required_capacity, self._vector_capacity + growth)
        new_storage = torch.empty(
            (new_capacity, self._key_dim), dtype=torch.int32, device=self.device
        )
        new_storage[: self._num_entries] = self.vector_keys
        self._vector_keys = new_storage
        self._vector_capacity = new_capacity

    def to_dict(self) -> dict:
        """Serializes the data arrays and metadata (transfers tensors to CPU as NumPy arrays)."""
        if self._table_kvs is None or self._vector_keys is None:
            table_kvs_np = None
            vec_keys_np = None
        else:
            # Ensure tensors are on CPU before converting to numpy
            table_kvs_np = self._table_kvs.cpu().numpy()
            vec_keys_np = self.vector_keys.cpu().numpy()

        return {
            "table_kvs": table_kvs_np,
            "vec_keys": vec_keys_np,
            "hash_method_value": self._hash_method_enum.value,
            "capacity": self._capacity,
            "key_dim": self._key_dim,
            "device": str(self.device),  # Store device as string
            "num_entries": self._num_entries,
            "vector_capacity": self._vector_capacity,
        }

    def from_dict(self, data: dict):
        """
        Loads data from a dict and re-initializes the hash table state.
        Assumes data arrays are numpy arrays. Tensors are created on the specified device.
        """
        required_keys = {
            "capacity",
            "hash_method_value",
            "key_dim",
            "table_kvs",
            "vec_keys",
            "device",
            "num_entries",
            "vector_capacity",
        }
        if not required_keys.issubset(data.keys()):
            raise ValueError(
                f"Data dictionary missing required keys. Need: {required_keys}"
            )

        capacity = data["capacity"]
        hash_method_value = data["hash_method_value"]
        key_dim = data["key_dim"]
        table_kvs_np = data["table_kvs"]
        vec_keys_np = data["vec_keys"]
        device_str = data["device"]
        target_device = torch.device(device_str)
        num_entries = int(data["num_entries"])
        vector_capacity = int(data["vector_capacity"])

        # Re-initialize the object with the correct capacity, hash method, and device
        self.__init__(
            capacity=capacity,
            hash_method=HashMethod(hash_method_value),
            device=target_device,
        )

        self._key_dim = key_dim
        self._vector_capacity = vector_capacity
        self._num_entries = num_entries

        # Load data if present, creating tensors on the target device
        if table_kvs_np is not None:
            self._table_kvs = torch.as_tensor(
                table_kvs_np, dtype=torch.int32, device=target_device
            )
            assert self._table_kvs.shape == (self._capacity, 2)
            assert self._table_kvs.dtype == torch.int32
        else:
            self._table_kvs = None

        if vec_keys_np is not None:
            # Allocate full storage and copy populated rows.
            self._vector_keys = torch.empty(
                (self._vector_capacity, self._key_dim),
                dtype=torch.int32,
                device=target_device,
            )
            vec_keys_t = torch.as_tensor(
                vec_keys_np, dtype=torch.int32, device=target_device
            )
            assert vec_keys_t.ndim == 2
            assert vec_keys_t.shape[1] == self._key_dim
            self._vector_keys[: self._num_entries] = vec_keys_t
        else:
            self._vector_keys = None

        # State is loaded, kernels are implicitly ready from __init__.
        return self
