import os.path as op
import numpy as np
import matplotlib.pyplot as plt

from sklearn.svm import SVR
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

from mne.io.meas_info import create_info
from mne.epochs import EpochsArray
from mne.decoding import GeneralizationAcrossTime


# GENERATE SYNTHETIC DATA
n_trial = 198
n_chan = 100
n_time = 4
angles = np.linspace(15, 165, 6)

# generate template topography for each angle
X0 = np.linspace(0, 2, np.sqrt(n_chan)) - 1
topos = list()
fig, axs = plt.subplots(len(angles), sharex=False)
for a, angle in enumerate(angles):
    Xm, Ym = np.meshgrid(X0, X0)
    Xm += np.cos(np.deg2rad(2 * angle))
    Ym += np.sin(np.deg2rad(2 * angle))
    topos.append(np.exp(-((Xm ** 2) + (Ym ** 2))))
    axs[a].matshow(topos[a], aspect='equal')
    axs[a].set_axis_off()
    axs[a].set_title('topo generated by angle %s' % angle)
plt.show()

# add noisy topo to each trial at time=1, 90° shift at time=2
snr = 10
data = np.random.randn(n_trial, n_chan, n_time) / snr
y = np.arange(n_trial) % len(angles)
trial_angles = y * 30 + 15
for trial in range(n_trial):
    data[trial, :, 1] += topos[y[trial]].flatten()
    data[trial, :, 2] += topos[y[(trial + len(angles)/2) % len(angles)]].flatten()

# export in mne structure
time = range(n_time)
chan_names = ['meg' + str(i) for i in range(n_chan)]
chan_types = ['grad'] * n_chan
info = create_info(chan_names, 1, chan_types)
events = np.c_[np.cumsum(np.ones(n_trial)), np.zeros(n_trial), np.zeros(n_trial)]
epochs = EpochsArray(data, info, events)


# RUN GAT
angle2circle = lambda angles: np.deg2rad(2 * (angles + 7.5))
circle2angle = lambda angles: np.rad2deg(2 * angles) / 2 - 7.5

cos = lambda angles: np.cos(angle2circle(angles))
sin = lambda angles: np.sin(angle2circle(angles))

scaler = StandardScaler()
svr = SVR(C=1, kernel='linear')
clf = Pipeline([('scaler', scaler), ('svr', svr)])

gats = list()
for transform in [cos, sin]:
    gat = GeneralizationAcrossTime(n_jobs=1, clf=clf)
    gat.fit(epochs, y=transform(trial_angles))
    gat.score(epochs, y=transform(trial_angles))
    gats.append(gat)


# ALIGN ANGLES
cart2pol = lambda x, y: np.arctan2(y, x)
pi = np.pi

gatx = gats[0]
gaty = gats[1]

# get true angle
true_x = gatx.y_train_
true_y = gaty.y_train_
true_angle = cart2pol(true_x, true_y)
# true_angle = (angle2circle(trial_angles) + pi) % (2 * pi) - pi
true_angle = angle2circle(trial_angles)

# get x and y regressors (cos and sin)
predict_x = np.squeeze(gatx.y_pred_)
predict_y = np.squeeze(gaty.y_pred_)
predict_angles = cart2pol(predict_x, predict_y)

# compute angle error
true_angles = np.tile(true_angle, (n_time, n_time, 1))
angle_errors = (predict_angles - true_angles + pi) % (2 * pi) - pi
plt.matshow(np.mean(angle_errors ** 2, axis=2), origin='lower')
plt.colorbar()
plt.show()

fig, ax = plt.subplots(len(angles))
for a in range(len(angles)):
    error = (predict_angles[0, 0, trial_angles==angles[a]] -
             angle2circle(angles[a]) + pi) % (2 * pi) - pi
    ax[a].hist(error, np.linspace(-pi, pi, 100))
plt.show()
