import numpy as np
import nengo
import nengo_spa as spa
import matplotlib.pyplot as plt
import semanticmapping

domain_dim = 2
d=97
bounds = 5*np.array([[-1,1],[-1,1]])
ssp_space = semanticmapping.HexagonalSSPSpace(domain_dim,ssp_dim=d, 
                 scale_min=0.1, scale_max=3.1,
                 domain_bounds=bounds, length_scale=0.1)

d=ssp_space.ssp_dim



T = 5
dt = 0.001
timesteps = np.arange(0, T, dt)
path = np.vstack([np.linspace(0,4,len(timesteps)), np.linspace(0,4,len(timesteps))]).T
#thetas = np.linspace(0,np.pi/2,len(timesteps))
#path = np.vstack([3*np.cos(thetas), 3*np.sin(thetas)]).T 
#path = path - path[0,:]
n_pause = 100
path = np.vstack([np.repeat(path[0,:].reshape(1,-1),n_pause,axis=0), path])
T = T + dt*n_pause
timesteps = np.arange(0, T, dt)
pathlen = path.shape[0]

plt.figure(figsize=(5,5))
plt.plot(path[:,0],path[:,1])


vels = (1/dt)*( path[(np.minimum(np.floor(timesteps/dt) + 1, pathlen-1)).astype(int),:] -
               path[(np.minimum(np.floor(timesteps/dt), pathlen-2)).astype(int),:])
real_freqs = (ssp_space.phase_matrix @ vels.T)
scale_fac = 2/(np.max(real_freqs) - np.min(real_freqs))
vels_scaled = vels*scale_fac


real_fssp = np.fft.fftshift(ssp_space.encode_fourier(path.T),axes=0)
init_features = np.array([[0,1],[-0.5,-0.5],[0.5,-0.5]])
init_ssp = ssp_space.unitary(np.sum(ssp_space.encode(init_features.T),axis=1))
real_fssp_i = np.fft.fftshift(ssp_space.unitary_fourier(np.sum(ssp_space.encode_fourier(init_features.T),axis=1)))

stim_fssp = np.ones(d-1)
stim_fssp[:d//2] = real_fssp_i.real[:d//2]
stim_fssp[d//2:] =real_fssp_i.imag[:d//2]

n_neurons = 5000
tau = 0.05
model = nengo.Network(seed=1)
with model:
    vel_input = nengo.Node(lambda t: vels_scaled[int(t/dt)-1,:], size_out=2)
    stim = nengo.Node(lambda t,x: stim_fssp if t<=n_pause*dt else np.zeros((d-1,)), size_in=d-1)
    pathintegrator = semanticmapping.networks.PathIntegration(ssp_space, tau, n_neurons, 
                   scaling_factor=scale_fac, stable=True)
    pathintegrator2 = semanticmapping.networks.PathIntegration_GCs9(ssp_space, tau, n_neurons, 
                  scaling_factor=scale_fac, stable=True,error_correction_factor=0)
    nengo.Connection(vel_input,pathintegrator.velocity_input)
    nengo.Connection(stim,pathintegrator.stim_input)
    
    nengo.Connection(vel_input,pathintegrator2.velocity_input)
    nengo.Connection(stim,pathintegrator2.stim_input)

    ssp_p = nengo.Probe(pathintegrator.output, synapse=0.05)
    ssp2_p = nengo.Probe(pathintegrator2.output, synapse=0.05)

sim = nengo.Simulator(model)
sim.run(T)

def plot_path(path, real_path, N_ma=None,ax=None):
    if ax is None:
        fig = plt.figure(figsize=(6.50127/3, 6.50127/3))
        ax = fig.add_subplot(111)
    if real_path is not None:
        ax.plot(real_path[:,0],real_path[:,1],c='lightgrey',linewidth=1,zorder = 1)
    if N_ma is None:
        ax.plot(path[:,0],path[:,1],linestyle='--',color="black",zorder = 1,linewidth=0.7)
    else:
        path_ma = np.vstack([np.convolve(path[:,0], np.ones(N_ma)/N_ma, mode='valid'), 
                              np.convolve(path[:,1], np.ones(N_ma)/N_ma, mode='valid')]).T
        ax.plot(path_ma[:,0],path_ma[:,1],linestyle='--',color="black",zorder = 1,linewidth=0.7)
    return ax
def get_binding_matrix(v):
    D = len(v)
    T = []
    for i in range(D):
        T.append([v[(i - j) % D] for j in range(D)])
    return np.array(T)

def get_inversion_matrix(d):
    return np.eye(d)[-np.arange(d)]

A = get_binding_matrix(ssp_space.encode(path[0,:]) ) @ get_binding_matrix(get_inversion_matrix(d) @ init_ssp)
sim_path, max_sims = ssp_space.decode_path(sim.data[ssp_p] @ A.T)
sim_path2, max_sims2 = ssp_space.decode_path(sim.data[ssp2_p] @ A.T)

Ai = get_binding_matrix(init_ssp ) @ get_binding_matrix(get_inversion_matrix(d) @ ssp_space.encode(path[0,:]))
real_ssp = Ai @ ssp_space.encode(path.T) 

fig=plt.figure(figsize=(8,8))
ax1 = fig.add_subplot(221)
ax1.plot(sim.trange(), np.sqrt(np.sum(sim.data[ssp_p]*real_ssp.T,axis=1)))
ax1.plot(sim.trange(), max_sims)
ax1.set_ylim(0,1)
ax2 = fig.add_subplot(222)
ax2 = plot_path(sim_path, path, N_ma=200,ax=ax2)
ax3 = fig.add_subplot(223)
ax3.plot(sim.trange(), np.sqrt(np.sum(sim.data[ssp2_p]*real_ssp.T,axis=1)))
ax3.plot(sim.trange(), max_sims2)
ax3.set_ylim(0,1)
ax4 = fig.add_subplot(224)
ax4 = plot_path(sim_path2, path, N_ma=200,ax=ax4)