#!/usr/bin/python
# -*- coding: utf8 -*-

from math import *
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon, Circle, Wedge
from matplotlib import animation
import numpy as np

# settings
fname = 'Mach-Zehnder_photons_animation'
width, height = 300, 220
nframes = 100
nphotons = 12
fps = 15

x0 = 100.5
x1 = 218.5
y0 = 200.5
y1 = 80.5
lx, lw, lh = 5, 46, 21 # laser
dtect = 62.5
t1, t2, tmove = 0.25, 0.9, 0.025
ymove = 24
rp = 2. # photon radius
cp1 = '#ff0000' # photon color
cp2 = '#ffaaaa' # splitphoton color

##
xstart = lx + lw / 2.
dx = x1 - x0
dy = y1 - y0
l = (x0 - xstart) + abs(dx) + abs(dy) + dtect + 2.*rp
xdet0 = (x0 + x1) / 2
fly_frac = 0.7
v = l / fly_frac
tdet0 = (xdet0 + 2.*rp - xstart) / v
tdet12 = l / v

# introduce artificial antibunching for illustration purpose
ptimes = (np.random.random() + np.sort(np.random.random(3*nphotons))[::3]) % 1

photons = [{} for i in range(nphotons)]
for i, p in enumerate(photons):
    p['t0'] = ptimes[i]
    if t1 <= (p['t0'] + tdet0) % 1 and (p['t0'] + tdet0) % 1 <= t2:
        # photon sees first detector
        if np.random.randint(2) == 0:
            # photon hits extra detector
            p['arm'] = 'none'
            p['det'] = 0
        else:
            # photon escapes first detector
            p['arm'] = 'lower'
            # => random detection at second beam splitter
            if np.random.randint(2) == 0:
                p['det'] = 1
            else:
                p['det'] = 2
    else:
        # photon sees standard Mach-Zehnder interferometer
        p['arm'] = 'both'
        p['det'] = 1
    
    if p['det'] == 0:
        p['tdet'] = (p['t0'] + tdet0) % 1
    else:
        p['tdet'] = (p['t0'] + tdet12) % 1
    p['click_frame'] = int(round(p['tdet'] * nframes)) % nframes

plt.close('all')
mpl.rc('path', snap=False)

def animate(nframe):
    # prepare a clean and image-filling canvas for each frame
    plt.clf()
    fig.gca().set_position((0, 0, 1, 1))
    plt.xlim(0, width)
    plt.ylim(0, height)
    plt.axis('off')
    
    t = float(nframe) / nframes
    
    # photons
    for p in photons:
        s0 = v * ((t - p['t0']) % 1)
        if s0 > l:
            continue
        s = s0 + start - x0
        if s <= 0:
            # from laser to first beam splitter
            x, y = x0 + s, y0
            fig.gca().add_patch(Circle((x, y), rp, color=cp1))
        elif s <= abs(dx) + abs(dy):
            # in the interferometer
            if s < abs(dx):
                xu, yu = x0 + copysign(s, dx), y0
            else:
                xu, yu = x1, y0 + copysign(s - abs(dx), dy)
            if s < abs(dy):
                xd, yd = x0, y0 + copysign(s, dy)
            else:
                xd, yd = x0 + copysign(s - abs(dy), dx), y1
                
            if s < xdet0 - x0 or p['arm'] == 'both':
                fig.gca().add_patch(Circle((xu, yu), rp, color=cp2))
                fig.gca().add_patch(Circle((xd, yd), rp, color=cp2))
            elif p['arm'] == 'lower':
                fig.gca().add_patch(Circle((xd, yd), rp, color=cp1))
        else:
            # after the interferometer
            x, y = x1 + (s - abs(dx) - abs(dy)), y1
            if p['arm'] == 'both':
                fig.gca().add_patch(Circle((x, y), rp, color=cp1))
            elif p['arm'] == 'lower':
                fig.gca().add_patch(Circle((x, y), rp, color=cp2))
                x, y = x1, y1 - (s - abs(dx) - abs(dy))
                fig.gca().add_patch(Circle((x, y), rp, color=cp2))
    
    
    # laser
    fig.gca().add_patch(
        Polygon([[lx, y0-lh/2.], [lx, y0+lh/2.],
                 [lx+lw, y0+lh/2.], [lx+lw, y0-lh/2.]],
            closed=True, facecolor='#cccccc', edgecolor='black'))
    plt.text(lx+lw/2., y0-2, 'laser', fontsize=12,
        horizontalalignment='center', verticalalignment='center')
    
    # beam splitters
    b = 12
    fig.gca().add_patch(
        Polygon([[x0-b, y0+b], [x0+b, y0+b], [x0+b, y0-b],
                 [x0-b, y0-b], [x0-b, y0+b], [x0+b, y0-b]],
            closed=True, facecolor='#88aadd', edgecolor='black',
            linewidth=2, alpha=0.4))
    fig.gca().add_patch(
        Polygon([[x1-b, y1+b], [x1+b, y1+b], [x1+b, y1-b],
                 [x1-b, y1-b], [x1-b, y1+b], [x1+b, y1-b]],
            closed=True, facecolor='#88aadd', edgecolor='black',
            linewidth=2, alpha=0.4))
    
    # mirrors
    m, mw = 12, 4
    fig.gca().add_patch(
        Polygon([[x1-m+mw/2., y0+m+mw/2.], [x1+m+mw/2., y0-m+mw/2.]],
            closed=False, edgecolor='#555555', linewidth=mw))
    fig.gca().add_patch(
        Polygon([[x0-m-mw/2., y1+m-mw/2.], [x0+m-mw/2., y1-m-mw/2.]],
            closed=False, edgecolor='#555555', linewidth=mw))
    
    # detectors
    c_off = '#cccccc'
    c_on = '#cc0000'
    c0 = c1 = c2 = c_off
    for p in photons:
        if p['click_frame'] == nframe:
            if p['det'] == 0: c0 = c_on
            if p['det'] == 1: c1 = c_on
            if p['det'] == 2: c2 = c_on
    if t1 <= t and t <= t2:
        yd = y0
    else:
        yd = y0 - min((t1-t)%1, tmove, (t-t2)%1) * ymove / float(tmove)
    fig.gca().add_patch(mpl.patches.Wedge((xdet0, yd), b, 270, 90, fc=c0))
    fig.gca().add_patch(mpl.patches.Wedge((x1 + dtect, y1), b, 270, 90, fc=c1))
    fig.gca().add_patch(mpl.patches.Wedge((x1, y1 - dtect), b, 180, 0, fc=c2))

fig = plt.figure(figsize=(width/100., height/100.))
anim = animation.FuncAnimation(fig, animate, frames=nframes)
anim.save(fname + '.gif', writer='imagemagick', fps=fps)