import numpy as np
from radnn import FileStore
from .constants import DataCacheKind, DataPurpose
from radnn.data.errors import *




class DataCachePickle(object):
  # --------------------------------------------------------------------------------------------------------------------
  def __init__(self, fs: FileStore, file_prefix: str, kind: DataCacheKind = DataCacheKind.TENSORS_LABELS, is_split: bool = True):
    self.fs = fs
    self.file_prefix = file_prefix
    self.kind = kind
    self.is_split = is_split
    self.has_ids = (self.kind.value % 10 == 2)
    self._existing_cache_kind = None
  # --------------------------------------------------------------------------------------------------------------------
  def prepare(self, kind: DataCacheKind = DataCacheKind.TENSORS_LABELS, is_split: bool = True):
    self.kind = kind
    self.is_split = is_split
    self.has_ids = (self.kind.value % 10 == 2)
  # --------------------------------------------------------------------------------------------------------------------
  def _load_manifest(self ):
    oFiles = self.fs.list_files("*.json", is_full_path=False)
    for sFile in oFiles:
      if sFile.endswith(".manifest.json"):
        dManifest = self.fs.json.load(sFile)
        self._existing_cache_kind = dManifest["cache_kind"]
        break
  # --------------------------------------------------------------------------------------------------------------------
  @property
  def exists(self):
    self._load_manifest()
    if self._existing_cache_kind is not None:
      bExists = self._existing_cache_kind.upper() == self.kind.name.upper()
    else:
      bExists = True
      
    if bExists:
      oFiles = self.fs.list_files("*.pkl", is_full_path=False)
      bExists = False
      for sFile in oFiles:
        if (sFile.upper().startswith(self.file_prefix.upper()) and sFile.lower().endswith(".pkl")):
          bExists = True
          break
    return bExists
  # --------------------------------------------------------------------------------------------------------------------
  def _get_suffix(self):
    if self.kind.value >= DataCacheKind.TENSORS_LAST_AXIS_FEATURES.value:
      sResult = ""
    elif self.kind.value >= DataCacheKind.TENSORS.value:
      sResult = "-torch"
    else:
      sResult = "-vec"
    return sResult
  # --------------------------------------------------------------------------------------------------------------------
  def data_file(self, subset_code: str, purpose: DataPurpose):
    sSuffix = self._get_suffix()
    
    if subset_code is None:
      subset_code = ""
    else:
      subset_code = "-" + subset_code
    
    if purpose == DataPurpose.SAMPLE_FEATURES:
      return f"{self.file_prefix}{subset_code}-Samples{sSuffix}.pkl"
    elif purpose == DataPurpose.SAMPLE_LABELS:
      return f"{self.file_prefix}{subset_code}-Labels{sSuffix}.pkl"
    elif purpose == DataPurpose.SAMPLE_IDS:
      return f"{self.file_prefix}{subset_code}-Ids{sSuffix}.pkl"
  # --------------------------------------------------------------------------------------------------------------------
  def load_cache(self, is_split=False):
    # .......................................................................
    def load_file(subset_code: str, purpose: DataPurpose):
      return self.fs.obj.load(self.data_file(subset_code, purpose))
    # .......................................................................
    nType = self.kind.value % 10
    
    if self.is_split:
      nTSSamples = load_file("TS", DataPurpose.SAMPLE_FEATURES)
      nVSSamples = load_file("VS", DataPurpose.SAMPLE_FEATURES)
      
      nTSLabels = load_file("TS", DataPurpose.SAMPLE_LABELS)
      nVSLabels = load_file("VS", DataPurpose.SAMPLE_LABELS)
      
      nTSIds = load_file("TS", DataPurpose.SAMPLE_IDS)
      nVSIds = load_file("VS", DataPurpose.SAMPLE_IDS)
      
      if self.has_ids:
        oLoadedPack = (nTSSamples, nVSSamples, nTSLabels, nVSLabels, nTSIds, nVSIds)
      else:
        oLoadedPack = (nTSSamples, nVSSamples, nTSLabels, nVSLabels)
    else:
      nSamples = load_file(None, DataPurpose.SAMPLE_FEATURES)
      nLabels = load_file(None, DataPurpose.SAMPLE_LABELS)
      nIds = load_file(None, DataPurpose.SAMPLE_IDS)
      
      if self.has_ids:
        oLoadedPack = (nSamples, nLabels, nIds)
      else:
        oLoadedPack = (nSamples, nLabels)
    
    return oLoadedPack
  # --------------------------------------------------------------------------------------------------------------------
  def compact_array(self, array: np.ndarray):
    nResult = array
    if np.issubdtype(nResult.dtype, np.integer):
      nMin = array.min()
      nMax = array.max()
      
      tByte = np.iinfo(np.uint8)
      tUShort = np.iinfo(np.uint16)
      tInt = np.iinfo(np.int32)
      
      if nMin >= tByte.min and nMax <= tByte.max:
        nResult = array.astype(np.uint8)
      elif nMin >= tUShort.min and nMax <= tUShort.max:
        nResult = array.astype(np.uint16)
      elif nMin >= tInt.min and nMax <= tInt.max:
        nResult = array.astype(np.int32)
    return nResult
  # --------------------------------------------------------------------------------------------------------------------
  def save_cache(self, data_pack=None, **kwargs):
    # .......................................................................
    def save_file(data, subset_code: str, purpose: DataPurpose):
      data = self.compact_array(data)
      self.fs.obj.save(data, self.data_file(subset_code, purpose))
    # .......................................................................

    bIsUnpacked = False
    if data_pack is None and (len(kwargs) > 0):
      for k,v in kwargs.items():
        if k == "ts_samples":
          nTSSamples = v
        elif k == "ts_labels":
          nTSLabels = v
        elif k == "ts_ids":
          nTSIds = v
        elif k == "vs_samples":
          nVSSamples = v
        elif k == "vs_labels":
          nVSLabels = v
        elif k == "vs_ids":
          nVSIds = v
        elif k == "samples":
          nSamples = v
        elif k == "labels":
          nLabels = v
        elif k == "ids":
          nIds = v
      bIsUnpacked = True
      
    if not bIsUnpacked:
      oSavePack = data_pack
      if oSavePack is not None:
        if isinstance(data_pack, tuple):
          oSavePack = list(data_pack)
        elif isinstance(data_pack, list):
          oSavePack = data_pack
        else:
          oSavePack = None
        
      if oSavePack is None:
        raise Exception(ERR_DS_CACHE_INVALID_PACK)
  
      if self.is_split:
        if self.has_ids:
          assert len(oSavePack) == 6, ERR_DS_CACHE_INVALID_PACK_LEN % (self.kind.name, "pre-split ", 6)
          nTSSamples, nVSSamples, nTSLabels, nVSLabels, nTSIds, nVSIds = oSavePack
        else:
          assert len(oSavePack) == 4, ERR_DS_CACHE_INVALID_PACK_LEN % ( self.kind.name, "pre-split ",4)
          nTSSamples, nVSSamples, nTSLabels, nVSLabels = oSavePack
      else:
        if self.has_ids:
          assert len(oSavePack) == 3, ERR_DS_CACHE_INVALID_PACK_LEN % (self.kind.name, "", 3)
          nSamples, nLabels, nIds = oSavePack
        else:
          assert len(oSavePack) == 2, ERR_DS_CACHE_INVALID_PACK_LEN % (self.kind.name, "", 2)
          nSamples, nLabels = oSavePack
      
    if self.is_split:
      save_file(nTSSamples, "TS", DataPurpose.SAMPLE_FEATURES)
      save_file(nVSSamples, "VS", DataPurpose.SAMPLE_FEATURES)

      save_file(nTSLabels, "TS", DataPurpose.SAMPLE_LABELS)
      save_file(nVSLabels, "VS", DataPurpose.SAMPLE_LABELS)
      
      if self.has_ids:
        save_file(nTSIds, "TS", DataPurpose.SAMPLE_IDS)
        save_file(nVSIds, "VS", DataPurpose.SAMPLE_IDS)
    else:
      save_file(nSamples, None, DataPurpose.SAMPLE_FEATURES)
      save_file(nLabels, None, DataPurpose.SAMPLE_LABELS)
      if self.has_ids:
        save_file(nIds, None, DataPurpose.SAMPLE_IDS)
  # --------------------------------------------------------------------------------------------------------------------