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)

[((414, 469), 'sklearn.model_selection.KFold', 'KFold', ([], {'n_splits': 'n_jobs', 'shuffle': '(True)', 'random_state': 'seed'}), '(n_splits=n_jobs, shuffle=True, random_state=seed)\n', (419, 469), False, 'from sklearn.model_selection import KFold, cross_val_score\n'), ((1494, 1548), 'os.path.join', 'op.join', (['cfg.path_outputs', '"""covs_allch_oas.float32.h5"""'], {}), "(cfg.path_outputs, 'covs_allch_oas.float32.h5')\n", (1501, 1548), True, 'import os.path as op\n'), ((1562, 1601), 'mne.externals.h5io.read_hdf5', 'mne.externals.h5io.read_hdf5', (['file_covs'], {}), '(file_covs)\n', (1590, 1601), False, 'import mne\n'), ((1696, 1725), 'mne.pick_types', 'mne.pick_types', (['info'], {'meg': 'meg'}), '(info, meg=meg)\n', (1710, 1725), False, 'import mne\n'), ((842, 866), 'numpy.linalg.eigh', 'np.linalg.eigh', (['covs_avg'], {}), '(covs_avg)\n', (856, 866), True, 'import numpy as np\n'), ((948, 983), 'numpy.zeros', 'np.zeros', (['(n_sub, n_fb, rank, rank)'], {}), '((n_sub, n_fb, rank, rank))\n', (956, 983), True, 'import numpy as np\n'), ((1264, 1305), 'numpy.zeros', 'np.zeros', (['(n_sub, n_fb, p * (p + 1) // 2)'], {}), '((n_sub, n_fb, p * (p + 1) // 2))\n', (1272, 1305), True, 'import numpy as np\n'), ((1872, 1914), 'os.path.join', 'op.join', (['cfg.path_data', '"""participants.csv"""'], {}), "(cfg.path_data, 'participants.csv')\n", (1879, 1914), True, 'import os.path as op\n'), ((2056, 2072), 'sklearn.preprocessing.StandardScaler', 'StandardScaler', ([], {}), '()\n', (2070, 2072), False, 'from sklearn.preprocessing import StandardScaler\n'), ((2147, 2250), 'sklearn.model_selection.cross_val_score', 'cross_val_score', (['ridge', 'X', 'y'], {'cv': 'cv', 'scoring': '"""neg_mean_absolute_error"""', 'n_jobs': 'n_jobs', 'verbose': '(True)'}), "(ridge, X, y, cv=cv, scoring='neg_mean_absolute_error',\n    n_jobs=n_jobs, verbose=True)\n", (2162, 2250), False, 'from sklearn.model_selection import KFold, cross_val_score\n'), ((637, 651), 'numpy.array', 'np.array', (['covs'], {}), '(covs)\n', (645, 651), True, 'import numpy as np\n'), ((1639, 1679), 'os.path.join', 'op.join', (['cfg.path_data', '"""info_allch.npy"""'], {}), "(cfg.path_data, 'info_allch.npy')\n", (1646, 1679), True, 'import os.path as op\n'), ((2111, 2134), 'numpy.logspace', 'np.logspace', (['(-3)', '(5)', '(100)'], {}), '(-3, 5, 100)\n', (2122, 2134), True, 'import numpy as np\n'), ((1155, 1167), 'numpy.eye', 'np.eye', (['rank'], {}), '(rank)\n', (1161, 1167), True, 'import numpy as np\n'), ((1357, 1391), 'pyriemann.tangentspace.TangentSpace', 'TangentSpace', ([], {'metric': '"""wasserstein"""'}), "(metric='wasserstein')\n", (1369, 1391), False, 'from pyriemann.tangentspace import TangentSpace\n')]