import numpy as np
from sklearn.model_selection import train_test_split

#创建一个数据集X和相应的标签y,X中样本数目为100
X, y = np.arange(200).reshape((100, 2)), range(100)

#用train_test_split函数划分出训练集和测试集，测试集占比0.33
X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.33, random_state=42)

#打印出原始样本集、训练集和测试集的数目
print("The length of original data X is:", X.shape[0])
print("The length of train Data is:", X_train.shape[0])
print("The length of test Data is:", X_test.shape[0])

from pywt import wavedec, waverec
from scipy import signal
import warnings
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from mne.decoding import CSP
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import make_pipeline
import warnings
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from mne.decoding import CSP
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import make_pipeline

import moabb
from moabb.evaluations import WithinSessionEvaluation
from moabb.paradigms import LeftRightImagery
from moabb.datasets import BNCI2014004
from moabb.paradigms import LeftRightImagery
import matplotlib
import numpy as np
import pywt
from matplotlib import pyplot as plt
from mne.channels import make_standard_montage
from mne.datasets import eegbci
from mne.preprocessing import ICA
from mne_features.feature_extraction import FeatureExtractor, extract_features
from scipy.signal import welch
import matplotlib
from sklearn.decomposition import PCA, FastICA
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import KFold, StratifiedKFold, ShuffleSplit, LeaveOneGroupOut
from sklearn.pipeline import make_pipeline, Pipeline
import moabb
from moabb.datasets import PhysionetMI
from moabb.evaluations import WithinSessionEvaluation, CrossSubjectEvaluation
from moabb.paradigms import LeftRightImagery, MotorImagery
import warnings
from mne.decoding import CSP, UnsupervisedSpatialFilter, SlidingEstimator, cross_val_multiscore
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA, LinearDiscriminantAnalysis
from mne.io import read_raw_edf
import mne
import seaborn as sns
from mne import events_from_annotations, concatenate_raws, pick_types, Epochs
from sklearn.model_selection import cross_val_score
from sklearn.preprocessing import LabelEncoder, StandardScaler
from mne_features.univariate import compute_energy_freq_bands, compute_pow_freq_bands, compute_mean
from scipy.stats import differential_entropy
from sklearn.tree import DecisionTreeClassifier
from spectrum import burg, arburg
from sklearn.feature_selection import SelectKBest, f_classif, chi2, mutual_info_classif
from sklearn.pipeline import make_pipeline, Pipeline
from mne.decoding import UnsupervisedSpatialFilter, CSP
from sklearn.decomposition import PCA
from mne_features.univariate import compute_spect_entropy
from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
import matplotlib.pyplot as plt
import math
from moabb.paradigms import FilterBankLeftRightImagery
from moabb.pipelines.utils import FilterBank
import moabb
from moabb.evaluations import WithinSessionEvaluation
from moabb.paradigms import LeftRightImagery
from moabb.datasets import BNCI2014004
from moabb.paradigms import LeftRightImagery
import matplotlib
import numpy as np
import pywt
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import FunctionTransformer
from matplotlib import pyplot as plt
from mne.channels import make_standard_montage
from mne.datasets import eegbci
from mne.preprocessing import ICA
from mne_features.feature_extraction import FeatureExtractor, extract_features, FeatureFunctionTransformer
from scipy.signal import welch
import matplotlib
from sklearn.decomposition import PCA, FastICA
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import KFold, StratifiedKFold, ShuffleSplit, LeaveOneGroupOut
from sklearn.pipeline import make_pipeline, Pipeline
import moabb
from moabb.datasets import PhysionetMI
from moabb.evaluations import WithinSessionEvaluation, CrossSubjectEvaluation
from moabb.paradigms import LeftRightImagery, MotorImagery
import warnings
from mne.decoding import CSP, UnsupervisedSpatialFilter, SlidingEstimator, cross_val_multiscore
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA, LinearDiscriminantAnalysis
from mne.io import read_raw_edf
import mne
import seaborn as sns
from mne import events_from_annotations, concatenate_raws, pick_types, Epochs
from sklearn.model_selection import cross_val_score
from sklearn.preprocessing import LabelEncoder, StandardScaler
from mne_features.univariate import compute_energy_freq_bands, compute_pow_freq_bands, compute_mean
from scipy.stats import differential_entropy
from spectrum import burg, arburg
from sklearn.feature_selection import SelectKBest, f_classif, chi2, mutual_info_classif
from sklearn.pipeline import make_pipeline, Pipeline
from mne.decoding import UnsupervisedSpatialFilter, CSP
from sklearn.decomposition import PCA
from mne_features.univariate import compute_spect_entropy
from sklearn.model_selection import GridSearchCV
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
import matplotlib.pyplot as plt
import math
from sklearn.preprocessing import FunctionTransformer
from mne_features.feature_extraction import FeatureExtractor, extract_features
from moabb.paradigms import FilterBankLeftRightImagery
from moabb.pipelines.utils import FilterBank
import copy
import matplotlib.pyplot as plt
import matplotlib as mpl
from mne import time_frequency

# def csp_params(data):
#     csp_ins = CSP(n_components=8)
#     X_new = csp_ins.fit_transform(data,label)
#     return X_new
"""
bci competition IV 2b 分类结果

"""


def plot_sub_result(acc_lda_scores, acc_knn_scores, acc_svm_scores):
    barWidth = 0.3

    acc_lda_scores.append(np.mean(acc_lda_scores))
    acc_knn_scores.append(np.mean(acc_knn_scores))
    acc_svm_scores.append(np.mean(acc_svm_scores))

    r1 = np.arange(len(acc_lda_scores))
    r2 = [x + barWidth for x in r1]
    r3 = [x + barWidth for x in r2]
    plt.bar(r1, acc_lda_scores, width=barWidth, edgecolor='white', label='LDA')
    plt.bar(r2, acc_knn_scores, width=barWidth, edgecolor='white', label='KNN')
    plt.bar(r3, acc_svm_scores, width=barWidth, edgecolor='white', label='SVM')
    plt.xticks([r + barWidth for r in range(len(acc_lda_scores))],
               ['sub' + str(i) for i in range(1, len(acc_lda_scores) + 1)] + ['平均'])
    # y轴
    plt.ylabel("准确率")
    # 创建图例

    plt.title("频域分类结果对比")
    # 展示图片
    plt.show()


def ar_coeff(data):
    """
    return AR coefficients
    :param data:
    :return:
    """
    model_order = 8
    features = [arburg(data[i], order=model_order)[0].real for i in range(data.shape[0])]
    # np.zeros((data.shape[0], model_order))
    # for i in range(data.shape[0]):
    # AR, rho, ref = arburg(data[i], model_order)
    # features[i] = AR.real
    features = np.mean(features, axis=0)

    return features


def differ_en(data):
    # compute differential entropy
    freqs, psd = welch(data, 160, nperseg=320, )
    idx_alpha = np.logical_and(freqs >= 1, freqs <= 45)
    alpha_psd = psd[:, idx_alpha]
    variance = np.var(alpha_psd, ddof=1, axis=-1)  # 计算指定数据(数组元素)的方差 ,ddof=1 除以（n-1）,ddof=0代表除以n，默认值为0
    result = np.log(2 * math.pi * math.e * variance) / 2
    return result


def dwt_coeff(data):
    """
    dwt coefficients
    :param data:
    :return:
    """
    cA4, cD4, cD3, _, _ = pywt.wavedec(data, 'db4', level=4)
    cA4_mean = np.mean(cA4, axis=1)
    cA_var = np.var(cA4, axis=-1)
    cD4_mean = np.mean(cD4, axis=1)
    cD4_var = np.var(cD4, axis=-1)
    cD3_mean = np.mean(cD3, axis=1)
    cD3_var = np.var(cD3, axis=-1)
    feature = np.concatenate((cA4_mean, cA_var, cD4_mean, cD4_var, cD3_mean, cD3_var), axis=-1)
    return feature


def wavepack_coeff_energy(data):
    n = 3
    wp = pywt.WaveletPacket(data, wavelet='db4', maxlevel=n)
    re = []  # 第n层所有节点的分解系数
    for i in [node.path for node in wp.get_level(n, 'freq')]:
        re.append(wp[i].data)
    # 第n层能量特征
    energy = []
    for i in re:
        energy.append(pow(np.linalg.norm(i, axis=-1), 2))
    D1 = energy[0] / np.sum(energy)
    D2 = energy[1] / np.sum(energy)
    D3 = energy[2] / np.sum(energy)
    D1_mean = np.mean(wp['aaa'].data, axis=-1)
    D2_mean = np.mean(wp['aad'].data, axis=-1)
    D3_mean = np.mean(wp['add'].data, axis=-1)
    D1_var = np.var(wp['aaa'].data, axis=-1)
    D2_var = np.var(wp['aad'].data, axis=-1)
    D3_var = np.var(wp['add'].data, axis=-1)
    features = np.concatenate((D1, D2, D3, D1_mean, D2_mean, D3_mean, D1_var, D2_var, D3_var), axis=-1)
    return features


def dwt_coeff(data):
    """
    dwt coefficients
    :param data:
    :return:
    """
    cA4, cD4, cD3, _, _ = pywt.wavedec(data, 'db4', level=4)
    cA4_mean = np.mean(cA4, axis=-1)
    cA_var = np.var(cA4, axis=-1)
    cD4_mean = np.mean(cD4, axis=-1)
    cD4_var = np.var(cD4, axis=-1)
    cD3_mean = np.mean(cD3, axis=-1)
    cD3_var = np.var(cD3, axis=-1)
    feature = np.concatenate((cA4_mean, cA_var, cD4_mean, cD4_var, cD3_mean, cD3_var), axis=-1)
    return feature


def delta(data):
    b, a = signal.butter(8, [4 / 160 * 2, 8 / 160 * 2], btype='bandpass')
    filteredData = signal.filtfilt(b, a, data)
    return filteredData


def process_bci(data):
    # session_condition = (data["session"]=="session_2")|(data["session"]=="session_3")|(data["session"]=="session_4")
    scores = []
    for i in range(1, 109):
        conditon = (data["subject"] == str(i))
        print(data[conditon]["score"])
        score = np.max(data[conditon]["score"])
        if score >= 0.74:
            scores.append(score)
        if len(scores) == 20:
            break
    return scores


def reconstruct_dwt(data):
    coeffs = wavedec(data, 'db4', level=4)
    newdata = waverec(coeffs, 'db4')
    return newdata


def psd_welch(data):
    psd, fres = time_frequency.psd_array_welch(data, sfreq=160, fmin=0.1, fmax=60, n_fft=160, n_overlap=80)
    print(psd.shape)
    return psd


def wavelet_packet(x):
    print(x.shape)
    mother_wavelet = 'db4'
    wp = pywt.WaveletPacket(data=x, wavelet=mother_wavelet, mode='symmetric', maxlevel=4)
    node_name_list = [node.path for node in wp.get_level(4, 'natural')]
    rec_results = []
    for i in node_name_list:
        new_wp = pywt.WaveletPacket(data=x, wavelet=mother_wavelet, mode='symmetric')
        new_wp[i] = wp[i].data
        x_i = new_wp.reconstruct(update=True)
        rec_results.append(x_i)
    output = np.array(rec_results)
    output = np.reshape(output, (x.shape[0], -1))
    print(output.shape)
    return output


def reshapeIN(data):
    print('expand')
    print(data.shape)
    ## 首先将inf值都替换为nan
    data[np.isinf(data) | np.isnan(data)] = 0.01
    # newdata = np.reshape(data,(-1,3,data.shape[-1]))
    # newdata = np.reshape(data, (-1, 21, data.shape[-1]))
    newdata = np.reshape(data, (-1, 14, data.shape[-1]))
    return newdata


def reshapeDE(data):
    print("shrink")
    print(data.shape)
    ## 首先将inf值都替换为nan
    data[np.isinf(data) | np.isnan(data)] = 0.01
    # newdata = np.reshape(data,(-1,3,data.shape[-1]))
    # newdata = np.reshape(data, (-1, 21, data.shape[-1]))
    newdata = np.reshape(data, (data.shape[0], -1))
    return newdata


def dp(data):
    print(data.shape)
    b, a = signal.butter(8, [0.1 / 160 * 2, 4 / 160 * 2], btype='bandpass')
    delta_dataset = signal.filtfilt(b, a, data)
    b, a = signal.butter(8, [4 / 160 * 2, 7 / 160 * 2], btype='bandpass')
    theta_dataset = signal.filtfilt(b, a, data)
    b, a = signal.butter(8, [7 / 160 * 2, 18 / 160 * 2], btype='bandpass')
    u_dataset = signal.filtfilt(b, a, data)
    b, a = signal.butter(8, [16 / 160 * 2, 24 / 160 * 2], btype='bandpass')
    beta_Data = signal.filtfilt(b, a, data)
    b, a = signal.butter(8, [30 / 160 * 2, 50 / 160 * 2], btype='bandpass')
    gamma_Data = signal.filtfilt(b, a, data)
    newData = np.concatenate((data, u_dataset, beta_Data), axis=-1)
    print(newData.shape)
    return newData


def select_channel(data):
    selected_channel = (8, 9, 10, 11, 12, 13,)
    # selected_channel = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20)
    selected_channel = (0, 2, 4, 6, 8, 12, 14, 20, 26, 33, 49, 51, 57, 63)
    newdata = data[:, selected_channel, :]
    return newdata


if __name__ == '__main__':
    # 设置中文显示字体
    matplotlib.rc("font", family='Microsoft YaHei')
    matplotlib.rcParams['font.sans-serif'] = ['SimHei']
    # feature extraction
    timedomain_features = {'variance', 'hjorth_mobility', 'hjorth_complexity', 'skewness', 'kurtosis',
                           'zero_crossings', 'higuchi_fd', 'rms', 'ptp_amp'}

    freqdomain_features = {('DE', differ_en), 'pow_freq_bands', 'energy_freq_bands'}
    freqdomain_params = {
        'pow_freq_bands__freq_bands': np.array([0.5, 4, 8, 15, 30, 45]),
        'pow_freq_bands__ratios': 'all',
        'energy_freq_bands__freq_bands': np.array([0.5, 4, 8, 14, 30, 45]),
    }
    timefredomain_features = {('dwt_coeff', dwt_coeff)}
    timefredomain_features = {('dwt_coeff', reconstruct_dwt)}
    delta_signal = {('delta', delta)}
    dataset = PhysionetMI(imagined=True, executed=False)
    sub_list = [i for i in range(1, 109)]
    dataset.subject_list = sub_list
    paradigm = LeftRightImagery()
    print("metric", paradigm.scoring)
    acc_lda_scores = []
    acc_knn_scores = []
    acc_svm_scores = []
    evaluation = WithinSessionEvaluation(
        paradigm=paradigm,
        datasets=[dataset],
        overwrite=True,
        hdf5_path=None,
    )
    expand = FunctionTransformer(reshapeIN)
    shrink = FunctionTransformer(reshapeDE)
    downsample = FunctionTransformer(dp)
    channel_selection = FunctionTransformer(select_channel)
    myfeature = {('psd', psd_welch)}
    fe = FeatureExtractor(160, selected_funcs=myfeature)
    pipeline = make_pipeline(channel_selection, downsample, CSP(n_components=8),
                             LDA())
    results_t = evaluation.process({"timedomain+lda": pipeline})
    print(results_t)

    # fe = FeatureExtractor(160, selected_funcs=freqdomain_features, params=freqdomain_params)
    pipeline = make_pipeline(channel_selection, downsample, CSP(n_components=8),
                             KNeighborsClassifier())
    results_k = evaluation.process({"freqdomain+lda": pipeline})

    # fe = FeatureExtractor(160, selected_funcs=timefredomain_features)
    pipeline = make_pipeline(channel_selection, downsample, CSP(n_components=8),
                             SVC())
    results_s = evaluation.process({"timefreq+lda": pipeline})

# delta_signal = {('delta', delta)}
# fe = FeatureExtractor(250, selected_funcs=delta_signal)
# pipeline = make_pipeline(CSP(n_components=8), SVC())
# results_csp = evaluation.process({"CSP+lda": pipeline})

l1 = process_bci(results_t)
l2 = process_bci(results_k)
l3 = process_bci(results_s)
# l4 = process_bci(results_csp)
l1 = np.array(l1)
l2 = np.array(l2)
l3 = np.array(l3)
# l4 = np.array(l4)

l1 = np.append(l1, np.mean(l1))
l2 = np.append(l2, np.mean(l3))
l3 = np.append(l3, np.mean(l3))
# l4 = np.append(l4, np.mean(l4))
barWidth = 0.3
r1 = np.arange(len(l1))
r2 = [x + barWidth for x in r1]
r3 = [x + barWidth for x in r2]
# r4 = [x + barWidth for x in r3]
plt.tight_layout()
plt.bar(r1, l1, width=barWidth, edgecolor='white', label='LDA')
plt.bar(r2, l2, width=barWidth, edgecolor='white', label='KNN')
plt.bar(r3, l3, width=barWidth, edgecolor='white', label='SVC')
# plt.bar(r4, l4, width=barWidth, edgecolor='white', label='CSP')
plt.xticks([r + barWidth for r in range(len(l1))], ['s' + str(i) for i in range(1, len(l1))] + ['平均'])
# y轴
plt.ylabel("准确率")
# 创建图例
plt.legend(loc="upper center")
# 标题
plt.title("TCSSF特征提取对比")
# 展示图片
plt.show()
print(l1)
print(l2)
print(l3)
# print(l4)