import time

import numpy as np
import scipy

import mne
import pyriemann
import sklearn

import tag_mne as tm
import pandas as pd

import utils


def func_proc_epochs(epochs):
    epochs = epochs.pick(picks="eeg").crop(tmin=0.25, tmax=5.0)
    return epochs


def load_epochs(files, concat=False):
    epochs_list = list()
    for file in files:
        epochs = mne.read_epochs(file)
        epochs_list.append(epochs)
    if concat:
        return tm.concatenate_epochs(epochs_list)
    else:
        return epochs_list


def recenter_cov(cov_session, scaling=False):
    mean_cov_session = pyriemann.utils.mean.mean_covariance(cov_session)
    n_covs, _, _ = cov_session.shape

    if scaling:
        d_tilda_sqrt = 0
        for m in range(n_covs):
            d_tilda_sqrt += pyriemann.utils.distance.distance(
                mean_cov_session, cov_session[m, :, :], squared=True
            )
        d_tilda = np.sqrt(d_tilda_sqrt)

    mean_cov_session_inv_sqrt = scipy.linalg.fractional_matrix_power(
        mean_cov_session, -0.5
    )

    recentered_covs = np.zeros(cov_session.shape)
    for m in range(n_covs):
        recentered_covs[m, :, :] = (
            mean_cov_session_inv_sqrt @ cov_session[m, :, :] @ mean_cov_session_inv_sqrt
        )
        if scaling:
            recentered_covs[m, :, :] = scipy.linalg.fractional_matrix_power(
                recentered_covs[m, :, :], 1 / d_tilda
            )

    return recentered_covs


def deeplearning_cross_subject(
    subjects_train,
    subjects_valid,
    subjects_test,
    func_get_fnames,
    criterion,
    optimizer,
    batch_size,
    n_epochs,
    optimizer_params=None,
    model=None,
    func_get_model=None,
    scheduler=None,
    scheduler_params=None,
    device="cpu",
    func_proc_epochs=None,
    label_keys={"event:left": 0, "event:right": 1},
    compile_test_subjects=False,
    enable_wandb_logging=False,
    wandb_params=None,
    checkpoint_fname=None,
    history_fname=None,
    early_stopping=None,
    name_classifier=None,
    desc=None,
):
    """
    subjects_train: list
    subjects_test: list
    files: callable
        subject名を引数とし，ファイル名を返す関数
    func_proc_epochs: function
        def func_proc_eochs(epochs)
            return epochs
    classifier: list of classifier, "tslr", "mdm", instance
    label_keys: dict
    compile_test_subjects: bool
        Trueにすると，テストsubjectのデータをまとめて，その精度とかを返す
        Falseにすると，各被験者ごとの精度をリストで返す

    """
    import torch

    if enable_wandb_logging:
        import wandb

    if not isinstance(subjects_train, list) or not isinstance(subjects_test, list):
        raise ValueError("type of subjects_train and subjects_test have to be list")

    # load data

    ## training data
    epochs_train = list()
    y_train = list()
    for subject in subjects_train:
        files = func_get_fnames(subject)
        epochs = load_epochs(files, True)
        if func_proc_epochs is not None:
            epochs = func_proc_epochs(epochs)
        y_train += utils.get_labels_from_epochs(epochs, label_keys)
        epochs_train.append(epochs)

    epochs_train = tm.concatenate_epochs(epochs_train)
    X_train = epochs_train.get_data()

    ## validation data
    epochs_valid = list()
    y_valid = list()
    for subject in subjects_valid:
        files = func_get_fnames(subject)
        epochs = load_epochs(files, True)
        if func_proc_epochs is not None:
            epochs = func_proc_epochs(epochs)
        y_valid += utils.get_labels_from_epochs(epochs, label_keys)
        epochs_valid.append(epochs)

    epochs_valid = tm.concatenate_epochs(epochs_valid)
    X_valid = epochs_valid.get_data()

    ## test data
    epochs_test = list()
    X_test = list()
    y_test = list()
    for subject in subjects_test:
        files = func_get_fnames(subject)
        epochs = load_epochs(files, True)
        if func_proc_epochs is not None:
            epochs = func_proc_epochs(epochs)
        y_test.append(utils.get_labels_from_epochs(epochs, label_keys))
        epochs_test.append(epochs)
        X_test.append(epochs.get_data())

    if compile_test_subjects:
        # compile y
        y_test_compiled = list()
        for y in y_test:
            y_test_compiled += y
        y_test = [y_test_compiled]

        # compile X
        X_test = [np.concatenate(X_test, axis=0)]

        # compile subjects_test list
        subjects_test = [subjects_test]

    # create dataloader

    (dataloader_train, dataloader_valid, dataloader_test) = utils.nd_to_dataloader(
        X_train,
        y_train,
        X_valid,
        y_valid,
        X_test,
        y_test,
        device=device,
        normalize=True,
        batch_size=batch_size,
    )

    # setup model

    if func_get_model is not None:
        model = func_get_model(X_train, y_train, device)

    if model is None:
        raise RuntimeError("model is None")

    if name_classifier is None:
        name_classifier = model.__class__.__name__

    # setup optimizer

    if optimizer_params is not None:
        optimizer = optimizer(
            params=model.parameters(),
            **optimizer_params,
        )
    else:
        optimizer = optimizer(params=model.parameters())

    # setup scheduler

    if scheduler is not None:
        if scheduler_params is not None:
            scheduler = scheduler(
                optimizer=optimizer,
                **scheduler_params,
            )
        else:
            scheduler = scheduler(optimizer=optimizer)

    # misc

    history = {
        "epoch": list(),
        "train_loss": list(),
        "valid_loss": list(),
        "train_acc": list(),
        "valid_acc": list(),
    }

    loss_best = [float("inf")]

    # wandb.init(project = project, name = f"within-session_sub-{subject}_run-{run}")
    if enable_wandb_logging:
        wandb.init(**wandb_params)

    tic = time.time()
    for epoch in range(n_epochs):
        valid_loss = utils.train_epoch(
            model=model,
            criterion=criterion,
            optimizer=optimizer,
            scheduler=scheduler,
            dataloader_train=dataloader_train,
            dataloader_valid=dataloader_valid,
            epoch=epoch,
            loss_best=loss_best,
            history=history,
            checkpoint_fname=checkpoint_fname,
            enable_wandb=enable_wandb_logging,
        )

        if early_stopping is not None:
            if early_stopping.step(valid_loss):
                print(f"Early stopping was triggered: epoch-{epoch+1}")
                break

    toc = time.time()
    elapsed_time = toc - tic

    if enable_wandb_logging:
        wandb.finish()

    if history_fname is not None:
        df_save = pd.DataFrame(history)
        df_save.to_pickle(f"{history_fname}.pkl")
        df_save.to_html(f"{history_fname}.html")

    print(f"Elapsed Time: {elapsed_time:.2f}s")

    # classify test data
    if checkpoint_fname is not None:
        checkpoint = torch.load(checkpoint_fname)
        model.load_state_dict(checkpoint["model_state_dict"])

    df_list = list()
    model.eval()
    with torch.no_grad():
        for dataloader, subject in zip(dataloader_test, subjects_test):
            # accuracy = utils.accuracy_score_dataloader(model, dataloader_test)

            preds, labels, logits, probas = utils.get_predictions(
                model,
                dataloader,
            )

            accuracy = sklearn.metrics.accuracy_score(labels, preds)

            df_results = pd.DataFrame()
            df_results["subjects_train"] = [subjects_train]
            df_results["subjects_valid"] = [subjects_valid]
            df_results["subjects_test"] = [subject]
            df_results["classifier"] = [name_classifier]
            df_results["accuracy"] = [accuracy]
            df_results["labels"] = [labels]
            df_results["preds"] = [preds]
            df_results["probas"] = [probas]
            df_results["logits"] = [logits]
            df_results["elapsed_time"] = [elapsed_time]
            df_results["desc"] = [desc]

            df_list.append(df_results)

    df = pd.concat(df_list, axis=0, ignore_index=True)

    return df


if __name__ == "__main__":
    debug = True
    import sys

    sys.path.append("../")
    import load

    import torch
    import braindecode

    def _func_proc_epochs(epochs):
        epochs = epochs.pick(picks="eeg").crop(tmin=0.25, tmax=5.0)
        return epochs

    def _func_get_fnames(subject):
        base_dir = (
            load.config["dir"]["deriv"]
            / "epochs"
            / "l_freq-8.0_h_freq-30.0_resample-128"
            / subject
        )

        fnames_list = list()
        fnames_list.append(base_dir / f"sub-{subject}_acquisition-epo.fif")
        fnames_list.append(base_dir / f"sub-{subject}_online-epo.fif")

        return fnames_list

    def _func_get_model(X, y, device):
        print(X.shape)
        _, n_chans, n_times = X.shape
        F1 = 4
        D = 2
        F2 = F1 * D

        model = braindecode.models.EEGNetv4(
            n_chans=n_chans,
            n_outputs=2,
            n_times=n_times,
            F1=F1,
            D=D,
            F2=F2,
            drop_prob=0.5,
        )
        model.to(device)

        return model

    # setup mpdels

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    lr = 0.001
    weight_decay = 0.0005
    n_epochs = 10
    batch_size = 64
    patience = 75

    early_stopping = utils.EarlyStopping(patience=patience)

    criterion = torch.nn.CrossEntropyLoss()
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR

    returns = deeplearning_cross_subject(
        subjects_train=["A1", "A2"],
        subjects_valid=["A3", "A4"],
        subjects_test=["A5", "A6"],
        func_get_fnames=_func_get_fnames,
        func_proc_epochs=_func_proc_epochs,
        func_get_model=_func_get_model,
        criterion=criterion,
        scheduler=scheduler,
        batch_size=batch_size,
        n_epochs=n_epochs,
        checkpoint_fname=None,
        early_stopping=early_stopping,
        device=device,
        history_fname=None,
        scheduler_params={"T_max": n_epochs, "eta_min": 1e-6},
        optimizer=torch.optim.AdamW,
        optimizer_params={"lr": lr, "weight_decay": weight_decay},
        compile_test_subjects=False,
    )
