import json
import os
import re
import sys
import tempfile
from pathlib import Path

import mne
import mne_bids
import numpy as np
import pandas as pd
import s3fs
from bids import BIDSLayout
from joblib import Parallel, delayed
from mne._fiff.utils import _find_channels, _read_segments_file
from mne.io import BaseRaw
from mne_bids import (
    BIDSPath,
)

from braindecode.datasets import BaseDataset


class EEGDashBaseDataset(BaseDataset):
    """Returns samples from an mne.io.Raw object along with a target.

    Dataset which serves samples from an mne.io.Raw object along with a target.
    The target is unique for the dataset, and is obtained through the
    `description` attribute.

    Parameters
    ----------
    raw : mne.io.Raw
        Continuous data.
    description : dict | pandas.Series | None
        Holds additional description about the continuous signal / subject.
    target_name : str | tuple | None
        Name(s) of the index in `description` that should be used to provide the
        target (e.g., to be used in a prediction task later on).
    transform : callable | None
        On-the-fly transform applied to the example before it is returned.
    """

    AWS_BUCKET = "s3://openneuro.org"

    def __init__(self, record, cache_dir, **kwargs):
        super().__init__(None, **kwargs)
        self.record = record
        self.cache_dir = Path(cache_dir)
        bids_kwargs = self.get_raw_bids_args()
        self.bidspath = BIDSPath(
            root=self.cache_dir / record["dataset"],
            datatype="eeg",
            suffix="eeg",
            **bids_kwargs,
        )
        self.s3file = self.get_s3path(record["bidspath"])
        self.filecache = self.cache_dir / record["bidspath"]
        self.bids_dependencies = record["bidsdependencies"]
        self._raw = None
        # if os.path.exists(self.filecache):
        #     self.raw = mne_bids.read_raw_bids(self.bidspath, verbose=False)

    def get_s3path(self, filepath):
        return f"{self.AWS_BUCKET}/{filepath}"

    def _download_s3(self):
        self.filecache.parent.mkdir(parents=True, exist_ok=True)
        filesystem = s3fs.S3FileSystem(
            anon=True, client_kwargs={"region_name": "us-east-2"}
        )
        filesystem.download(self.s3file, self.filecache)
        self.filenames = [self.filecache]

    def _download_dependencies(self):
        filesystem = s3fs.S3FileSystem(
            anon=True, client_kwargs={"region_name": "us-east-2"}
        )
        for dep in self.bids_dependencies:
            s3path = self.get_s3path(dep)
            filepath = self.cache_dir / dep
            if not filepath.exists():
                filepath.parent.mkdir(parents=True, exist_ok=True)
                filesystem.download(s3path, filepath)

    def get_raw_bids_args(self):
        desired_fields = ["subject", "session", "task", "run"]
        return {k: self.record[k] for k in desired_fields if self.record[k]}

    def check_and_get_raw(self):
        if not os.path.exists(self.filecache):  # not preload
            if self.bids_dependencies:
                self._download_dependencies()
            self._download_s3()
        if self._raw is None:
            self._raw = mne_bids.read_raw_bids(self.bidspath, verbose=False)

    def __getitem__(self, index):
        # self.check_and_get_raw()

        X = self.raw[:, index][0]
        y = None
        if self.target_name is not None:
            y = self.description[self.target_name]
        if isinstance(y, pd.Series):
            y = y.to_list()
        if self.transform is not None:
            X = self.transform(X)
        return X, y

    def __len__(self):
        if self._raw is None:
            return int(self.record["ntimes"] * self.record["sampling_frequency"])
        else:
            return len(self._raw)

    @property
    def raw(self):
        if self._raw is None:
            self.check_and_get_raw()
        return self._raw

    @raw.setter
    def raw(self, raw):
        self._raw = raw


class EEGDashBaseRaw(BaseRaw):
    r"""MNE Raw object from EEG-Dash connection with Openneuro S3 file.

    Parameters
    ----------
    input_fname : path-like
        Path to the S3 file
    eog : list | tuple | 'auto'
        Names or indices of channels that should be designated EOG channels.
        If 'auto', the channel names containing ``EOG`` or ``EYE`` are used.
        Defaults to empty tuple.
    %(preload)s
        Note that preload=False will be effective only if the data is stored
        in a separate binary file.
    %(uint16_codec)s
    %(montage_units)s
    %(verbose)s

    See Also
    --------
    mne.io.Raw : Documentation of attributes and methods.

    Notes
    -----
    .. versionadded:: 0.11.0
    """

    AWS_BUCKET = "s3://openneuro.org"

    def __init__(
        self,
        input_fname,
        metadata,
        eog=(),
        preload=False,
        *,
        cache_dir="./.eegdash_cache",
        bids_dependencies: list = [],
        uint16_codec=None,
        montage_units="auto",
        verbose=None,
    ):
        """
        Get to work with S3 endpoint first, no caching
        """
        # Create a simple RawArray
        sfreq = metadata["sfreq"]  # Sampling frequency
        n_times = metadata["n_times"]
        ch_names = metadata["ch_names"]
        ch_types = []
        for ch in metadata["ch_types"]:
            chtype = ch.lower()
            if chtype == "heog" or chtype == "veog":
                chtype = "eog"
            ch_types.append(chtype)
        info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=ch_types)
        self.s3file = self.get_s3path(input_fname)
        self.cache_dir = Path(cache_dir)
        self.filecache = self.cache_dir / input_fname
        self.bids_dependencies = bids_dependencies

        if preload and not os.path.exists(self.filecache):
            self._download_s3()
            preload = self.filecache

        super().__init__(
            info,
            preload,
            last_samps=[n_times - 1],
            orig_format="single",
            verbose=verbose,
        )

    def get_s3path(self, filepath):
        return f"{self.AWS_BUCKET}/{filepath}"

    def _download_s3(self):
        self.filecache.parent.mkdir(parents=True, exist_ok=True)
        filesystem = s3fs.S3FileSystem(
            anon=True, client_kwargs={"region_name": "us-east-2"}
        )
        filesystem.download(self.s3file, self.filecache)
        self.filenames = [self.filecache]

    def _download_dependencies(self):
        filesystem = s3fs.S3FileSystem(
            anon=True, client_kwargs={"region_name": "us-east-2"}
        )
        for dep in self.bids_dependencies:
            s3path = self.get_s3path(dep)
            filepath = self.cache_dir / dep
            if not filepath.exists():
                filepath.parent.mkdir(parents=True, exist_ok=True)
                filesystem.download(s3path, filepath)

    def _read_segment(
        self, start=0, stop=None, sel=None, data_buffer=None, *, verbose=None
    ):
        if not os.path.exists(self.filecache):  # not preload
            if self.bids_dependencies:
                self._download_dependencies()
            self._download_s3()
        else:  # not preload and file is not cached
            self.filenames = [self.filecache]
        return super()._read_segment(start, stop, sel, data_buffer, verbose=verbose)

    def _read_segment_file(self, data, idx, fi, start, stop, cals, mult):
        """Read a chunk of data from the file."""
        _read_segments_file(self, data, idx, fi, start, stop, cals, mult, dtype="<f4")


class EEGBIDSDataset:
    ALLOWED_FILE_FORMAT = ["eeglab", "brainvision", "biosemi", "european"]
    RAW_EXTENSIONS = {
        ".set": [".set", ".fdt"],  # eeglab
        ".edf": [".edf"],  # european
        ".vhdr": [".eeg", ".vhdr", ".vmrk", ".dat", ".raw"],  # brainvision
        ".bdf": [".bdf"],  # biosemi
    }
    METADATA_FILE_EXTENSIONS = [
        "eeg.json",
        "channels.tsv",
        "electrodes.tsv",
        "events.tsv",
        "events.json",
    ]

    def __init__(
        self,
        data_dir=None,  # location of bids dataset
        dataset="",  # dataset name
    ):
        if data_dir is None or not os.path.exists(data_dir):
            raise ValueError("data_dir must be specified and must exist")
        self.bidsdir = Path(data_dir)
        self.dataset = dataset
        assert str(self.bidsdir).endswith(self.dataset)
        self.layout = BIDSLayout(data_dir)

        # get all recording files in the bids directory
        self.files = self.get_recordings(self.layout)
        assert len(self.files) > 0, ValueError(
            "Unable to construct EEG dataset. No EEG recordings found."
        )
        assert self.check_eeg_dataset(), ValueError("Dataset is not an EEG dataset.")
        # temp_dir = (Path().resolve() / 'data')
        # if not os.path.exists(temp_dir):
        #     os.mkdir(temp_dir)
        # if not os.path.exists(temp_dir / f'{dataset}_files.npy'):
        #     self.files = self.get_files_with_extension_parallel(self.bidsdir, extension=self.RAW_EXTENSION[self.raw_format])
        #     np.save(temp_dir / f'{dataset}_files.npy', self.files)
        # else:
        #     self.files = np.load(temp_dir / f'{dataset}_files.npy', allow_pickle=True)

    def check_eeg_dataset(self):
        return self.get_bids_file_attribute("modality", self.files[0]).lower() == "eeg"

    def get_recordings(self, layout: BIDSLayout):
        files = []
        for ext, exts in self.RAW_EXTENSIONS.items():
            files = layout.get(extension=ext, return_type="filename")
            if files:
                break
        return files

    def get_relative_bidspath(self, filename):
        bids_parent_dir = self.bidsdir.parent.absolute()
        return str(Path(filename).relative_to(bids_parent_dir))

    def get_property_from_filename(self, property, filename):
        import platform

        if platform.system() == "Windows":
            lookup = re.search(rf"{property}-(.*?)[_\\]", filename)
        else:
            lookup = re.search(rf"{property}-(.*?)[_\/]", filename)
        return lookup.group(1) if lookup else ""

    def merge_json_inheritance(self, json_files):
        """
        Merge list of json files found by get_bids_file_inheritance,
        expecting the order (from left to right) is from lowest level to highest level,
        and return a merged dictionary
        """
        json_files.reverse()
        json_dict = {}
        for f in json_files:
            json_dict.update(json.load(open(f)))
        return json_dict

    def get_bids_file_inheritance(self, path, basename, extension):
        """
        Get all files with given extension that applies to the basename file
        following the BIDS inheritance principle in the order of lowest level first
        @param
            basename: bids file basename without _eeg.set extension for example
            extension: e.g. channels.tsv
        """
        top_level_files = ["README", "dataset_description.json", "participants.tsv"]
        bids_files = []

        # check if path is str object
        if isinstance(path, str):
            path = Path(path)
        if not path.exists:
            raise ValueError("path {path} does not exist")

        # check if file is in current path
        for file in os.listdir(path):
            # target_file = path / f"{cur_file_basename}_{extension}"
            if os.path.isfile(path / file):
                # check if file has extension extension
                # check if file basename has extension
                if file.endswith(extension):
                    filepath = path / file
                    bids_files.append(filepath)

        # check if file is in top level directory
        if any(file in os.listdir(path) for file in top_level_files):
            return bids_files
        else:
            # call get_bids_file_inheritance recursively with parent directory
            bids_files.extend(
                self.get_bids_file_inheritance(path.parent, basename, extension)
            )
            return bids_files

    def get_bids_metadata_files(self, filepath, metadata_file_extension):
        """
        (Wrapper for self.get_bids_file_inheritance)
        Get all BIDS metadata files that are associated with the given filepath, following the BIDS inheritance principle.

        Args:
            filepath (str or Path): The filepath to get the associated metadata files for.
            metadata_files_extensions (list): A list of file extensions to search for metadata files.

        Returns:
            list: A list of filepaths for all the associated metadata files
        """
        if isinstance(filepath, str):
            filepath = Path(filepath)
        if not filepath.exists:
            raise ValueError("filepath {filepath} does not exist")
        path, filename = os.path.split(filepath)
        basename = filename[: filename.rfind("_")]
        # metadata files
        meta_files = self.get_bids_file_inheritance(
            path, basename, metadata_file_extension
        )
        return meta_files

    def scan_directory(self, directory, extension):
        result_files = []
        directory_to_ignore = [".git", ".datalad", "derivatives", "code"]
        with os.scandir(directory) as entries:
            for entry in entries:
                if entry.is_file() and entry.name.endswith(extension):
                    print("Adding ", entry.path)
                    result_files.append(entry.path)
                elif entry.is_dir():
                    # check that entry path doesn't contain any name in ignore list
                    if not any(name in entry.name for name in directory_to_ignore):
                        result_files.append(entry.path)  # Add directory to scan later
        return result_files

    def get_files_with_extension_parallel(
        self, directory, extension=".set", max_workers=-1
    ):
        result_files = []
        dirs_to_scan = [directory]

        # Use joblib.Parallel and delayed to parallelize directory scanning
        while dirs_to_scan:
            print(f"Scanning {len(dirs_to_scan)} directories...", dirs_to_scan)
            # Run the scan_directory function in parallel across directories
            results = Parallel(n_jobs=max_workers, prefer="threads", verbose=1)(
                delayed(self.scan_directory)(d, extension) for d in dirs_to_scan
            )

            # Reset the directories to scan and process the results
            dirs_to_scan = []
            for res in results:
                for path in res:
                    if os.path.isdir(path):
                        dirs_to_scan.append(path)  # Queue up subdirectories to scan
                    else:
                        result_files.append(path)  # Add files to the final result
            print(f"Current number of files: {len(result_files)}")

        return result_files

    def load_and_preprocess_raw(self, raw_file, preprocess=False):
        print(f"Loading {raw_file}")
        EEG = mne.io.read_raw_eeglab(raw_file, preload=True, verbose="error")

        if preprocess:
            # highpass filter
            EEG = EEG.filter(l_freq=0.25, h_freq=25, verbose=False)
            # remove 60Hz line noise
            EEG = EEG.notch_filter(freqs=(60), verbose=False)
            # bring to common sampling rate
            sfreq = 128
            if EEG.info["sfreq"] != sfreq:
                EEG = EEG.resample(sfreq)
            # # normalize data to zero mean and unit variance
            # scalar = preprocessing.StandardScaler()
            # mat_data = scalar.fit_transform(mat_data.T).T # scalar normalize for each feature and expects shape data x features

        mat_data = EEG.get_data()

        if len(mat_data.shape) > 2:
            raise ValueError("Expect raw data to be CxT dimension")
        return mat_data

    def get_files(self):
        return self.files

    def resolve_bids_json(self, json_files: list):
        """
        Resolve the BIDS JSON files and return a dictionary of the resolved values.
        Args:
            json_files (list): A list of JSON files to resolve in order of leaf level first

        Returns:
            dict: A dictionary of the resolved values.
        """
        if len(json_files) == 0:
            raise ValueError("No JSON files provided")
        json_files.reverse()  # TODO undeterministic

        json_dict = {}
        for json_file in json_files:
            with open(json_file) as f:
                json_dict.update(json.load(f))
        return json_dict

    def get_bids_file_attribute(self, attribute, data_filepath):
        entities = self.layout.parse_file_entities(data_filepath)
        bidsfile = self.layout.get(**entities)[0]
        attributes = bidsfile.get_entities(metadata="all")
        attribute_mapping = {
            "sfreq": "SamplingFrequency",
            "modality": "datatype",
            "task": "task",
            "session": "session",
            "run": "run",
            "subject": "subject",
            "ntimes": "RecordingDuration",
            "nchans": "EEGChannelCount",
        }
        attribute_value = attributes.get(attribute_mapping.get(attribute), None)
        return attribute_value

    def channel_labels(self, data_filepath):
        channels_tsv = pd.read_csv(
            self.get_bids_metadata_files(data_filepath, "channels.tsv")[0], sep="\t"
        )
        return channels_tsv["name"].tolist()

    def channel_types(self, data_filepath):
        channels_tsv = pd.read_csv(
            self.get_bids_metadata_files(data_filepath, "channels.tsv")[0], sep="\t"
        )
        return channels_tsv["type"].tolist()

    def num_times(self, data_filepath):
        eeg_jsons = self.get_bids_metadata_files(data_filepath, "eeg.json")
        eeg_json_dict = self.merge_json_inheritance(eeg_jsons)
        return int(
            eeg_json_dict["SamplingFrequency"] * eeg_json_dict["RecordingDuration"]
        )

    def subject_participant_tsv(self, data_filepath):
        """Get participants_tsv info of a subject based on filepath"""
        participants_tsv = pd.read_csv(
            self.get_bids_metadata_files(data_filepath, "participants.tsv")[0], sep="\t"
        )
        # if participants_tsv is not empty
        if participants_tsv.empty:
            return {}
        # set 'participant_id' as index
        participants_tsv.set_index("participant_id", inplace=True)
        subject = f"sub-{self.get_bids_file_attribute('subject', data_filepath)}"
        return participants_tsv.loc[subject].to_dict()

    def eeg_json(self, data_filepath):
        eeg_jsons = self.get_bids_metadata_files(data_filepath, "eeg.json")
        eeg_json_dict = self.merge_json_inheritance(eeg_jsons)
        return eeg_json_dict

    def channel_tsv(self, data_filepath):
        channels_tsv = pd.read_csv(
            self.get_bids_metadata_files(data_filepath, "channels.tsv")[0], sep="\t"
        )
        channel_tsv = channels_tsv.to_dict()
        # 'name' and 'type' now have a dictionary of index-value. Convert them to list
        for list_field in ["name", "type", "units"]:
            channel_tsv[list_field] = list(channel_tsv[list_field].values())
        return channel_tsv
