#!/usr/bin/env python
# coding: utf-8

# Note that the first ten cells do nothing more than set up the grid of roughly equidistant points on the sphere, the actual code for the quantum simulation starts at cell 11
# (for a slightly better icosahedral-based grid, see http://iopscience.iop.org/article/10.1086/310310/pdf )
# 
# [Note also this wasn't done very carefully nor necessarily cleanly ...]

# In[1]:


get_ipython().run_line_magic('pylab', 'inline')
from mpl_toolkits.mplot3d import Axes3D


# In[2]:


#first a single equilateral triangle
v=[(-.5,-.5/sqrt(3)),(.5,-.5/sqrt(3)),(0,sqrt(3)/2.-.5/sqrt(3))]
v=map(array,v)


# In[3]:


def intv(v):  #find the interior points for dividing sides into m subsegments
    u=[v[1]-v[0],v[2]-v[0]]
    uw=[]
    for i in range(2,m):
      w0=v[0]+u[0]*i/float(m)
      w1=v[0]+u[1]*i/float(m)
      uw += [w0 + (w1-w0)*j/float(i) for j in range(1,i)]
    return uw


# In[4]:


# total number of gridpoints as function of m
for i in range(1,10): #vertices + edges + interior
    print i, 12 + 30*(i-1) + 20*(i-1)*(i-2)/2


# In[5]:


m=7  #492 should be enough
#check graphically that intv(v) works as advertised

plot(*zip(*(v+[v[0]])))
plot(*zip(*v),marker='o',ls='None',c='r')
w=[v[1]-v[0],v[2]-v[1],v[0]-v[2]] #three side vectors
for i in range(3):
    vw=[v[i]+j*w[i]/float(m) for j in range(1,m)]
    plot(*zip(*vw),marker='o',ls='None',c='b')
    
grid('on')
plot(*zip(*intv(v)),marker='o',ls='None',c='g');  #intv(v) in green


# In[6]:


#coordinates for the 12 icosahedral vertices
phi=(1+sqrt(5))/2
h=sqrt(1+phi**2)  #main diag, for top and bottom vertices
h1=phi/h          #height of pentagon slices
phi5=arange(0,2*pi,2*pi/5)
rp=1/sin(pi/5)
vp1=array([(rp*cos(p),rp*sin(p),-h1) for p in phi5]) #lower five vertices
vp2=array([(rp*cos(p),rp*sin(p),h1) for p in phi5+pi/5])  #upper five vertices
rc=sqrt(phi**2 +1) -1  #middle radius


# In[7]:


iedges=[]  #list of the 30 icosahedral edges
for v in vp1: iedges.append((v,(0,0,-h)))  #five bottom
for v in vp2: iedges.append((v,(0,0,h)))   # five top
for i in range(5):
    iedges.append((vp1[i],vp1[(i+1)%5]))  #lower pentagon
    iedges.append((vp2[i],vp2[(i+1)%5]))  #upper pentagon
    iedges.append((vp1[i],vp2[(i-1)%5]))  #forward middle edges
    iedges.append((vp1[i],vp2[i]))        #backward middle edges


# In[8]:


#now use intv(v) to add all of the interior grid points
#recall m=7 seems sufficient
mverts = list(vp1)+list(vp2) + [(0,0,-h),(0,0,h)] # will be full grid
iverts=list(vp1)+list(vp2) + [(0,0,-h),(0,0,h)]   #just icos vertices

for e in iedges:  #vertices along edges
    u= (e[1] - e[0])/float(m)
    mverts += [e[0]+i*u for i in range(1,m)]

for i in range(5):  # middle lower five faces
    mverts += intv((vp1[i],vp1[(i+1)%5],vp2[i]))
for i in range(5):  # middle upper five faces
    mverts += intv((vp2[i],vp2[(i+1)%5],vp1[(i+1)%5]))
for i in range(5):  #upper five faces
    mverts+= intv([(0,0,h),vp2[i],vp2[(i+1)%5]])
for i in range(5):  #lower five faces
    mverts+= intv([(0,0,-h),vp1[i],vp1[(i+1)%5]])


# In[9]:


sverts = [mv/norm(mv) for mv in mverts] #normalize to unit radius
iverts = [mv/norm(mv) for mv in iverts]
sgrid=array(sverts)
n=len(sverts)
n  #gridpoints


# In[10]:


#see if it all worked, including icos coords
fig=figure(figsize=(25,12))

ax = fig.add_subplot(121, projection='3d')
for e in iedges:
    ax.plot(*zip(*e),color='b',linestyle='--')
ax.scatter(*zip(*mverts) ,c='r',s=40)
xlim(-1.5,1.5)
ylim(-1.5,1.5)
ax.set_zlim(-1.5,1.5)
ax.elev=20 #20
#ax.azim=-36  #-60

ax = fig.add_subplot(122, projection='3d')
ax.scatter(*zip(*sverts),s=40,c='r')
ax.scatter(*zip(*iverts),s=160,c='b')
ax.elev=20 #20
#ax.azim=-36  #-60
xlim(-1,1)
ylim(-1,1)
ax.set_zlim(-1,1);


# In[ ]:





# All of the above was just the elementary coordinate geometry to set up some form of roughly equal-spaced grid on the Bloch sphere.
# 
# Here is where the actual code begins:

# In[11]:


def H(p):  #entropy
    p = p[p>0]
    return -p.dot(log2(p))


# In[12]:


def show(p,nh=None,x=None,npng=None,elev=30,azim=-60,myax=None):
  """current probs, measurement axis, value"""
  ax=myax
  ax.elev=elev
  ax.azim=azim
  ax.scatter(*zip(*sverts),c=p,cmap=cm.coolwarm,s=40)  #use red/blue colormap for probability
  titletext='S={:.2f}'.format(H(p))
  if nh is not None:  #show the measurement with green background
    if nh[1]>nh[0]:
        nh=-nh
        x=(x+1)%2
    ax.plot(*zip(nh),marker='o',markersize=25,color='g',alpha=.4)
    ax.text(nh[0],nh[1],nh[2],str(x),color='yellow',fontsize=20,ha='center',va='center')
  if npng: titletext += ', m={}'.format(npng)
  ax.set_title(titletext)


# In[13]:


def fegain(prob):
#calculates the info gain for every sphere grid point
  egain=[]
  s0=H(prob)
  for nh in sgrid:
    prob0=prob.dot((1.+sgrid.dot(nh))/2.)  #cos^2(th/2) = (1+cos(th))/2
    prob1=1-prob0
    p0=prob*((1.+sgrid.dot(nh))/2.)/prob0
    p1=prob*((1.-sgrid.dot(nh))/2.)/prob1
    einfo=prob0*H(p0)+prob1*H(p1)
    egain.append(s0-einfo)
  return array(egain)


# In[14]:


def update(p,nh,rs):  #do the random measurement
    prob0=prob.dot((1.+sgrid.dot(nh))/2.)
    prob1=1-prob0
    x = 0. if rand() < (nh.dot(rs)+1.)/2. else 1.  #cos^2(th/2) = (1+cos(th))/2
    if x==0: return (0,p*((1.+sgrid.dot(nh))/2.)/prob0)
    else: return (1,p*((1.-sgrid.dot(nh))/2.)/prob1)


# In[15]:


rs=array([1.,0.,1.])/sqrt(2.)
argmin([norm(sv-rs) for sv in sgrid]) #use closest grid state to (1,0,1)/sqrt(2)


# In[16]:


#initialize
prob=ones(n)/n
rs=sgrid[414] #closest to (1,0,1)/sqrt(2)
m=0
sa=[H(prob)]
x=0


# In[17]:


from JSAnimation import IPython_display
from matplotlib import animation


# In[18]:


fig=figure(figsize=(6,6))
ax = fig.add_subplot(111, projection='3d')
s=1
ax.set_xlim(-s,s)
ax.set_ylim(-s,s)
ax.set_zlim(-s,s)
ax.axis('off')

def init():
   show(prob,array([0,0,1]),0,myax=ax)

def animate(f):  #takes frame num as arg
    global prob,m,ax
    ax.clear()
    ax.set_xlim(-s,s)
    ax.set_ylim(-s,s)
    ax.set_zlim(-s,s)
    print f, #sa[-1],m
    if f<2 and m==0:
        show(prob,array([0,0,1]),0,myax=ax,npng=m)
        return()
    egain=fegain(prob)
    mnh=sgrid[argmax(egain)]
    if m==0: mnh=array([0,0,1])
    x,newprob=update(prob,mnh,rs)
    while (sa[-1] < 5 and m%10 != 0):
          prob=newprob
          sa.append(H(prob))
          egain=fegain(prob)
          mnh=sgrid[argmax(egain)]
          x,newprob=update(prob,mnh,rs)
          m+=1
    show(prob,mnh,x,myax=ax,npng=m)
    m+=1
    prob=newprob
    sa.append(H(prob))

anim = animation.FuncAnimation(fig, animate, frames=75, interval=500, blit=False)

anim


# In[19]:


plot(sa); #entropy of prob distribution as a function of #measurements


# In[20]:


#reinitialize for simulation with deterministic choice of measurement axes
prob=ones(n)/n
rs=sgrid[414] #closest to (1,0,1)/sqrt(2)
m=0
sa=[H(prob)]
x=0

fig=figure(figsize=(6,6))
ax = fig.add_subplot(111, projection='3d')
s=1
ax.set_xlim(-s,s)
ax.set_ylim(-s,s)
ax.set_zlim(-s,s)
ax.axis('off')

def animate(f):  #takes frame num as arg
    global prob,m,ax
    ax.clear()
    ax.set_xlim(-s,s)
    ax.set_ylim(-s,s)
    ax.set_zlim(-s,s)
    print f, #sa[-1],m
    if f<2 and m==0:
        show(prob,array([0,0,1]),0,myax=ax,npng=m)
        return()
    mnh=array([[0,0,1],[0,1,0],[1,0,0]][m%3])    #alternate z,y,x axes
    if m==0: mnh=array([0,0,1])
    x,newprob=update(prob,mnh,rs)
    while (sa[-1] < 5 and m%10 != 0):
          prob=newprob
          sa.append(H(prob))
          mnh=array([[0,0,1],[0,1,0],[1,0,0]][m%3])
          x,newprob=update(prob,mnh,rs)
          m+=1
    show(prob,mnh,x,myax=ax,npng=m)
    m+=1
    prob=newprob
    sa.append(H(prob))

anim = animation.FuncAnimation(fig, animate, frames=75, interval=500, blit=False)

anim


# In[21]:


plot(sa); #entropy of prob distribution as a function of #measurements


# In[ ]: