import torch
import torch.nn as nn
import numpy as np

from .kernels import GetAddressOfIDCUDA
from .kernels import GetDivOfAddressCUDA
from .kernels import GetIOACUDA
from .kernels import GetWriteAddressCUDA

class IVFPQBase(nn.Module):
  def __init__(
    self,
    code_size,
    n_cq_clusters,
    blocksize,
    verbose,
    distance,
    device
    ):
    """
    this is a base class for IVFPQ variants
    for more details, you can search 'Product Quantization for Nearest Neighbor Search'

    Parameters:
      d_vector: int
        dimentionality of vectors to be quantized.

      code_size: int, default : 8
        byte size of quantized code, must to be a multiple of 4
        maximum possible code_size depends on GPU architecture
          GPU Architecture:   fp32  fp16 
          Ampere (GA100)      163   326  
          Turing (TU102 etc.) 64    128  
          Volta (GV100)       96    192  
          Pascal and before   48    96   
      
      n_cq_clusters: int, default : 128
        number coarse quantizer clusters
        recommended value is between 4*sqrt(n_data) ~ 16*sqrt(n_data)

      blocksize: int, default : 64
        number of vectors that can be asigned to each cluster of coarse_quantizer initially
        can be expanded using .expand methods, .add method will automatically cal .expand if necessary
        recommended value is (n_data_to_be_stored / n_cq_clusters)

      verbose: int, default : 0
        verbosity

      distance: str, default : 'euclidean'
        type of distance metric
        can be one of: ['euclidean', 'cosine', 'inner']

      device: str, default : 'cuda:0'
        which device to use
        can be one of: ['cpu', 'cuda', 'cuda:X'] *X represents cuda device index
      
    """
    super(IVFPQBase, self).__init__()
    n_cs = 4
    self.n_cs = n_cs
    self.code_size = code_size
    assert code_size % self.n_cs == 0, f"code_size needs to be a multiple of {n_cs}"
    self.n_cq_clusters = n_cq_clusters
    self.blocksize = blocksize
    self.verbose = verbose
    if distance.lower() in ["euclidean", "l2"]:
      distance = "euclidean"
    elif distance.lower() in ["dot", "inner"]:
      distance = "inner"
    elif distance.lower() in ["cosine"]:
      distance = "cosine"
    self.distance = distance
      
    self.register_buffer("_is_trained", torch.tensor(False))
    self.device = device

    storage = torch.zeros(
      code_size // n_cs,
      blocksize * n_cq_clusters,
      n_cs,
      device=self.device,
      dtype=torch.uint8,
    )
    address2id = torch.ones(
      blocksize * n_cq_clusters,
      device=self.device,
      dtype=torch.long,
    ) * -1
    is_empty = torch.ones(
      blocksize * n_cq_clusters,
      device=self.device,
      dtype=torch.uint8
    )
    div_start = torch.arange(
      n_cq_clusters,
      device=self.device
    ) * blocksize
    div_capacity = torch.ones(
      n_cq_clusters,
      device=self.device,
      dtype=torch.long,
    ) * blocksize
    div_size = torch.zeros(
      n_cq_clusters,
      device=self.device,
      dtype=torch.long
    )
    self.register_buffer("storage", storage)
    self.register_buffer("div_start", div_start)
    self.register_buffer("div_capacity",div_capacity)
    self.register_buffer("div_size", div_size)
    self.register_buffer("address2id", address2id)
    self.register_buffer("is_empty", is_empty)

    self._get_address_of_id_cuda = GetAddressOfIDCUDA(
      tpb=256,
    )
    self._get_div_of_address_cuda = GetDivOfAddressCUDA(
      ta=1,
      tpb=256,
    )
    self._get_ioa_cuda = GetIOACUDA(
      tpb=256,
    )
    self._get_write_address_cuda = GetWriteAddressCUDA(
      tpb=256,
    )

  def __getitem__(self, ids):
    """
      ids: int, or an iterable of integers
    """
    if type(ids) is not torch.Tensor:
      ids = torch.LongTensor(ids)
    if len(ids.shape) == 0:
      ids = ids[None]
    elif len(ids.shape) > 1:
      raise ValueError("only 1D array/tensors are supported")
    return self.get_data_of_id(ids)

  def __repr__(self):
    def get_bytesize_str(bytesize):
      if bytesize < 1024*1:
        bytesize = f"{bytesize} B"
      elif bytesize < 1024**2:
        bytesize = f"{round(bytesize / 1024, 2)} KB"
      elif bytesize < 1024**3:
        bytesize = f"{round(bytesize / 1024**2, 2)} MB"
      elif bytesize < 1024**4:
        bytesize = f"{round(bytesize / 1024**3, 2)} GB"
      elif bytesize < 1024**5:
        bytesize = f"{round(bytesize / 1024**4, 2)} TB"
      return bytesize

    bytesize = self.bytesize
    cpu_bytesize = get_bytesize_str(bytesize[0])
    gpu_bytesize = get_bytesize_str(bytesize[1])

    txt =  f"""
    {type(self).__name__}(
      number of stored items = {self.tot_size},
      total capacity of storage = {self.tot_capacity},
      byte size on CPU = {cpu_bytesize},
      byte size on GPU = {gpu_bytesize},
      code size = {self.code_size} B,
      is trained = {self._is_trained.item()}
    )
    """
    txt = '\n'.join([i[4:] for i in txt.split('\n')])
    return txt

  def __len__(self):
    return self.tot_size

  @staticmethod
  def get_cc():
    if torch.cuda.is_available():
      device_id = torch.cuda.current_device()
      gpu_properties = torch.cuda.get_device_properties(device_id)
      result = (gpu_properties.major, gpu_properties.minor)
    else:
      result = (10, 0)
    return result
  
  @staticmethod
  def normalize(a, dim=-1):
    """
      normalize given tensor along given dimention
    """
    a_norm = a.norm(dim=dim, keepdim=True) + 1e-9
    return a / a_norm

  @property
  def tot_size(self):
    """
      total number of items stored
    """
    return self.div_size.sum().item()

  @property
  def tot_capacity(self):
    """
      total capacity of storage
    """
    return self.div_capacity.sum().item()

  @property
  def bytesize(self):
    """
      the bytesize of cpu and gpu memory that are being used by module
      returns: (cpu_bytesize, gpu_bytesize)
    """
    cpu_bytesize = 0
    gpu_bytesize = 0
    for p in self.state_dict().values():
      if type(p) is torch.Tensor:
        if p.device.type == "cpu":
          cpu_bytesize += p.numel() * p.element_size()
        elif p.device.type == "cuda":
          gpu_bytesize += p.numel() * p.element_size()
    return (cpu_bytesize, gpu_bytesize)

  def _get_address_of_id_old(self, ids):
    batch_size = ids.shape[0]
    address = torch.zeros(batch_size, device=self.device, dtype=torch.long)
    for i in range(batch_size):
      id = ids[i]
      adr = torch.nonzero(self.address2id == id)
      if adr.shape[0] > 0:
        address[i] = adr[0, 0]
      else:
        address[i] = -1
    return address

  def get_address_of_id(self, ids):
    """
      ids: torch.Tensor, shape : [batch_size], dtype : int64
      if there are multiple matching ids, only the address of the first one will be returned
      if there are no matching ids, the address is -1
    """
    if self.address2id.device.type == "cuda":
      return self._get_address_of_id_cuda(self.address2id, ids)
    elif self.address2id.device.type == "cpu":
      return self._get_address_of_id_old(ids)

  def get_id_of_address(self, address):
    """
      address: torch.Tensor, shape : [batch_size], dtype : int64
    """
    return self.address2id[address]

  def _get_div_of_address_old(self, address):
    assert address.max() < self.tot_capacity
    batch_size = address.shape[0]
    div_start = self.div_start
    div_end = div_start + self.div_capacity
    # for adr in address:
    #   mask1 = div_start <= adr
    #   mask2 = div_ends > adr
    #   mask = mask1 & mask2
    #   div = torch.nonzero(mask)
    mask1 = div_start[None, ] <= address[:, None]
    mask2 = div_end[None, ] > address[:, None] # [batch_size, n_cq_clusters]
    mask = mask1 & mask2
    not_found = mask.sum(dim=1) == 0
    mask[not_found, 0] = True
    divs = torch.nonzero(mask)
    divs[not_found, 1] = -1
    return divs[:, 1]

  def get_div_of_address(self, address):
    """
      address: torch.Tensor, shape : [n_address], dtype : int64
      if address is not in range [0, tot_capacity), its div will be -1
    """
    if self.div_start.device.type == "cuda":
      div_start = self.div_start
      div_end = div_start + self.div_capacity
      return self._get_div_of_address_cuda(address, div_start, div_end)
    elif self.div_start.device.type == "cpu":
      return self._get_div_of_address_old(address)

  def get_data_of_address(self, address):
    """
      address: torch.Tensor, shape : [n_address], dtype : int64
      returns: torch.Tensor, shape : [code_size, n_address], dtype : uint8
      address must be in range [0, tot_capacity)
    """
    n_address = address.shape[0]
    assert address.dtype == torch.long
    mask = address < 0
    data = self.storage[:, address] #[code_size//n_cs, n_address, n_cs]
    data = data.transpose(1,2).reshape(self.code_size, n_address)
    if mask.sum() > 0:
      data[:, mask] = 0
    return data

  def get_data_of_id(self, ids):
    """
      ids: torch.Tensor, shape : [n_ids], dtype : int64
      returns: torch.Tensor, [code_size, n_ids]
      if an id doesn't exist, its retrieved data will be zeros.
    """
    address = self.get_address_of_id(ids)
    data = self.get_data_of_address(address)
    return data

  def set_data_of_address(self, data, address):
    """
      data: torch.Tensor, shape : [code_size // n_cs, n_data, n_cs], dtype : uint8
      address: torch.Tensor, shape : [n_data], dtype : int64
    """
    self.storage[:, address] = data

  def set_data_of_id(self, data, ids):
    """
      data: torch.Tensor, shape : [code_size // n_cs, n_data, n_cs], dtype : uint8
      ids: torch.Tensor, shape : [n_data], dtype : int64
    """
    address = self.get_address_of_id(ids)
    set_data_of_address(data, address)

  def _get_ioa_old(self, labels, unique_labels=None):
    if unique_labels is None:
      unique_labels = torch.unique(labels) #[n_unique_clusters]
    expanded_labels = labels[:, None].expand(-1, unique_labels.shape[0]) #[n_data, n_unique_clusters]
    mask = expanded_labels == unique_labels[None, :] #[n_data, n_unique_clusters]
    mcs = mask.cumsum(dim=0)
    mcs[[~mask]] = 0
    ioa = mcs.sum(dim=1) - 1

  def get_ioa(self, labels, unique_labels=None):
    if labels.device.type == "cuda":
      return self._get_ioa_cuda(labels, unique_labels)
    elif labels.device.type == "cpu":
      return self._get_ioa_old(labels, unique_labels)

  def _get_write_address_old(self, empty_adr, div_of_empty_adr, labels, ioa):
    write_address = torch.zeros_like(labels)
    for i in range(n_labels):
      div_mask = div_of_empty_adr == labels[i]
      write_adr = empty_adr[div_mask]
      write_adr = write_adr[ioa[i] ]
      write_address[i] = write_adr
    return write_address

  def get_write_address(self, empty_adr, labels, div_of_empty_adr=None, ioa=None):
    if div_of_empty_adr is None:
      div_of_empty_adr = self.get_div_of_address(empty_adr)
    if ioa is None:
      ioa = self.get_ioa(labels)

    if labels.device.type == "cuda":
      return self._get_write_address_cuda(empty_adr, div_of_empty_adr, labels, ioa)
    elif labels.device.type == "cpu":
      return self._get_write_address_old(empty_adr, div_of_empty_adr, labels, ioa)

  def expand(self, clusters):
    """
        clusters: an iterable of integers, len(clusters) <= n_cq_clusters
    """
    if self.verbose > 1:
      print(f"Expanding {len(clusters)} clusters")
    tot = 0
    storage = self.storage
    address2id = self.address2id
    is_empty = self.is_empty
    del self.storage, self.address2id, self.is_empty
    for cluster_index in clusters:
      div_start = self.div_start[cluster_index]
      div_cap = self.div_capacity[cluster_index].item()
      div_end = div_start + div_cap
      new_block = torch.zeros(
        self.code_size // self.n_cs,
        div_cap,
        self.n_cs, device=self.device,
        dtype=torch.uint8,
      )
      storage = torch.cat([
        storage[:, :div_end],
        new_block,
        storage[:, div_end:]
      ], dim=1)

      new_a2i = torch.ones(div_cap, device=self.device, dtype=address2id.dtype) * -1
      address2id = torch.cat([
        address2id[:div_end],
        new_a2i,
        address2id[div_end:]
      ], dim=0)

      new_is_empty = torch.ones(div_cap, device=self.device, dtype=is_empty.dtype)
      is_empty = torch.cat([
        is_empty[:div_end],
        new_is_empty,
        is_empty[div_end:]
      ], dim=0)

      self.div_capacity[cluster_index] += div_cap
      arange = torch.arange(start=cluster_index+1, end=self.n_cq_clusters)
      self.div_start[arange] += div_cap
      # self.div_start[:] = self.div_capacity.cumsum(dim=0)
      tot += div_cap
    self.register_buffer("storage", storage)
    self.register_buffer("address2id", address2id)
    self.register_buffer("is_empty", is_empty)
    if self.verbose > 1:
      print(f"Storage capacity is increased by {tot}, total capacity: {self.storage.shape[1]}")
  
  def add(self, input, input_ids=None):
    """
      has to be overrided in child class
    """
    raise NotImplementedError("method needs to be overrided in child class")

  def remove(self, remove_ids):
    """
      remove_ids: torch.Tensor, shape : [n_ids], dtype : int64
    """
    assert self._is_trained == True, "Module is not trained"
    n_ids = remove_ids.shape[0]
    address = self.get_address_of_id(remove_ids)
    address = address[address >= 0] #filter out -1 values
    divs = self.get_div_of_address(address)

    self.is_empty[address] = 1
    self.address2id[address] = -1

    unique_divs, counts = divs.unique(return_counts=True)
    self.div_size[unique_divs] -= counts
    if self.verbose > 1:
      print(f"{len(address)} items have been successfully removed, {n_ids - len(address)} items are ignored.")

  def train(self, input, force_retrain=False):
    """
      has to be overrided in child class
    """
    raise NotImplementedError("method needs to be overrided in child class")
  
  def encode(self, input):
    """
      has to be overrided in child class
    """
    raise NotImplementedError("method needs to be overrided in child class")

  def decode(self, code):
    """
      has to be overrided in child class
    """
    raise NotImplementedError("method needs to be overrided in child class")
    
  def topk(self, query, k, mode=1):
    """
      has to be overrided in child class
    """
    raise NotImplementedError("method needs to be overrided in child class")
    