"""
Script for parsing the data files generated by the biological motion protocol.
author: Juan Jesus Torre Tresols
e-mail: juan-jesus.torre-tresols@inria.fr
"""

import argparse
import glob
import json
import os

import numpy as np
import pandas as pd
import scipy.io as sio

parser = argparse.ArgumentParser(description='Parameters for the logfiles')
parser.add_argument('-t', '--type', metavar='SubjectType', default='sub',
                    choices=['sub', 'pilot'],
                    help="Subject type. It can only be 'pilot' or 'sub'. "
                         "Choices: %(choices)s. Default: %(default)s")
parser.add_argument('-n', '--number', metavar='SubjectNum', type=int,
                    help="Subject number. It will be formatted as a "
                         "2-digit number (Max. 99)")

args = parser.parse_args()
sub_type = args.type
sub_num = "%02d" % args.number
# session = "%02d" % args.session


class Logfile:
    """
    Class to hold the logfile object. It will open and parse the original .mat
    file, keeping the relevant info as attributes

    Parameters
    ----------

    path: str or path object
          Absolute or relative path to the .mat file

    Attributes
    ----------

    path: str or path object
          Absolute or relative path to the .mat file

    sub_id: str
            ID of the subject

    run: int
         Number of the run

    acq_type: int
              Type of sequence (1 or 2)

    blocks: np.array
            Array indicating the block sequence of the run


    """

    def __init__(self, path: str):
        self.path = path

    def get_data(self, field: str) -> int or str or np.array:
        """Get a particular field from the .mat file"""
        data_struct = sio.loadmat(self.path)['Data']
        data_array = data_struct[field]
        value = data_array[0][0][0]
        return value

    @property
    def sub_id(self):
        sub_id = self.get_data('SubjectID')
        return sub_id

    @property
    def run(self):
        run_num = int(self.get_data('RunNum')[0])
        return run_num

    @property
    def acq_type(self):
        acq_type = int(self.get_data('RunType')[0])
        return acq_type

    @property
    def blocks(self):
        block_types = self.get_data('BlockCond')
        return block_types


class BiomoConverter:
    """
    Opens and modifies a logfile according to BIDS-specification

    Parameters
    ----------

    output_path: path to a directory
                 Absolute or relative path to the directory where the new
                 files will be stored

    parameters: dict
                Extra information (trial duration, inter-stim interval, etc.)

    """

    def __init__(self, output_path: str, parameters: dict):
        self.out_path = output_path
        self.params = parameters

    def _create_block(self, onset: int, block_id: int = 0):
        """
        Create a block of data

        Parameters
        ----------

        block_id: int, default 0
                  Identifier for the block. It represents the type of stim
                  presented

        onset: int
               Onset time at the beginning of the block

        Returns
        -------

        data_block: list
                    Information about onset, duration and trial_type for
                    the block
        """

        data_block = []
        for trial in range(self.params['trials_per_block']):
            data_trial = [onset, self.params['trial_duration'],
                          self.params['cond_dict'][str(block_id)]]
            onset += self.params['trial_duration']

            isi = [onset, self.params['inter_stim_interval'], 'fixation']
            onset += self.params['inter_stim_interval']

            data_block.append(data_trial)
            data_block.append(isi)

        return data_block

    def get_log(self, logfile) -> pd.DataFrame:
        """
        Create a dataframe with the relevant information

        Parameters
        ----------

        logfile: Logfile object
                 Logfile instance to retrieve the relevant info

        Returns
        -------

        data_log: pd.DataFrame
                  Dataframe containing information about onset, duration and
                  trial_type in BIDS specification
        """
        onset = 0

        data = [[onset, onset + 5.0, 'initial_fix']]
        onset += 5.0

        inital_fix_block = self._create_block(onset)
        data.extend(inital_fix_block)

        onset = data[-1][0] + data[-1][1]

        for block in logfile.blocks:
            data_block = self._create_block(onset, block_id=int(block))
            data.extend(data_block)
            onset = data[-1][0] + data[-1][1]

            fixation_block = self._create_block(onset)
            data.extend(fixation_block)
            onset = data[-1][0] + data[-1][1]

        data_log = pd.DataFrame(data, columns=['onset', 'duration',
                                               'trial_type'])

        return data_log


path = os.getcwd()
input_path = os.path.join(path, '../octave_protocol/Data')
output_path = os.path.join(path, 'events_files')

if not os.path.exists(output_path):
    os.mkdir(output_path)

with open('params.json', 'r') as params:
    settings = json.load(params)

glob_str = sub_type + '-' + sub_num + '*'
glob_path = os.path.join(input_path, glob_str)
logfile_list = [Logfile(os.path.join(input_path, filename)) for
                filename in glob.glob(glob_path)]

converter = BiomoConverter(output_path, settings)

for logfile in logfile_list:
    out_df = converter.get_log(logfile)
    filename = f"{sub_type}-{sub_num}_task-BiologicalMotion" \
               f"{logfile.acq_type}_run{logfile.run}_events.tsv"
    out_df.to_csv(os.path.join(output_path, filename), sep='\t', index=False)
