from pylab import *
import subprocess
import sys
import os

figformat = '.png'
seterr(divide='ignore')
rcParams['font.size'] = 9

#define color map that is transparent for low values, and dark blue for high values.
# weighted to show low probabilities well
cdic = {'red':   [(0,0,0),(1,0,0)],
        'green': [(0,0,0),(1,0,0)],
        'blue':  [(0,0.7,0.7),(1,0.7,0.7)],
        'alpha': [(0,0,0),
                  (0.1,0.4,0.4),
                  (0.2,0.6,0.6),
                  (0.4,0.8,0.8),
                  (0.6,0.9,0.9),
                  (1,1,1)]}
cm_prob = matplotlib.colors.LinearSegmentedColormap('prob',cdic,N=640)

### System dynamics ###

# potential is a polynomial
potential_coefs = array([1,0,0,4,-5,-4,0],'d')
def potential(x,t):
    return polyval(potential_coefs,x)

# force function is its derivative.
force_coefs = (potential_coefs*arange(len(potential_coefs)-1,-1,-1))[:-1]
def force(x,t):
    """ derivative of potential(x) """
    return polyval(force_coefs,x)
invmass = 1.0
dt = 0.03

def motion(t,x,p):
    """ returns dx/dt, dp/dt """
    return p*invmass, -force(x,t)

cur_x = -0.1
cur_p = 0

def rkky_step(t, x_i, p_i, dt):
    kx1,kp1 = motion(t, x_i, p_i)
    dt2 = 0.5*dt
    kx2,kp2 = motion(t+dt2, x_i+dt2*kx1, p_i+dt2*kp1)
    kx3,kp3 = motion(t+dt2, x_i+dt2*kx2, p_i+dt2*kp2)
    kx4,kp4 = motion(t+dt, x_i+dt*kx3, p_i+dt*kp3)
    newx = x_i + (dt/6.0)*(kx1 + 2.0*kx2 + 2.0*kx3 + kx4)
    newp = p_i + (dt/6.0)*(kp1 + 2.0*kp2 + 2.0*kp3 + kp4)
    return newx, newp

### Setup ensemble points ###

# most are randomly chosen
x = 0 + 0.5*rand(20000)
p = -1.0 + 2.0*rand(20000)

# the pilot points are set manually
x[0] = 0;    p[0] = 0
x[1] = 0.4;  p[1] = 0.0
pilots = [0,1]
pilot_colors = {
       0: (0,0.7,0),
       1: (0.7,0,0)}
E = potential(x,0) + 0.5*invmass*p**2

### set up plot limits and histogram bins ###
xedges = linspace(-2.1,1.7,151)
pedges = linspace(-7.5,7.5,151)
Eedges = linspace(-9,9,151)
pix = 150
extent = [xedges[0], xedges[-1], pedges[-1], pedges[0]]
H = histogram2d(x,p,bins=[xedges,pedges])[0].transpose()
cmax = amax(H)*0.8

extenten = [xedges[0], xedges[-1], Eedges[-1], Eedges[0]]
Hen = histogram2d(x,E,bins=[xedges,Eedges])[0].transpose()
cmaxen = amax(Hen)*0.3

fig = figure(1)
ysize = 2.6
xsize = 1.3
fig.set_size_inches(xsize,ysize)

### Prepare lower plot ###
axen = axes((0.2/xsize,0.2/ysize,1.0/xsize,1.0/ysize),frameon=True)
axen.xaxis.set_ticks([])
axen.xaxis.labelpad = 2
axen.yaxis.set_ticks([])
axen.yaxis.labelpad = 2
xlim(-2.1,1.7)
ylim(-9,9)
xlabel('position $x$')
ylabel('energy')
potx = linspace(-2.1,1.7,151)

### Prepare upper plot ###
ax = axes((0.2/xsize,1.5/ysize,1.0/xsize,1.0/ysize),frameon=True)
ax.xaxis.set_ticks([])
ax.xaxis.labelpad = 2
ax.yaxis.set_ticks([])
ax.yaxis.labelpad = 2
xlim(-2.1,1.7)
ylim(-7.5,7.5)
xlabel('position $x$')
ylabel('momentum $p$')

### Start running simulation ###
frames = list()
delays = list()
framemod = 5
frame = "frames/background"+figformat
savefig(frame,dpi=pix)
frames.append(frame)
delays.append(16)

print "generating frames...  0%",
sys.stdout.flush()
savesteps = range(0,401,framemod) + [600, 1000, 2000, 6000]
delays += [10]*len(savesteps)
delays[1] = 200
delays[-5:] = [100,200,200,200,400]
totalsteps = max(savesteps)+1
for step in range(totalsteps):
    if step % 20 == 0:
        print "\b\b\b\b\b{0:3}%".format(int(round(step*100.0/totalsteps))),
        sys.stdout.flush()
    if step in savesteps:
        # Every several frames, do a plot
        remlist = list()

        sca(ax)
        H = histogram2d(x,p,bins=[xedges,pedges])[0].transpose()
        remlist.append(imshow(H, extent=extent, cmap=cm_prob, interpolation='none', aspect='auto'))
        remlist[-1].set_clim(0,cmax)
        for i in pilots:
            remlist += plot(x[i], p[i], '.', color=pilot_colors[i], markersize=3)

        E = potential(x,step*dt) + 0.5*invmass*p**2
        sca(axen)
        pot = potential(potx,step*dt)
        remlist += plot(potx,pot,color='r',zorder=0)
        Hen = histogram2d(x,E,bins=[xedges,Eedges])[0].transpose()
        remlist.append(imshow(Hen, extent=extenten, cmap=cm_prob, interpolation='none', aspect='auto',zorder=1))
        remlist[-1].set_clim(0,cmaxen)
        for i in pilots:
            remlist += plot(x[i], E[i], '.', color=pilot_colors[i], markersize=3)

        frame = "frames/frame"+str(step)+figformat
        savefig(frame,dpi=pix)
        frames.append(frame)
        # Clear out updated stuff.
        for r in remlist: r.remove()
    x, p = rkky_step(step*dt, x, p,dt)
print "\b\b\b\b\b      done"

assert(len(delays) == len(frames))

### Assemble animation using ImageMagick ###
calllist = 'convert -dispose Background'.split()
for delay,frame in zip(delays,frames):
    calllist += ['-delay',str(delay)]
    calllist += [frame]
calllist += '-loop 0 -layers Optimize _animation.gif'.split()
f = open('anim_command.txt','w')
f.write(' '.join(calllist)+'\n')
f.close()

print "composing into animated gif...",
sys.stdout.flush()
subprocess.call(calllist)
print "      done"
os.rename('_animation.gif','animation.gif')