import os
import re
from functools import partial
from multiprocessing import Manager, Pool, Process, Queue
from typing import Callable, Union

import scipy.io as scio
from torcheeg.io import EEGSignalIO, MetaInfoIO
from tqdm import tqdm

MAX_QUEUE_SIZE = 1024


def transform_producer(file_path: str, chunk_size: int, overlap: int,
                       num_channel: int, before_trial: Union[None, Callable],
                       transform: Union[None,
                                        Callable], after_trial: Union[Callable,
                                                                      None],
                       write_info_fn: Callable, queue: Queue):
    session_id = os.path.basename(os.path.dirname(file_path))
    _, file_name = os.path.split(file_path)

    subject = int(
        os.path.basename(file_name).split('.')[0].split('_')[0])  # subject (15)
    date = int(
        os.path.basename(file_name).split('.')[0].split('_')[1])  # period (3)

    samples = scio.loadmat(file_path, verify_compressed_data_integrity=False
                           )  # trial (15), channel(62), timestep(n*200)
    # label file
    labels = [[
        1, 2, 3, 0, 2, 0, 0, 1, 0, 1, 2, 1, 1, 1, 2, 3, 2, 2, 3, 3, 0, 3, 0, 3
    ], [
        2, 1, 3, 0, 0, 2, 0, 2, 3, 3, 2, 3, 2, 0, 1, 1, 2, 1, 0, 3, 0, 1, 3, 1
    ], [
        1, 2, 2, 1, 3, 3, 3, 1, 1, 2, 1, 0, 2, 3, 3, 0, 2, 3, 0, 0, 2, 0, 1, 0
    ]]  # The labels with 0, 1, 2, and 3 denote the ground truth, neutral, sad, fear, and happy emotions, respectively.
    session_labels = labels[int(session_id) - 1]

    trial_name_ids = [(trial_name,
                       int(re.findall(r".*_eeg(\d+)", trial_name)[0]))
                      for trial_name in samples.keys() if 'eeg' in trial_name]

    write_pointer = 0
    # loop for each trial
    for trial_name, trial_id in trial_name_ids:
        trial_samples = samples[trial_name]  # channel(62), timestep(n*200)
        if before_trial:
            trial_samples = before_trial(trial_samples)

        # record the common meta info
        trial_meta_info = {
            'subject_id': subject,
            'trial_id': trial_id,
            'session_id': session_id,
            'emotion': int(session_labels[trial_id - 1]),
            'date': date
        }

        # extract experimental signals
        start_at = 0
        if chunk_size <= 0:
            chunk_size = trial_samples.shape[1] - start_at

        # chunk with chunk size
        end_at = chunk_size
        # calculate moving step
        step = chunk_size - overlap

        trial_queue = []
        while end_at <= trial_samples.shape[1]:
            clip_sample = trial_samples[:num_channel, start_at:end_at]

            t_eeg = clip_sample
            if not transform is None:
                t_eeg = transform(eeg=clip_sample)['eeg']

            clip_id = f'{file_name}_{write_pointer}'
            if after_trial:
                trial_queue.append({'eeg': t_eeg, 'key': clip_id})
            else:
                queue.put({'eeg': t_eeg, 'key': clip_id})
            write_pointer += 1

            # record meta info for each signal
            record_info = {
                'start_at': start_at,
                'end_at': end_at,
                'clip_id': clip_id
            }
            record_info.update(trial_meta_info)
            write_info_fn(record_info)

            start_at = start_at + step
            end_at = start_at + chunk_size

        if len(trial_queue) and after_trial:
            trial_queue = after_trial(trial_queue)
            for obj in trial_queue:
                assert 'eeg' in obj and 'key' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg` and `key` as keys.'
                queue.put(obj)


def io_consumer(write_eeg_fn, queue):
    while True:
        item = queue.get()
        if not item is None:
            eeg = item['eeg']
            key = item['key']
            write_eeg_fn(eeg, key)
        else:
            break


class SingleProcessingQueue:
    def __init__(self, write_eeg_fn):
        self.write_eeg_fn = write_eeg_fn

    def put(self, item):
        eeg = item['eeg']
        key = item['key']
        self.write_eeg_fn(eeg, key)


def seed_iv_constructor(root_path: str = './eeg_raw_data',
                        chunk_size: int = 800,
                        overlap: int = 0,
                        num_channel: int = 62,
                        before_trial: Union[None, Callable] = None,
                        transform: Union[None, Callable] = None,
                        after_trial: Union[Callable, None] = None,
                        io_path: str = './io/seed_iv',
                        num_worker: int = 0,
                        verbose: bool = True,
                        cache_size: int = 64 * 1024 * 1024 * 1024) -> None:
    # init IO
    meta_info_io_path = os.path.join(io_path, 'info.csv')
    eeg_signal_io_path = os.path.join(io_path, 'eeg')

    if os.path.exists(io_path) and not os.path.getsize(meta_info_io_path) == 0:
        print(
            f'The target folder already exists, if you need to regenerate the database IO, please delete the path {io_path}.'
        )
        return

    os.makedirs(io_path, exist_ok=True)

    meta_info_io_path = os.path.join(io_path, 'info.csv')
    eeg_signal_io_path = os.path.join(io_path, 'eeg')

    info_io = MetaInfoIO(meta_info_io_path)
    eeg_io = EEGSignalIO(eeg_signal_io_path, cache_size=cache_size)

    # loop to access the dataset files
    session_list = ['1', '2', '3']
    file_path_list = []
    for session in session_list:
        session_root_path = os.path.join(root_path, session)
        for file_name in os.listdir(session_root_path):
            file_path_list.append(os.path.join(session_root_path, file_name))

    if verbose:
        # show process bar
        pbar = tqdm(total=len(file_path_list))
        pbar.set_description("[SEED-IV]")

    if num_worker > 1:
        manager = Manager()
        queue = manager.Queue(maxsize=MAX_QUEUE_SIZE)
        io_consumer_process = Process(target=io_consumer,
                                      args=(eeg_io.write_eeg, queue),
                                      daemon=True)
        io_consumer_process.start()

        partial_mp_fn = partial(transform_producer,
                                chunk_size=chunk_size,
                                overlap=overlap,
                                num_channel=num_channel,
                                before_trial=before_trial,
                                transform=transform,
                                after_trial=after_trial,
                                write_info_fn=info_io.write_info,
                                queue=queue)

        for _ in Pool(num_worker).imap(partial_mp_fn, file_path_list):
            if verbose:
                pbar.update(1)

        queue.put(None)

        io_consumer_process.join()
        io_consumer_process.close()

    else:
        for file_path in file_path_list:
            transform_producer(file_path=file_path,
                               chunk_size=chunk_size,
                               overlap=overlap,
                               num_channel=num_channel,
                               before_trial=before_trial,
                               transform=transform,
                               after_trial=after_trial,
                               write_info_fn=info_io.write_info,
                               queue=SingleProcessingQueue(eeg_io.write_eeg))
            if verbose:
                pbar.update(1)

    if verbose:
        pbar.close()
        print('Please wait for the writing process to complete...')
