import os.path as op

import numpy as np
import pandas as pd
from sklearn.pipeline import make_pipeline
from sklearn.linear_model import RidgeCV
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import KFold, cross_val_score
import mne
from pyriemann.tangentspace import TangentSpace

import config_drago as cfg

meg = 'mag'
scale = 1e22
rank = 65
reg = 1e-6
seed = 42
n_jobs = 10
cv = KFold(n_splits=n_jobs, shuffle=True, random_state=seed)


def proj_covs_common(covs, picks, scale=scale, rank=rank, reg=reg):
    covs = [d['covs'][:, picks][:, :, picks] for d in covs if 'subject' in d]
    covs = scale * np.array(covs)
    n_sub, n_fb, n_ch, n_ch = covs.shape

    # covs2 = covs.reshape(n_sub*n_fb, n_ch, n_ch)
    # covs_avg = np.mean(covs2, axis=0)
    covs_avg = covs.mean(axis=1).mean(axis=0)
    d, V = np.linalg.eigh(covs_avg)
    d = d[::-1]
    V = V[:, ::-1]
    proj_mat = V[:, :rank].T

    covs_proj = np.zeros((n_sub, n_fb, rank, rank))
    for sub in range(n_sub):
        for fb in range(n_fb):
            covs_proj[sub, fb] = proj_mat @ covs[sub, fb] @ proj_mat.T
            covs_proj[sub, fb] += reg * np.eye(rank)
    return covs_proj


def proj_covs_ts(covs):
    n_sub, n_fb, p, _ = covs.shape
    covs_ts = np.zeros((n_sub, n_fb, (p*(p+1))//2))
    for fb in range(n_fb):
        covs_ts[:, fb, :] = TangentSpace(metric="wasserstein").fit(
                covs[:, fb, :, :]).transform(covs[:, fb, :, :])
    return covs_ts


file_covs = op.join(cfg.path_outputs, 'covs_allch_oas.float32.h5')
covs_allch = mne.externals.h5io.read_hdf5(file_covs)  # (sub, fb, ch, ch)

info = np.load(op.join(cfg.path_data, 'info_allch.npy')).item()
picks = mne.pick_types(info, meg=meg)

covs = proj_covs_common(covs_allch, picks, scale=scale, rank=rank, reg=reg)
X = proj_covs_ts(covs)
X = X.reshape(len(X), -1)

info = pd.read_csv(op.join(cfg.path_data, 'participants.csv'))
subjects = [d['subject'] for d in covs_allch if 'subject' in d]
y = info.set_index('Observations').age.loc[subjects]

ridge = make_pipeline(StandardScaler(),
                      RidgeCV(alphas=np.logspace(-3, 5, 100)))
score = - cross_val_score(ridge, X, y, cv=cv,
                          scoring="neg_mean_absolute_error", n_jobs=n_jobs,
                          verbose=True)

[((20, 5, 20, 60), 'sklearn.model_selection.KFold', 'KFold', (), '', False, 'from sklearn.model_selection import KFold, cross_val_score\n'), ((53, 12, 53, 66), 'os.path.join', 'op.join', ({(53, 20, 53, 36): 'cfg.path_outputs', (53, 38, 53, 65): '"""covs_allch_oas.float32.h5"""'}, {}), "(cfg.path_outputs, 'covs_allch_oas.float32.h5')", True, 'import os.path as op\n'), ((54, 13, 54, 52), 'mne.externals.h5io.read_hdf5', 'mne.externals.h5io.read_hdf5', ({(54, 42, 54, 51): 'file_covs'}, {}), '(file_covs)', False, 'import mne\n'), ((57, 8, 57, 37), 'mne.pick_types', 'mne.pick_types', (), '', False, 'import mne\n'), ((31, 11, 31, 35), 'numpy.linalg.eigh', 'np.linalg.eigh', ({(31, 26, 31, 34): 'covs_avg'}, {}), '(covs_avg)', True, 'import numpy as np\n'), ((36, 16, 36, 51), 'numpy.zeros', 'np.zeros', ({(36, 25, 36, 50): '(n_sub, n_fb, rank, rank)'}, {}), '((n_sub, n_fb, rank, rank))', True, 'import numpy as np\n'), ((46, 14, 46, 51), 'numpy.zeros', 'np.zeros', ({(46, 23, 46, 50): '(n_sub, n_fb, p * (p + 1) // 2)'}, {}), '((n_sub, n_fb, p * (p + 1) // 2))', True, 'import numpy as np\n'), ((63, 19, 63, 61), 'os.path.join', 'op.join', ({(63, 27, 63, 40): 'cfg.path_data', (63, 42, 63, 60): '"""participants.csv"""'}, {}), "(cfg.path_data, 'participants.csv')", True, 'import os.path as op\n'), ((67, 22, 67, 38), 'sklearn.preprocessing.StandardScaler', 'StandardScaler', ({}, {}), '()', False, 'from sklearn.preprocessing import StandardScaler\n'), ((69, 10, 71, 39), 'sklearn.model_selection.cross_val_score', 'cross_val_score', (), '', False, 'from sklearn.model_selection import KFold, cross_val_score\n'), ((25, 19, 25, 33), 'numpy.array', 'np.array', ({(25, 28, 25, 32): 'covs'}, {}), '(covs)', True, 'import numpy as np\n'), ((56, 15, 56, 55), 'os.path.join', 'op.join', ({(56, 23, 56, 36): 'cfg.path_data', (56, 38, 56, 54): '"""info_allch.npy"""'}, {}), "(cfg.path_data, 'info_allch.npy')", True, 'import os.path as op\n'), ((68, 37, 68, 60), 'numpy.logspace', 'np.logspace', ({(68, 49, 68, 51): '-3', (68, 53, 68, 54): '5', (68, 56, 68, 59): '100'}, {}), '(-3, 5, 100)', True, 'import numpy as np\n'), ((40, 40, 40, 52), 'numpy.eye', 'np.eye', ({(40, 47, 40, 51): 'rank'}, {}), '(rank)', True, 'import numpy as np\n'), ((48, 28, 48, 62), 'pyriemann.tangentspace.TangentSpace', 'TangentSpace', (), '', False, 'from pyriemann.tangentspace import TangentSpace\n')]