from radnn import MLGenre, FileStore
from radnn.data.sample_set_simple import SampleSet
from collections import OrderedDict
import radnn
import platform

class DataManifest(dict):
  def __init__(self, dataset_code):
    # --------------------------------------------------------------------------------------------------------------------
    self.format = "MLD1"
    self.dataset_code = dataset_code
    
    self.dict = OrderedDict()
    
    self.dict["format"] = self.format
    self.dict["name"] = None
    self.dict["variant"] = None
    self.dict["cache_kind"] = None
    self.dict["task"] = OrderedDict()
    
    self.dict["feature_count"] = None
    self.dict["sample_shape"] = None
    self.dict["label_shape"] = None
    self.dict["is_vector_samples"] = None
    self.dict["are_features_on_last_axis"] = None

    self.dict["random_seed"] = None
    self.dict["is_split"] = False
    self.dict["ground_truth_set"] = OrderedDict()
    self.dict["ground_truth_set"]["sample_files"] = []
    self.dict["ground_truth_set"]["label_files"] = []
    self.dict["ground_truth_set"]["id_files"] = []
    self.dict["ground_truth_set"]["all_sample_count"] = None
    self.dict["training_set"] = None
    self.dict["validation_set"] = None
    self.dict["unknown_test_set"] = None
    self.dict["pickle_version"] = platform.python_version()
    self.dict["radnn_version"] = radnn.__version__
  # --------------------------------------------------------------------------------------------------------------------
  def set_shapes(self, samples, labels):
    self.dict["sample_shape"] = list(samples[0].shape)
    if labels is not None:
      oShape = list(labels[0].shape)
      if oShape == []:
        oShape = [1]
      self.dict["label_shape"] = oShape
  
  # --------------------------------------------------------------------------------------------------------------------
  def set_ts_info(self, ts: SampleSet):
    if ts is not None:
      self.dict["training_set"] = OrderedDict()
      self.dict["training_set"]["sample_files"] = []
      self.dict["training_set"]["label_files"] = []
      self.dict["training_set"]["id_files"] = []
      self.dict["training_set"]["sample_count"] = ts.sample_count
      self.dict["training_set"]["mean"] = ts.mean
      self.dict["training_set"]["std"] = ts.std
      self.set_shapes(ts.samples, ts.labels)
  # --------------------------------------------------------------------------------------------------------------------
  def set_vs_info(self, vs: SampleSet):
    if vs is not None:
      self.dict["validation_set"] = OrderedDict()
      self.dict["validation_set"]["sample_files"] = []
      self.dict["validation_set"]["label_files"] = []
      self.dict["validation_set"]["id_files"] = []
      self.dict["validation_set"]["sample_count"] = vs.sample_count
      self.dict["validation_set"]["mean"] = vs.mean
      self.dict["validation_set"]["std"] = vs.std
  # --------------------------------------------------------------------------------------------------------------------
  def load(self, fs: FileStore):
    dLoaded = fs.json.load(f"{self.dataset_code}.manifest.json")
    if dLoaded is not None:
      self.dict = dLoaded
      self.format = self.dict["format"]
      
    return dLoaded is not None
  # --------------------------------------------------------------------------------------------------------------------
  def save(self, fs: FileStore):
    fs.json.save(self.dict, f"{self.dataset_code}.manifest.json")
  # --------------------------------------------------------------------------------------------------------------------
