"""
SimpleDataLoader is the most basic interface between the data on the disk and the codebase.
Created to support LMI training with any type of dataset irregadless of having queries / knn-gt / pivots file.

=========
EXAMPLE USAGE:
=========
from utils import load_yaml
config = load_yaml('./config/config-simple.yml')

loader = SimpleDataLoader(config['data'])
df = loader.load_descriptors()
"""
from typing import Dict
from os import path
import pandas as pd
import numpy as np

from sklearn import preprocessing
from lmi_proteins.utils import isfile, load_json

import logging
import time


class SimpleDataLoader:
    """ The most simple form of data loader. Supports only unsupervised LMI training.

    In case of any specific needs in data loading, rewrite the `pd.read_csv` line or
    create a new class inheriting from this one with a custom `load_dataset` function.

    Attributes
    ----------
    base_data_dir : str
        Full path to the base directory of the dataset
    dataset_path : str
        Full path to the dataset
    queries_filename
        Full path to the randomly chosen queries (objects from the dataset)
    knn_gt_filename
        Full path to the 30-NN ground truth (used for experiment evaluation)
    """
    def __init__(self, data: Dict):
        """
        Parameters
        ----------
        data : Dict
            `data` portion of the json config file
        """

        self.base_data_dir = data['data-dir']
        assert 'dataset-file' in data, \
            'Expected to find `dataset-file` in `config`.'\
            'Are you using the correct config file (e.g. `config-simple.yml`)?'

        self.dataset_path = path.join(self.base_data_dir, data['dataset-file'])

        if 'queries' in data:
            self.queries_filename = path.join(self.base_data_dir, data['queries'])

        if 'knn-gt' in data:
            self.knn_filename = path.join(self.base_data_dir, data['knn-gt'])

        self._shuffle = data['shuffle']
        self._normalize = data['normalize']

        logging.basicConfig(level=logging.INFO)
        self.LOG = logging.getLogger(__name__)

    def load_descriptors(self) -> pd.DataFrame:
        """ Loads the descriptors from the disk into the memory.
        Assumes that headers are present, the values are floating point numbers,
        fills any missing values with '0'.
        If `normalize=True` is specified in config, preforms z-score normalization.

        Returns
        ----------
        df: pd.DataFrame
            Dataset of the loaded descriptors with index set to object_id.
        """
        assert isfile(self.dataset_path), f'Expected {self.dataset_path} to exist.'

        self.LOG.info(f'Loading dataset from {self.dataset_path}.')
        time_start = time.time()

        #df = pd.read_csv(self.dataset_path)
        df = pd.read_pickle(self.dataset_path)
        #df = df.set_index(['protein_id'])

        self.LOG.debug(
            f'Loaded dataset of shape: {df.shape}.'
            f' Loading took {round(time.time() - time_start, 2)}s.'
        )

        df = df.fillna(0)
        if self._normalize:
            scaler = preprocessing.StandardScaler()
            df[df.columns] = scaler.fit_transform(df[df.columns])
        if self._shuffle:
            df = df.sample(frac=1)
        return df

    def get_knn_ground_truth(self):
        return load_json(self.knn_filename)

    def get_queries(self, data):
        queries = []
        file_queries = pd.read_csv(self.queries_filename, header=None).values.flatten()
        for query in file_queries:
            queries.append((query, data.loc[query, ].values))
        return queries
    
