from epilepsy_bids_loader import ManifestBIDSLoader, save_yaml
from collections import defaultdict
from copy import deepcopy
from pathlib import Path
from typing import Dict, List


def merge_run_dicts(
    dict1: Dict[str, List[str]],
    dict2: Dict[str, List[str]]
):
    new_dict = {
        key: dict1.get(key, []) + dict2.get(key, [])
        for key in sorted(list(set(dict1.keys()).union(set(dict2.keys()))))
    }
    return new_dict


def create_sub_spec_manifest(
    loader: ManifestBIDSLoader,
    save_path: Path
):
    manifest = defaultdict(lambda: # sub-xx
        defaultdict(lambda: # fold-xx
            defaultdict(lambda: # train/dev/test
                defaultdict(list) # ses-xx: [run-xx, ...]
            )
        )
    )
    bids_tree = loader.data
    events = loader.events.to_dict()["eventType"]
    for sub, sub_data in bids_tree.items():
        train_runs = defaultdict(list)
        dev_runs = defaultdict(list)
        test_runs = defaultdict(list)
        fold = 0
        sz_count = 0

        for ses, ses_data in sub_data.items():

            for run in ses_data.keys():
                # Starting condition
                cumulate_runs = (
                    train_runs if sz_count < 1
                    else dev_runs if sz_count < 2
                    else test_runs
                )
                cumulate_runs[ses].append(run)

                # Check if seizure within run
                is_sz = events[(sub, ses, run)].startswith("sz")
                if is_sz:
                    sz_count += 1

                if is_sz and sz_count > 2:
                    # Write manifest
                    fold_key = f"fold-{fold:02d}"
                    manifest[sub][fold_key]["train"] = deepcopy(train_runs)
                    manifest[sub][fold_key]["dev"] = deepcopy(dev_runs)
                    manifest[sub][fold_key]["test"] = deepcopy(test_runs)

                    # Merge train and dev
                    train_runs = merge_run_dicts(train_runs, dev_runs)
                    dev_runs = deepcopy(test_runs)
                    test_runs = defaultdict(list)

                    # Increment fold
                    fold += 1

                continue
            continue
        continue

    save_yaml(
        path=save_path,
        obj={
            "subject_specific": manifest
        }
    )

    return

def main(args):
    loader = ManifestBIDSLoader(
        data_path=args.path,
        manifest_path=args.manifest,
        enforce_chronology=False,
        allowed_subjects=None
    )
    create_sub_spec_manifest(
        loader=loader,
        save_path=Path(args.outpath)
    )
    return


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-p", "--path", type=str,
        # required=True,
        # default="data/BIDS_Siena"
        default="data/BIDS_CHB-MIT"
    )
    parser.add_argument(
        "-m", "--manifest", type=str,
        # required=True,
        # default="manifests/siena.yaml"
        default="manifests/chbmit.yaml"
    )
    parser.add_argument(
        "-o", "--outpath", type=str,
        # required=True,
        # default="manifests/_siena.yaml"
        default="manifests/_chbmit.yaml"
    )
    args = parser.parse_args()
    main(args)