import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
#rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)

from dmp.dmp_cartesian import DMPs_cartesian as dmp
from dmp.rotation_matrix import roto_dilatation

num_traj = 50
traj_set = []
t_set = []

# The trajectories will be generated by numerically integrating a dynamical
# system

def fun(x):
    x1 = x[0]
    x2 = x[1]
    f1 = x1 ** 3 + x2 ** 2 * x1 - x1 - x2
    f2 = x2 ** 3 + x1 ** 2 * x2 + x1 - x2
    return np.array([f1, f2])

def RK4(x0, m, tf):
    X = np.zeros([m, len(x0)])
    X[0] = x0.copy()
    dt = tf / (m - 1)
    x = x0
    for n in range(m - 1):
        K1 = fun(x)
        K2 = fun(x + dt * K1 / 2.)
        K3 = fun(x + dt * K2 / 2.)
        K4 = fun(x + dt * K3)
        x += dt * (K1 + 2 * K2 + 2 * K3 + K4) / 6.
        X[n+1] = x
    return X

for i in range(num_traj):
    ## Random select x0, tf, and m
    theta = np.random.rand() * 2. * np.pi
    rho = 1. - np.random.rand() / 5.
    x0 = rho * np.array([np.cos(theta), np.sin(theta)])
    tf = 6. + 2. * (np.random.rand() - 0.5)
    m = int (500 + np.floor(500 * np.random.rand()))
    t_set.append(np.linspace(0, tf, m))
    # Execute the trajectory
    X = RK4(x0, m, tf)
    traj_set.append(X.copy())
    # Plot
    plt.figure(1)
    plt.plot(X[:, 0], X[:, 1], '-b', lw = 0.5)
    plt.axis('equal')
    # Plot after translation and roto dilatation
    Z = X - X[0]
    old_pos = Z[-1]
    R = roto_dilatation(old_pos, np.array([1,1]))
    Z = np.dot(Z, R.transpose())
    plt.figure(2)
    plt.plot(Z[:, 0], Z[:, 1], '-b', lw = 0.5)
    plt.axis('equal')

MP = dmp(n_dmps = 2, n_bfs = 50, K = 1000, alpha_s = 4.,rescale = 'rotodilatation', T = 2.)
MP.paths_regression(traj_set, t_set)
x_track, _, _, _ = MP.rollout()
plt.figure(2)
plt.plot(x_track[:, 0], x_track[:, 1], '-r', lw = 2)
plt.xlabel(r'$x_1$')
plt.ylabel(r'$x_2$')
plt.plot(0, 0, '.k', markersize = 10)
plt.plot(1, 1, '*k', markersize = 10)
plt.title('Scaled reference frame')
# Plot from an arbitrary position
theta = 2 * np.pi * np.random.rand()
rho = 1. - np.random.rand() / 5.
MP.x_0 = rho * np.array([np.cos(theta), np.sin(theta)])
MP.x_goal = np.zeros(2)
x_track, _, _, _ = MP.rollout()
plt.figure(1)
plt.plot(x_track[:, 0], x_track[:, 1], '-r', lw = 2)
plt.xlabel(r'$x_1$')
plt.ylabel(r'$x_2$')
plt.title('Unscaled reference frame')

# With noise
var = 0.0001
for i in range(num_traj):
    traj_set[i] += np.random.randn(np.shape(traj_set[i])[0], np.shape(traj_set[i])[1]) * np.sqrt(var)

MP.paths_regression(traj_set, t_set)
MP_single = dmp(n_dmps = 2, n_bfs = 50, K = 1000, alpha_s = 4.,rescale = 'rotodilatation', T = 2.)
MP_single.imitate_path(x_des = traj_set[0], t_des = t_set[0])
x_track, dx_track, ddx_track, _ = MP.rollout()
MP_single.x_0 = np.zeros(2)
MP_single.x_goal = np.ones(2)
x_track_single, _, _, _ = MP_single.rollout()
plt.figure(3)
plt.plot(x_track[:,0], x_track[:,1], '-r')
plt.plot(x_track_single[:,0], x_track_single[:,1], '--g')
plt.plot(x_track[0][0], x_track[0][1], '.k', markersize = 10)
plt.plot(x_track[-1][0], x_track[-1][1], '*k', markersize = 10)
plt.xlabel(r'$x_1$', fontsize = 14)
plt.ylabel(r'$x_2$', fontsize = 14)

plt.show()