#  Copyright (c) 2024. Affects AI LLC
#
#  Licensed under the Creative Common CC BY-NC-SA 4.0 International License (the "License");
#  you may not use this file except in compliance with the License. The full text of the License is
#  provided in the included LICENSE file. If this file is not available, you may obtain a copy of the
#  License at
#
#       https://creativecommons.org/licenses/by-nc-sa/4.0/deed.en
#
#  Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
#  express or implied. See the License for the specific language governing permissions and limitations
#  under the License.

import os
import logging
from pathlib import Path

import scipy

from ardt import config
from ardt.datasets import AERDataset
from .AscertainTrial import AscertainTrial

CONFIG = config['datasets']['ascertain']
DEFAULT_ASCERTAIN_PATH = Path(CONFIG['path'])
ASCERTAIN_RAW_FOLDER = Path(CONFIG['raw_data_path'])
ASCERTAIN_FEATURES_FOLDER = Path(CONFIG['features_data_path'])
ASCERTAIN_NUM_MEDIA_FILES = 36
ASCERTAIN_NUM_PARTICIPANTS = 58

logger = logging.getLogger('AscertainDataset')
logger.level = logging.DEBUG

expected_classifications = {
            # Taken from DECAF+ mediafile_offset: MEG-BASED MULTIMODAL DATABASE FOR DECODING AFFECTIVE PHYSIOLOGICAL RESPONSES
            # Amusing
            1: 1,   # Ace Ventura: Pet Detective
            2: 1,   # The Gods Must Be Crazy II
            4: 1,   # Airplane
            5: 1,   # When Harry Met Sally

            # Funny:
            3: 1,   # Liar Liar
            6: 1,   # The Gods Must Be Crazy
            7: 1,   # The Hangover
            9: 1,   # Hot Shots

            # Happy:
            8: 4,   # Up
            10: 4,  # August Rush
            11: 4,  # Truman Show
            12: 4,  # Wall-E
            13: 4,  # Love Actually
            14: 4,  # Remember the Titans
            16: 4,  # Life is Beautiful
            17: 4,  # Slumdog Millionaire
            18: 4,  # House of Flying Daggers

            # Exciting
            15: 1,  # Legally Blonde
            33: 1,  # The Untouchables

            # Angry
            19: 3,  # Ghandi
            21: 3,  # Lagaan
            23: 3,  # My Bodyguard
            35: 3,  # Crash

            # Disgusting
            28: 2,  # The Exorcist
            34: 2,  # Pink Flamingos

            # Fear:
            30: 2,  # The Shining
            36: 2,  # Black Swan

            # Sad
            20: 3,  # My Girl
            22: 3,  # Bambi
            24: 3,  # Up
            25: 3,  # Life is Beautiful
            26: 3,  # Remember the Titans
            27: 3,  # Titanic
            31: 3,  # Prestige

            # Shock
            29: 2,  # Mulholland Drive
            32: 2,  # Alien
        }

default_signal_metadata = {'ECG': {
                'sample_rate': 256,
                'n_channels': 2,
            }
}

class AscertainDataset(AERDataset):
    def __init__(self, ascertain_path=None, signals=None, participant_offset=0, mediafile_offset=0):
        """
        Construct a new AscertainDataset object for a given ascertainPath.

        :param ascertain_path: Path to the extracted ASCERTAIN dataset
        :param signals: A list of signals to load, e.g. ['ECG','EEG'] to load ECG and EEG data. If None, the folder
        given by ascertain_path will be checked for subfolders named '<SIGNAL>Data', and each one found will be loaded.
        For example, if ascertain/ECGData/ is found, the ECG is automatically loaded.
        :param participant_offset: Constant value added to each participant identifier within this dataset. For example,
        if participant_offset is 32, then Participant 1 from this dataset's raw data will be returned as Participant 33.
        :param mediafile_offset: Constant value added to each media identifier within this dataset. For example, if
        mediafile_offset is 12, then Movie 1 from this dataset's raw data will be reported as Media ID 13.
        """
        super().__init__(signals=signals,
                         participant_offset=participant_offset,
                         mediafile_offset=mediafile_offset,
                         signal_metadata=default_signal_metadata,
                         expected_responses=expected_classifications)

        if ascertain_path is None:
            ascertain_path = CONFIG.get('path')

        if ascertain_path is None or not os.path.exists(ascertain_path):
            raise ValueError(
                f'Invalid path to ASCERTAIN dataset: {ascertain_path}. Please correct and try again.')

        logger.debug(f'Loading ASCERTAIN from {ascertain_path} with signals {signals}.')

        self.ascertain_path = Path(ascertain_path)
        self.ascertain_raw_path = self.ascertain_path / ASCERTAIN_RAW_FOLDER
        self.ascertain_features_path = self.ascertain_path / ASCERTAIN_FEATURES_FOLDER
        self.media_index_to_name = {}           # Maps media index back to name

        if not self.ascertain_path.exists():
            raise ValueError('Path to ASCERTAIN does not exist: {}'.format(ascertain_path))

        # Load signals
        if signals is not None:
            for signal in signals:
                if not (self.ascertain_raw_path / f'{signal}Data').exists():
                    raise ValueError(
                        f'{signal}Data does not exist, unable to load {signal} Signal. Please correct and try again.')
        else:
            for p in sorted(self.ascertain_raw_path.rglob("*Data")):
                if p.is_dir():
                    self.signals.append(str(p.name).replace("Data", ""))


    def _preload_dataset(self):
        pass

    def load_trials(self):
        # Load ascertain data files...
        # Map< participantId, Map< movieId, data_file_path >>
        ascertain_datafiles = {}

        dt_selfreports_path = os.path.join(self.ascertain_features_path, "Dt_SelfReports.mat")
        dt_selfreports = scipy.io.loadmat(dt_selfreports_path)

        for matlab_file in self.ascertain_raw_path.rglob("*Clip*.mat"):
            movie_folder = matlab_file.parents[0].name
            signal_folder = matlab_file.parents[1].name

            signal_type = signal_folder.replace("Data", "")
            if signal_type not in self._signals:
                continue

            participant_id = int(movie_folder.split("_P")[1]) # + self.participant_offset
            movie_id = int(matlab_file.name.upper().replace(f'{signal_type}_CLIP', '').replace('.MAT', ''))
            # movie_id += self.media_file_offset

            self.media_index_to_name[movie_id] = movie_id   # no names, just ids... 1:1 map
            if participant_id not in ascertain_datafiles.keys():
                ascertain_datafiles[participant_id] = {}

            if movie_id not in ascertain_datafiles[participant_id].keys():
                ascertain_datafiles[participant_id][movie_id] = {}

            ascertain_datafiles[participant_id][movie_id][signal_type] = matlab_file.resolve()

        def _to_quadrant(a,v):
            q=-1
            if a >= 3:  # A is high
                if v >= 0:  # A is High, V is Neg = Quad 0
                    q = 1
                else:
                    q = 2
            else:
                if v < 0:
                    q = 3
                else:
                    q = 4
            return q

        for participant_id in ascertain_datafiles:
            for movie_id in ascertain_datafiles[participant_id]:
                arousal = dt_selfreports['Ratings'][0][participant_id - 1 - self.participant_offset][movie_id - 1 - self.media_file_offset]
                valence = dt_selfreports['Ratings'][1][participant_id - 1 - self.participant_offset][movie_id - 1 - self.media_file_offset]

                trial = AscertainTrial(self, participant_id, movie_id, _to_quadrant(arousal, valence))
                trial.signal_data_files = ascertain_datafiles[participant_id][movie_id]
                trial.signal_preprocessors = self.signal_preprocessors
                self.trials.append(trial)

    def get_media_name_by_movie_id(self, movie_id):
        return None


