# ......................................................................................
# 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 import FileStore
from radnn import mlsys
from .constants import DataPreprocessingKind
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler, StandardScaler

# ======================================================================================================================
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):
    # ..................// Instance Fields \\.........................
    self.name = name
    self.variant = variant
    self.fs = None
    self._determine_local_filestore(file_store)
    assert self.fs is not None, ERR_DS_MUST_PROVIDE_LOCAL_FILESTORE
    self.random_seed = random_seed
    
    self.all_ids = None
    self.all_samples = None
    self.all_labels = None
    
    self.feature_count = None
    self.class_count = None
    self.sample_count = 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.is_split = False
    # ................................................................

  # --------------------------------------------------------------------------------------
  def split(self, validation_samples_pc=0.10,
            preprocessing: DataPreprocessingKind | None = DataPreprocessingKind.STANDARDIZE,
            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)
    if preprocessing == DataPreprocessingKind.MIN_MAX_NORMALIZE:
      self.preprocessor = MinMaxScaler().fit(nTS_Samples)
    elif preprocessing == DataPreprocessingKind.STANDARDIZE:
      self.preprocessor = StandardScaler().fit(nTS_Samples)
    else:
      self.preprocessor = None

    if self.preprocessor is not None:
      nTS_Samples = self.preprocessor.transform(nTS_Samples)
      nVS_Samples = self.preprocessor.transform(nVS_Samples)

    # (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)
    if preprocessing == DataPreprocessingKind.STANDARDIZE:
      self.ts.mean = self.preprocessor.mean_
      self.ts.std = self.preprocessor.scale_

    self.vs = SampleSet(self, nVS_Samples, nVS_Labels, kind=SampleSetKind.VALIDATION_SET)
    
    self.is_split = True
    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
      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)
    else:
      assert mlsys.filesys is not None, ERR_MLSYS_FILESYS_NOT_INITIALIZED
      
      self.fs: FileStore = mlsys.filesys.datasets.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 load_cache(self, is_vector_samples=True, is_last_axis_features=True):
    nSuffix = ""
    if is_vector_samples:
      nSuffix = "-vec"
    elif not is_last_axis_features:
      nSuffix = "-torch"

    nTSSamples = self.fs.obj.load(f"{self.dataset_code}-TS-Samples{nSuffix}.pkl")
    nVSSamples = self.fs.obj.load(f"{self.dataset_code}-VS-Samples{nSuffix}.pkl")

    nTSLabels = self.fs.obj.load(f"{self.dataset_code}-TS-Labels{nSuffix}.pkl")
    nVSLabels = self.fs.obj.load(f"{self.dataset_code}-VS-Labels{nSuffix}.pkl")

    return nTSSamples, nVSSamples, nTSLabels, nVSLabels
  # --------------------------------------------------------------------------------------------------------------------
  def save_cache(self, ts_samples, vs_samples, ts_labels, vs_labels, is_vector_samples=True, is_last_axis_features=True):
    nSuffix = ""
    if is_vector_samples:
      nSuffix = "-vec"
    elif not is_last_axis_features:
      nSuffix = "-torch"
    self.fs.obj.save(ts_samples, f"{self.dataset_code}-TS-Samples{nSuffix}.pkl")
    self.fs.obj.save(vs_samples, f"{self.dataset_code}-VS-Samples{nSuffix}.pkl")

    self.fs.obj.save(ts_labels, f"{self.dataset_code}-TS-Labels{nSuffix}.pkl")
    self.fs.obj.save(vs_labels, f"{self.dataset_code}-VS-Labels{nSuffix}.pkl")
  # --------------------------------------------------------------------------------------------------------------------
  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
        
    if self.us is not None:
      assert self.us.kind == SampleSetKind.UNKNOWN_TEST_SET, ERR_DS_SUBSET_INVALID_SETUP
      
    return self
  # --------------------------------------------------------------------------------------------------------------------
  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()
  # --------------------------------------------------------------------------------------------------------------------
