import numpy as np
from radnn.data.errors import ERR_DS_SAMPLE_INDEX_MISMATCH

class SampleSearchIndex(object):
  # --------------------------------------------------------------------------------------------------------------------
  def __init__(self, ids: np.ndarray, samples: np.ndarray, labels: np.ndarray = None, class_names = None):
    if len(ids) != len(samples):
      raise Exception(ERR_DS_SAMPLE_INDEX_MISMATCH)
    
    self.ids = ids.squeeze()
    self.samples = samples
    self.labels = labels
    self.class_names = class_names
    
    self.sorted_order = np.argsort(self.ids).squeeze()
    self.sorted_ids = ids[self.sorted_order].squeeze()
  # --------------------------------------------------------------------------------------------------------------------
  def sample_by_id(self, search_id):
    search_id = int(search_id)
    nSample = None
    idx = np.searchsorted(self.sorted_ids, search_id)
    
    if (idx < len(self.sorted_ids)) and (self.sorted_ids[idx] == search_id):
      original_index = self.sorted_order[idx]
      if self.samples.dtype == object:
        nSample = self.samples[original_index, 0]
      else:
        nSample = self.samples[original_index]
      
    return nSample
  # --------------------------------------------------------------------------------------------------------------------
  def __getitem__(self, search_id):
    search_id = int(search_id)
    idx = np.searchsorted(self.sorted_ids, search_id)
    
    if (idx < len(self.sorted_ids)) and (self.sorted_ids[idx] == search_id):
      original_index = self.sorted_order[idx]
      if self.samples.dtype == object:
        nSample = self.samples[original_index, 0]
      else:
        nSample = self.samples[original_index]
      
      dSampleRecord = {
        "type": "generic",
        "id": search_id,
        "sample": nSample,
        "label": None,
        "class_name": None
      }

      if self.labels is not None:
        nLabel = int(self.labels[original_index])
        dSampleRecord["label"] = nLabel
        if self.class_names is not None:
          dSampleRecord["class_name"] = self.class_names[nLabel]
          
    return dSampleRecord
  # --------------------------------------------------------------------------------------------------------------------
  def __contains__(self, key):
    idx = np.searchsorted(self.sorted_ids, key)
    bContains = (idx < len(self.sorted_ids)) and (self.sorted_ids[idx] == key)
    return bContains
  # --------------------------------------------------------------------------------------------------------------------
  def __len__(self):
    return len(self.sorted_ids)
  # --------------------------------------------------------------------------------------------------------------------