# ......................................................................................
# MIT License

# Copyright (c) 2019-2025 Pantelis I. Kaplanoglou

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# ......................................................................................
import os
import numpy as np
import pandas as pd
from abc import ABC, abstractmethod
from .sample_set_simple import SampleSet
from .sample_set_kind import SampleSetKind
from .sample_preprocessor import SamplePreprocessor, VoidPreprocessor
from .errors import *
from radnn.errors import ERR_MLSYS_FILESYS_NOT_INITIALIZED
from radnn import FileStore
from radnn import mlsys
from radnn.utils import equal_dicts
from .constants import DataPreprocessingKind
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from .meta.data_manifest import DataManifest
from .cache import DataCachePickle
from .cache.constants import DataPurpose




# ======================================================================================================================
class DataSetCallbacks(object):
  # --------------------------------------------------------------------------------------------------------------------
  def __init__(self, download_method=None, seed_method=None):
    self._lazy_download = download_method
    self._random_seed = seed_method
  # --------------------------------------------------------------------------------------------------------------------
  def lazy_download(self, fs):
    self._lazy_download(fs)
  # --------------------------------------------------------------------------------------------------------------------
  def random_seed(self, seed: int):
    self._random_seed(seed)
  # --------------------------------------------------------------------------------------------------------------------
# ======================================================================================================================s





# ======================================================================================================================
class DataSetBase(ABC):
  # --------------------------------------------------------------------------------------------------------------------
  # Constructor
  def __init__(self, name: str, variant: str | None = None, file_store=None, random_seed: int | None = None,
               callbacks: DataSetCallbacks | None = None, display_name: str | None = None,):
    # ..................// Instance Fields \\.........................
    self.name = name
    self.variant = variant
    self.display_name = display_name
    self.fs: FileStore|None = None
    self._cache_fs: FileStore|None = None
    self._determine_local_filestore(file_store)
    assert self.fs is not None, ERR_DS_MUST_PROVIDE_LOCAL_FILESTORE
    self.cache: DataCachePickle = DataCachePickle(self._cache_fs, self.dataset_code)

    self.random_seed = random_seed
    
    self.all_ids = None
    self.all_samples = None
    self.all_labels = None
    
    self.is_vector_samples = True
    self.are_features_on_last_axis = False
    
    self.all_sample_count = None
    self.feature_count = None
    self.class_count = None
    self._class_names = None

    self.callbacks: DataSetCallbacks = callbacks
    
    self.hprm: dict | None = None
    self.ts: SampleSet | None = None
    self.vs: SampleSet | None = None
    self.us: SampleSet | None = None
    self.preprocessor: SamplePreprocessor = VoidPreprocessor(self)
    self.has_been_split = False
    # ................................................................
  
  # --------------------------------------------------------------------------------------
  @property
  def class_names(self):
    return self._class_names
    # --------------------------------------------------------------------------------------
  @class_names.setter
  def class_names(self, value):
    self._class_names = value
    if isinstance(self._class_names, list):
      self.class_count = len(self._class_names)
    elif isinstance(self._class_names, dict):
      self.class_count = len(self._class_names.keys())
  # --------------------------------------------------------------------------------------
  @property
  def class_names_list(self):
    if isinstance(self._class_names, dict):
      return [v for k,v in self._class_names.items()]
    else:
      return self._class_names
  # --------------------------------------------------------------------------------------
  @property
  def is_split(self):
    return self.cache.is_split
  # --------------------------------------------------------------------------------------
  def split(self, validation_samples_pc=0.10,
            preprocessing: DataPreprocessingKind | None = None,
            random_seed: int=2021):
    
    nTS_Samples, nVS_Samples, nTS_Labels, nVS_Labels = train_test_split(self.all_samples, self.all_labels,
                                                                        test_size=validation_samples_pc,
                                                                        random_state=random_seed)

    # (Re)creating the subsets of the dataset after the splits have been created
    self.ts = SampleSet(self, nTS_Samples, nTS_Labels, kind=SampleSetKind.TRAINING_SET)
    self.vs = SampleSet(self, nVS_Samples, nVS_Labels, kind=SampleSetKind.VALIDATION_SET)

    self.preprocess(preprocessing)
    self.has_been_split = True
    return self
  # --------------------------------------------------------------------------------------------------------------------
  def preprocess(self, kind: DataPreprocessingKind | None = DataPreprocessingKind.STANDARDIZE, preprocessor = None):
    # [TODO] Cover in-disk samples with generators
    if kind is None:
      return

    self.preprocessor = preprocessor
    if self.ts is not None:
      if self.preprocessor is not None:
        self.preprocessor.fit(self.ts.samples)
      elif kind == DataPreprocessingKind.MIN_MAX_NORMALIZE:
        self.preprocessor = MinMaxScaler().fit(self.ts.samples)
      elif kind == DataPreprocessingKind.STANDARDIZE:
        self.preprocessor = StandardScaler().fit(self.ts.samples)
        self.ts.mean = self.preprocessor.mean_
        self.ts.std = self.preprocessor.scale_
        self.update_manifest()

    if self.preprocessor is not None:
      self.ts.samples = self.preprocessor.transform(self.ts.samples)

      if self.vs is not None:
        self.vs.samples = self.preprocessor.transform(self.vs.samples)

      if self.us is not None:
        self.us.samples = self.preprocessor.transform(self.us.samples)
    return self
  # --------------------------------------------------------------------------------------------------------------------
  @property
  def dataset_code(self):
    sUniqueName = f"{self.name.upper()}"
    if self.variant is not None:
      sUniqueName += f"_{self.variant.upper()}"
    return sUniqueName
  # --------------------------------------------------------------------------------------------------------------------
  def _determine_local_filestore(self, file_store):
    if (file_store is not None):
      if isinstance(file_store, FileStore):
        self.fs = file_store
        self._cache_fs = self.fs
      elif isinstance(file_store, str):
        if not os.path.exists(file_store):
          raise Exception(ERR_DS_FOLDER_NOT_FOUND % file_store)
        self.fs = FileStore(file_store)
        self._cache_fs = self.fs
    else:
      assert mlsys.filesys is not None, ERR_MLSYS_FILESYS_NOT_INITIALIZED
      
      self.fs = mlsys.filesys.datasets.subfs(self.dataset_code)
      self._cache_fs = mlsys.filesys.caches.subfs(self.dataset_code)
  # --------------------------------------------------------------------------------------------------------------------
  @property
  def filesystem_folder(self):
    return self.fs.absolute_path
  # --------------------------------------------------------------------------------------------------------------------
  def read_hyperparams(self):
    pass # Optionally override
  # --------------------------------------------------------------------------------------------------------------------
  @abstractmethod
  def load_data(self):
    pass # Must implement
  # --------------------------------------------------------------------------------------------------------------------
  def prepare(self, hyperparams: dict | None = None):
    self.hprm = hyperparams
    
    # VIRTUAL CALL: Reads the hyperparameters into instance variables
    if self.hprm is not None:
      self.read_hyperparams()
    
    if (self.callbacks is not None):
      assert self.callbacks._lazy_download is not None, ERR_DS_NO_RANDOM_SEED_INITIALIZER_CALLBACK
      if self.callbacks._lazy_download is not None:
        self.callbacks.lazy_download(self.fs)
    
    if (self.random_seed is not None):
      bIsInitRandomSeed = False
      if self.callbacks is not None:
        if self.callbacks._random_seed is not None:
          self.callbacks.random_seed(self.random_seed)
          bIsInitRandomSeed = True
      if not bIsInitRandomSeed:
        mlsys.random_seed_all(self.random_seed)

    self.ts = None
    self.vs = None
    self.us = None
    

    
    # VIRTUAL CALL:  Imports the dataset from the source local/remote filestore to the local cache.
    self.load_data()
    
    if self.is_split:
      assert self.ts is not None, ERR_DS_SUBSET_MUST_HAVE_TS
      assert self.ts.kind == SampleSetKind.TRAINING_SET, ERR_DS_SUBSET_INVALID_SETUP
      if self.vs is not None:
        assert self.vs.kind == SampleSetKind.VALIDATION_SET, ERR_DS_SUBSET_INVALID_SETUP
      self.has_been_split = True

    # Count the total samples of a pre-splitted dataset
    if self.all_sample_count is None:
      if self.is_split:
        self.all_sample_count = self.ts.sample_count
        if self.vs is not None:
          self.all_sample_count += self.vs.sample_count
      else:
        self.all_sample_count = len(self.all_samples)
        
        
    if self.us is not None:
      assert self.us.kind == SampleSetKind.UNKNOWN_TEST_SET, ERR_DS_SUBSET_INVALID_SETUP

      
    self.manifest = DataManifest(self.dataset_code)
    self.update_manifest()

    return self
  # --------------------------------------------------------------------------------------------------------------------
  def get_manifest(self):
    oNewManifest = DataManifest(self.dataset_code)
    dManifest = oNewManifest.dict
    dManifest["is_split"] = self.is_split
    dManifest["name"] = self.name
    dManifest["variant"] = self.variant
    dManifest["cache_kind"] = self.cache.kind.name
    dManifest["random_seed"] = self.random_seed
    dManifest["ground_truth_set"]["all_sample_count"] = self.all_sample_count
    dManifest["feature_count"] = self.feature_count
    
    if self.class_count is not None:
      dManifest["task"]["class_count"] = self.class_count
      dManifest["task"]["class_names"] = self.class_names
    
    dManifest["is_vector_samples"] = self.is_vector_samples
    dManifest["are_features_on_last_axis"] = self.are_features_on_last_axis
    
    if not self.is_split:
      dManifest["ground_truth_set"]["sample_files"].append(self.cache.data_file(None, DataPurpose.SAMPLE_FEATURES))
      dManifest["ground_truth_set"]["label_files"].append(self.cache.data_file(None, DataPurpose.SAMPLE_LABELS))
      if self.cache.has_ids:
        dManifest["ground_truth_set"]["id_files"].append(self.cache.data_file(None, DataPurpose.SAMPLE_IDS))
      
      if self.all_samples is not None:
        oNewManifest.set_shapes(self.all_samples, self.all_labels)
    else:
      if self.ts is not None:
        oNewManifest.set_ts_info(self.ts)
        dManifest["training_set"]["sample_files"].append(self.cache.data_file("TS", DataPurpose.SAMPLE_FEATURES))
        dManifest["training_set"]["label_files"].append(self.cache.data_file("TS", DataPurpose.SAMPLE_LABELS))
        if self.cache.has_ids:
          dManifest["training_set"]["id_files"].append(self.cache.data_file("TS", DataPurpose.SAMPLE_IDS))
      
      if self.vs is not None:
        oNewManifest.set_vs_info(self.vs)
        dManifest["validation_set"]["sample_files"].append(self.cache.data_file("VS", DataPurpose.SAMPLE_FEATURES))
        dManifest["validation_set"]["label_files"].append(self.cache.data_file("VS", DataPurpose.SAMPLE_LABELS))
        if self.cache.has_ids:
          dManifest["validation_set"]["id_files"].append(self.cache.data_file("VS", DataPurpose.SAMPLE_IDS))
          
    return oNewManifest
  # --------------------------------------------------------------------------------------------------------------------
  def update_manifest(self):
    bMustCreateManifest = not self.manifest.load(self._cache_fs)
    
    oCurrentManifest = self.get_manifest()
 
    bMustSave = False
    if bMustCreateManifest:
      self.manifest = oCurrentManifest
      bMustSave = True
    else:
      if not equal_dicts(self.manifest.dict, oCurrentManifest.dict):
        # Update only if it is the same format
        if self.manifest.format == oCurrentManifest.format:
          self.manifest = oCurrentManifest
          bMustSave = True
        else:
          # TODO: Convert to newer format or enable compatibility layer
          pass
        
    if bMustSave:
      self.manifest.save(self._cache_fs)
  # --------------------------------------------------------------------------------------------------------------------
  def assign(self, data, label_columns: range):
    self.all_samples, self.all_labels, self.all_ids = None, None, None
    if isinstance(data, tuple):
      self.all_samples, self.all_labels = data
    elif isinstance(data, np.ndarray):
      self.all_samples = data
    elif isinstance(data, dict):
      if ("samples" in dict):
        self.all_samples = data["samples"]
      if ("labels" in dict):
        self.all_labels = data["labels"]
      if ("ids" in dict):
        self.all_ids = data["ids"]
    elif isinstance(data, pd.DataFrame):
      if isinstance(data.columns, pd.Index):
        nData = data.iloc[1:].to_numpy()
      else:
        nData = data.to_numpy()
      
      if label_columns is None:
        self.all_samples = nData
      else:
        if label_columns.start >= 0:
          if label_columns.stop is None:
            self.all_labels = nData[:, label_columns.start]
            self.all_samples = nData[:, label_columns.start + 1:]
          else:
            self.all_labels = nData[:, label_columns.start:label_columns.stop + 1]
            self.all_samples = nData[:, label_columns.stop + 1:]
        else:
          self.all_samples = nData[:, :label_columns.start]
          self.all_labels = nData[:, label_columns.start:]
      
    if self.all_ids is None:
      self.all_ids = range(len(self.all_samples)) + 1
          
    return self
  # --------------------------------------------------------------------------------------------------------------------
  def print_info(self):
    print(f"Dataset [{self.dataset_code}]")
    self.ts.print_info()
    if self.vs is not None:
      self.vs.print_info()
    if self.us is not None:
      self.us.print_info()
  # --------------------------------------------------------------------------------------------------------------------
