#%%

import numpy as np

class DynSys:
    def __init__(self, order, dt=0.01, dim=2):
        self.order = order
        self.dim = dim
        self.dt = dt

        ## construct system matrices
        # continuous single
        A = np.zeros((order,order))
        A[0:order-1, 1:order] = np.eye(order-1)
        B = np.zeros((order,1))
        B[-1,:] = 1.

        # Euler discretization
        A = np.eye(A.shape[0]) + A * dt
        B = B * dt

        # Multiple output
        # Add dimensions
        A = np.kron(A, np.eye(dim))
        B = np.kron(B, np.eye(dim) )

        # Sensor matrix (assuming only position is observed) 
        # Actually unused here and directly setting unobserved Q entries to 0 weight
        C = np.kron( np.hstack([np.ones((1,1)), np.zeros((1, order-1))]),
                    np.eye(dim))
        
        self.A, self.B, self.C = A, B, C
        
    def y(self, x):
        ''' y = Cx'''
        return np.dot(self.C, x)
    
    def num_timesteps(self, duration):
        return int(duration/self.dt)

    def matrices(self):
        return self.A, self.B

sys = DynSys(2)
latex(sys.A,'\\bm{A}'), latex(sys.B, '\\bm{B}'), latex(sys.C,'\\bm{C}') # show latex

#%%

def stepwise(m, n):
    T = range(n)
    I = np.linspace(0, float(m)-0.1, n).astype(int)
    return I, T

# For 30 time steps and 3 via points:
Ts = np.ones(30).astype(int)
I, T = stepwise(3, 30)
Ts[T] = I

plt.figure(figsize=(14,0.5))
plt.gca().pcolormesh([[ t if t >= 0 else '.' for t in Ts]], cmap='spring', edgecolor='k')
for i, t in enumerate(Ts):
    plt.text(0.4 + i, 0.5,str(t))
plt.axis('off')
plt.show()

def make_reference(Mu, Sigma, n, sys, reference=stepwise, end_weight = 1e10):
    ''' Make reference and state cost matrices'''
    order, dim = sys.order, sys.dim
    cDim = order*dim
    muDim = Mu.shape[0]
    m = Mu.shape[1]
    
    # precision matrices
    Lambda = np.zeros((muDim, m*muDim))
    for i in range(m):
        Lambda[:,i*muDim:muDim*(i+1)] = np.linalg.inv(Sigma[:,:,i])

    Q = np.zeros((cDim*n, cDim*n)) # Precision matrix
    MuQ = np.zeros((cDim, n))

    I, T = reference(m, n) # Get reference indices
    for i, t in zip(I, T):
        Q[ t*cDim:t*cDim+muDim, t*cDim:t*cDim+muDim] = Lambda[:,i*muDim:(i+1)*muDim] 
        MuQ[:muDim, t] = Mu[:,i]

    # end indices (forces movement to a stop with higher order systems)
    ind = T[-1]
    if end_weight > 0.0: # and muDim <= dim:
        for i in range(muDim, cDim):
            Q[ind*cDim+i, ind*cDim+i] = end_weight
    return MuQ, Q

#%%

import matplotlib.pyplot as plt

def load_gauss(path, scale=1.): 
    Mu, Sigma = load_pkl(path)
    Mu *= scale
    Sigma *= scale**2
    return Mu, Sigma

Mu, Sigma = load_gauss('./m_1.pkl', scale=0.01) 

# Plot it
plt.figure()
plot_gauss(Mu, Sigma)
plt.plot(Mu[0,:], Mu[1,:], ':', color=cfg.plan_color, linewidth=1.)
#plt.plot(Mu[0,:], Mu[1,:], 'ro')
plt_setup()
plt.show()

#%%

order = 4
duration = Mu.shape[1]*0.2 # 0.2 seconds per state
sys = DynSys(order, dt=0.005)
n = sys.num_timesteps(duration)
endw = 1e10
r = 1e-10 # Regularisation 
MuQ, Q = make_reference(Mu, Sigma, n, sys, reference=stepwise, end_weight=1e10) 
x = iterative_mpc(sys, MuQ, Q, r=r)
y = sys.y(x)

# plot it
plt.figure(figsize=cfg.figsize)
plot_gauss(Mu, Sigma)
plt.plot(Mu[0,:], Mu[1,:], ':', color=cfg.plan_color, linewidth=1.)
plt.plot(y[0,:], y[1,:], 'k')
plt_setup()
plt.show()

#%%

from scipy.linalg import eigh

def largest_eigs(X, k):
    ''' Return k eigencomponents with decreasing eigenvalues'''
    N = X.shape[0]
    return eigh(X, eigvals=(N-k,N-1))

def stochastic_sample(x, Cov, mse, n_samples, n_eigs, sig=1.):
    ''' Generate samples from the trajectory distribution'''
    samps = []
    D, V  =  largest_eigs(Cov, n_eigs) 
    D = np.real(np.diag(D))
    V = np.real(V)
    for i in range(n_samples):
         sigma =  np.sqrt(mse)*sig
         o = np.random.randn(n_eigs, 1)*sigma
         offset =  mul([ V, D**.5, o ])
         O = np.reshape( offset, (x.shape[1], x.shape[0])).T
         samp = x + O
         samps.append(samp)

    return samps

#%%

import random
np.random.seed(2311)
order = 5

duration = Mu.shape[1]*0.2
sys = DynSys(order, dt=0.01)
n = sys.num_timesteps(duration)
MuQ, Q = make_reference(Mu, Sigma, n, sys, reference=stepwise, end_weight=1e10)
r = SHM_r(d, order, duration, Mu.shape[1])
x, Cov, mse = batch_mpc(sys, MuQ, Q, r=r, compute_covariance=True)

samps = stochastic_sample(x, Cov, mse, 40, 7, sig=2.5)
samps = [sys.y(samp) for samp in samps]
#x = iterative_mpc(sys, MuQ, Q, r=r)

y = sys.y(x)

# time 
T = np.linspace(0, duration, n)
cfg.plan_color = [0.5, 0.5, 0.5]
plt.figure(figsize=(cfg.fig_w*2.0, cfg.fig_h*0.7))
plt.subplot(1, 2, 1)
#plt.title('Trajectory samples')
plot_gauss(Mu, Sigma) # Note this function def is hidden for saving space
#plt.plot(Mu[0,:], Mu[1,:], ':', color=cfg.plan_color)
np.random.seed(32)
rand_grey = lambda: np.ones(3) * np.random.uniform(0.1, 0.9)
for s in samps:
    plt.plot(s[0,:], s[1,:], 'm', label='samples', linewidth=0.2, color=rand_grey() )
plt.plot(y[0,:], y[1,:], 'k', label='mean trajectory')
plt_setup()
#plt.legend()

plt.subplot(1, 2, 2)
#plt.title('Speed')
np.random.seed(32)
for s in samps:
    S = norm(deriv(s*100, sys.dt, 1))
    plt.plot(np.linspace(0, duration, S.size), S, 'm', linewidth=0.2, color=rand_grey())
S = norm(deriv(y*100, sys.dt, 1))
plt.plot(np.linspace(0, duration, S.size), S, 'k')
plt.gca().set_xlabel('$t$')
plt.gca().set_ylabel(r'$\left\| \dot{x} \right\|$')
plt.savefig('stochastic_n.pdf')
plt.show()

nrows = 2
ncols = 5
n_samps = nrows*ncols
figw = 1.6
plt.figure(figsize=(figw*n_samps, 0.3*figw*n_samps))
random.seed(72)
random.seed(911972)

for i in range(n_samps):
    plt.subplot(nrows, ncols, i+1)
    s = random.choice(samps)
    plt.plot(s[0,:], s[1,:], 'k')
    plt_setup(False)
plt.savefig('stochastic_n_samples.pdf')
plt.show()

#%%

import svgpathtools as svg

def to_pt(c):
    ''' convert complex number to np vector'''
    return np.array([c.real, c.imag])
    
def to_bezier(piece):
    ''' convert a line or Bezier segment to control points'''
    one3d = 1./3
    if type(piece)==svg.path.Line:
        a, b = to_pt(piece.start), to_pt(piece.end)
        return [a, a+(b-a)*one3d, b+(a-b)*one3d, b]
    elif type(piece)==svg.path.CubicBezier:
        return [to_pt(piece.start),
                to_pt(piece.control1),
                to_pt(piece.control2),
                to_pt(piece.end)]
    raise ValueError

def path_to_bezier(path):
    ''' convert SVG path to a Bezier control points'''
    pieces = [to_bezier(piece) for piece in path]
    bezier = [pieces[0][0]] + sum([piece[1:] for piece in pieces],[]) 
    return np.vstack(bezier).T

def svg_to_beziers(path):
    ''' Load Bezier curves from a SVG file'''
    paths, attributes = svg.svg2paths(path) 
    beziers = [path_to_bezier(path) for path in paths]
    return beziers

def num_control_points(n_bezier):
    return 1 + 3*n_bezier

def num_bezier(n_ctrl):
    return int((n_ctrl - 1) / 3)

def num_targets(n_ctrl):
    return num_bezier(n_ctrl)-1

def bezier_piecewise(Cp, subd=1000):
    ''' sample a piecewise Bezier curve given a sequence of control points'''
    num = num_bezier(Cp.shape[1]) 
    X = []
    for i in range(num):
        P = Cp[:,i*3:i*3+4].T
        t = np.linspace(0, 1., subd)[:-1].reshape(-1,1)
        t = np.hstack([t, t])
        Y = (1.0-t)**3 * P[0] + 3*(1.0-t)**2 * t * P[1] + 3*(1.0-t)* t**2 * P[2] + t**3 * P[3]
        X += [Y.T] 
    X = np.hstack(X)
    return X

def plot_control_polygon(Cp):
    n_bezier = num_bezier(Cp.shape[1])
    for i in range(n_bezier):
        cp = cp = Cp[:,i*3:i*3+4]
        plt.plot(cp[0,0:2], cp[1,0:2], ':', color=cfg.plan_color, linewidth=1.)
        plt.plot(cp[0,2:], cp[1,2:], ':', color=cfg.plan_color, linewidth=1.)
        plt.plot(cp[0,:], cp[1,:], 'o', color=[0, 0.5, 1.])

Cp = svg_to_beziers('./a.svg')[0]
yb = bezier_piecewise(Cp)

plt.figure(figsize=cfg.figsize)
plot_control_polygon(Cp)
plt.plot(yb[0,:], yb[1,:], 'k')
plt_setup()
plt.show()

#%%

def mpc_bezier_stepwise(Cp, order, ctrl_ratio, max_d, dt=0.005):
    ''' Mimic a Bezier curve with isotropic Gaussians and a stepwise tracking reference''' 
    sys = DynSys(order, dt=dt)

    m = Cp.shape[1]
    
    T = 0.25 * m #1.0*num_bezier(m)
    n = int(T / sys.dt)

    Mu = np.zeros((2,m))
    Sigma = np.zeros((2,2,m))

    for i in range(num_bezier(m)):
        
        p1,p2,p3,p4 = [Cp[:,j+i*3] for j in range(4)]
        
        d = (p1 - p2 + p4 - p3)/2

        dr = 0.0 
        Mu[:,i*3] = p1
        Mu[:,i*3+3] = p4
        Mu[:,i*3+1] = p2 + d*dr #- d1*dr
        Mu[:,i*3+2] =  p3 + d*dr #- d2*dr

        sigma = np.eye(2)
        Sigma[:,:,i*3]   = sigma
        Sigma[:,:,i*3+1] = np.eye(2)*ctrl_ratio
        Sigma[:,:,i*3+2] = np.eye(2)*ctrl_ratio
        Sigma[:,:,i*3+3] = sigma
        
    MuQ, Q = make_reference(Mu, Sigma, n, sys, reference=stepwise, end_weight=1e10)

    r = SHM_r(max_d, order, T, m)
    
    x = iterative_mpc(sys, MuQ, Q, r=r) 
    y = sys.y(x)
    return y, Mu, Sigma

plt.figure(figsize=(cfg.fig_w*2, cfg.fig_h))
# Plot Bezier and MPC overlaid
plt.subplot(1, 2, 1)
plot_control_polygon(Cp)
y, Mu, Sigma = mpc_bezier_stepwise(Cp, 
                                   order=4, 
                                   ctrl_ratio=2.0, 
                                   max_d=0.01)
plot_gauss(Mu, Sigma, scale=100.)
plt.plot(y[0,:], y[1,:], 'k', label='MPC')
plt.plot(yb[0,:], yb[1,:], 'r', linewidth=1.5, label='Bezier')
plt_setup()

plt.subplot(1, 2, 2)
# compute speed and normalize
S = norm(deriv(y, 0.01, 1))
Sb = norm(deriv(yb, 0.01, 1))
S = S/np.max(S)
Sb = Sb/np.max(Sb)
plt.plot(np.linspace(0, 1., S.shape[0]), S, 'k')
plt.plot(np.linspace(0, 1., Sb.shape[0]), Sb, 'r', linewidth=1.5)
plt.gca().set_xlabel('normalized time')
plt.gca().set_ylabel('normalized speed')
plt.show()

#%%
def via_points(m, n):
    ''' Get via-point reference indices
    returns target indices, and corresponding time steps in reference for each index'''
    a, b = 0., n-1
    return range(m), [ int(a + i*(b-a)/(m-1)) for i in range(m) ]

# For 30 time steps and 3 via points:
Ts = -np.ones(30).astype(int)
I, T = via_points(3, 30)
Ts[T] = I

plt.figure(figsize=(14,0.5))
plt.gca().pcolormesh([[ float(t)/3+0.5 if t >= 0 else 0 for t in Ts]], cmap='Spectral', edgecolor='k', vmin=-1., vmax=1.)
for i, t in enumerate(Ts):
    if t >= 0:
        plt.text(0.4 + i, 0.5,str(t))
plt.axis('off')
plt.show()

def load_and_interpolate(path, order, stroke_duration=0.3, scale=1./100000):
    Mu, Sigma = load_gauss(path, scale=scale) 
    duration = Mu.shape[1]*stroke_duration
    sys = DynSys(order, dt=0.005)
    n = sys.num_timesteps(duration)
    
    endw = 1.
    r = 1e-10

    MuQ, Q = make_reference(Mu, Sigma, n, sys, reference=via_points, end_weight=endw) #1e-15) #0.)
    x = iterative_mpc(sys, MuQ, Q, r=r)
    
    return Mu, Sigma, sys.y(x) 
    
dt = 0.005
orders = range(2,6)

plt.figure(figsize=(cfg.fig_w*len(orders), cfg.fig_h))
for i,order in enumerate(orders):
    plt.subplot(1, len(orders), i+1)
    plt.title('$n=' + str(order) + '$')
    Mu, Sigma, x = load_and_interpolate('e.pkl', order)
    plt.plot(Mu[0,:], Mu[1,:], ':', color=cfg.plan_color, linewidth=1.)
    plt.plot(Mu[0,:], Mu[1,:], 'o', color=[0,0.5,1.], linewidth=1.)
    plt.plot(x[0,:], x[1,:], 'k')
    plt_setup()

plt.show()

#%%

def min_jerk_via( x0,x1,xf, t, t1, tf ):
    tau1 = t1 / tf
    tau = t / tf
    tf5 = tf**5
    ta2 = tau**2 
    ta3 = tau**3
    ta4 = tau**4
    ta5 = tau**5
    ta1_2 = tau1**2 
    ta1_3 = tau1**3
    ta1_4 = tau1**4
    ta1_5 = tau1**5
    
    c = ( (1 / (tf5*ta1_2*(1 - tau1)**5)) * ((xf - x0)*(300.*ta1_5 - 1200.*ta1_4 + 1600.*ta1_3)
        + ta1_2*(-720.*xf + 120.*x1 + 600.*x0) 
            + (x0 - x1)*(300*tau1 - 200)))
    
    p = ( (1 / (tf5*ta1_5*(1 - tau1)**5)) * ((xf - x0)*(120*ta1_5 
            - 300.*ta1_4 + 200.*ta1_3) - 20*(x1 - x0)))
            
    x = ((tf5/720.) * (p*(ta1_4*(15.*ta4 - 30.*ta3) 
    + ta1_3*(80.*ta3 - 30.*ta4) - 60.*ta3*ta1_2 + 30.*ta4*tau1 - 6.*ta5)
                  + c*(15.*ta4 - 10.*ta3 - 6.*ta5)) + x0)
    if t > t1:
        x = x + p * ( tf5 * (tau - tau1)**5 ) / 120.
    return x

import sympy as sym
def solve_tau1(X):
    vals = sym.symbols('x_0, y_0, x_1, y_1, x_f, y_f', real=True)
    x_0, y_0, x_1, y_1, x_f, y_f = vals
    v = [X[0,0], X[1,0], X[0,1], X[1,1], X[0,2], X[1,2]]
    tau_1 = sym.symbols('tau_1', real=True)
    # quite a beast...
    poly = 6*tau_1**9*x_0**2 - 12*tau_1**9*x_0*x_f + 6*tau_1**9*x_f**2 + 6*tau_1**9*y_0**2 \
           - 12*tau_1**9*y_0*y_f + 6*tau_1**9*y_f**2 - 27*tau_1**8*x_0**2 + 54*tau_1**8*x_0*x_f \
           - 27*tau_1**8*x_f**2 - 27*tau_1**8*y_0**2 + 54*tau_1**8*y_0*y_f - 27*tau_1**8*y_f**2 \
           + 40*tau_1**7*x_0**2 - 80*tau_1**7*x_0*x_f + 40*tau_1**7*x_f**2 + 40*tau_1**7*y_0**2 \
           - 80*tau_1**7*y_0*y_f + 40*tau_1**7*y_f**2 - 8*tau_1**6*x_0**2 - 12*tau_1**6*x_0*x_1 \
           + 28*tau_1**6*x_0*x_f + 12*tau_1**6*x_1*x_f - 20*tau_1**6*x_f**2 - 8*tau_1**6*y_0**2 \
           - 12*tau_1**6*y_0*y_1 + 28*tau_1**6*y_0*y_f + 12*tau_1**6*y_1*y_f - 20*tau_1**6*y_f**2 \
           - 36*tau_1**5*x_0**2 + 36*tau_1**5*x_0*x_1 + 36*tau_1**5*x_0*x_f - 36*tau_1**5*x_1*x_f \
           - 36*tau_1**5*y_0**2 + 36*tau_1**5*y_0*y_1 + 36*tau_1**5*y_0*y_f - 36*tau_1**5*y_1*y_f \
           + 34*tau_1**4*x_0**2 - 34*tau_1**4*x_0*x_1 - 34*tau_1**4*x_0*x_f + 34*tau_1**4*x_1*x_f \
           + 34*tau_1**4*y_0**2 - 34*tau_1**4*y_0*y_1 - 34*tau_1**4*y_0*y_f + 34*tau_1**4*y_1*y_f \
           - 8*tau_1**3*x_0**2 + 8*tau_1**3*x_0*x_1 + 8*tau_1**3*x_0*x_f - 8*tau_1**3*x_1*x_f \
           - 8*tau_1**3*y_0**2 + 8*tau_1**3*y_0*y_1 + 8*tau_1**3*y_0*y_f - 8*tau_1**3*y_1*y_f \
           - 2*tau_1*x_0**2 + 4*tau_1*x_0*x_1 - 2*tau_1*x_1**2 - 2*tau_1*y_0**2 + 4*tau_1*y_0*y_1 \
           - 2*tau_1*y_1**2 + x_0**2 - 2*x_0*x_1 + x_1**2 + y_0**2 - 2*y_0*y_1 + y_1**2
    eval_poly = poly.subs(dict(zip(vals,v)))
    expr = (sym.expand(eval_poly.as_expr()))
    realroot = sym.nsolve(expr, 0.5) 
    return float(realroot)

def via_points_tau1(m, n, tau_1):
    ''' Get via-point reference indices
    returns target indices, and corresponding time steps in reference for each index'''
    a, b = 0., n-1
    I = range(m)
    t1_n = int(tau_1*n)
    T = [0, t1_n, n-1]
    return I, T #range(m), [ int(a + i*(b-a)/(m-1)) for i in range(m) ]

def rel_err(x, x_hat):
    '''Relative error [0-1]'''
    e = (np.sum(np.abs(x-x_hat)) / np.sum(x_hat)) / len(x_hat)
    return e


sys = DynSys(3, 0.005)
duration = 2.
n = sys.num_timesteps(duration)

# MPC trajectory
plt.rc('text', usetex = False)
m = 3
np.random.seed(123)
Mu = np.floor(np.random.uniform(-100, 100, size=(2, m)))
Sigma = np.dstack([np.eye(2,2) for i in range(m)])

# Solve for via point relative time
tau_1 = solve_tau1(Mu)

MuQ, Q = make_reference(Mu, Sigma, n, sys, reference=lambda m,n: via_points_tau1(m, n, tau_1), end_weight=1.)
x = iterative_mpc(sys, MuQ, Q, r=1e-10)
y = sys.y(x)

# time steps
T = np.linspace(0, duration, n)

# Minimum jerk trajectory
yj = np.array([ min_jerk_via(Mu[:,0], Mu[:,1], Mu[:,2], t, tau_1*duration, duration) for t in T]).T

plt.figure(figsize=(cfg.fig_w*2, 1.3*cfg.fig_h)) #(15,5))
vmargin = 0.4
plt.subplots_adjust(hspace=vmargin) #, top=vmargin+2.)
plt.subplot2grid((2, 2), (0,0), rowspan=2)
plt.plot(Mu[0,:], Mu[1,:], ':', color=cfg.plan_color, linewidth=1.)
plt.plot(Mu[0,:], Mu[1,:], 'o', color=[0,0.5,1.])
plt.plot(Mu[0,0:1], Mu[1,0:1], 'ro') #, color=[0,0.5,1.])
plt.plot(y[0,:], y[1,:], 'k', label='MPC')
plt.plot(yj[0,:], yj[1,:], 'r', label='MJ', linewidth=1.5)

plt_setup()
plt.legend()

# Speed error
plt.subplot2grid((2, 2), (0,1))
#plt.title('Speed')
S = norm(deriv(y, sys.dt, 1))
Sj = norm(deriv(yj, sys.dt, 1))
# size hack for time
T = np.linspace(0, duration, len(Sj))
err = rel_err(S, Sj)
errstr = ('%.3f'%(err*100)) + '%'
plt.plot(T, S, 'k')
plt.plot(T, Sj, 'r', linewidth=1.5)
plt.text(2.0, 130, 'Speed rel. error: ' + errstr, horizontalalignment='right')
plt.gca().set_xlabel('$t$')
plt.gca().set_ylabel(r'$\left\| \dot{x} \right\|$')

# Jerk error
plt.subplot2grid((2, 2), (1,1))
#plt.title('Jerk')
S = norm(deriv(y, sys.dt, 3))[1:]
Sj = norm(deriv(yj, sys.dt, 3))[1:]
# size hack for time
T = np.linspace(0, duration, len(Sj))
err = rel_err(S, Sj)
errstr = ('%.3f'%(err*100)) + '%'
plt.plot(T, Sj, 'r')
plt.plot(T, S, 'k')
plt.text(2., 2500.0, 'Jerk rel. error: ' + errstr, horizontalalignment='right')
plt.gca().set_xlabel(r'$t$')
plt.gca().set_ylabel(r'$\left\| \dot\ddot{x} \right\|$')

# Save and show
plt.savefig('mj.pdf')
plt.show()