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

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import qutip
import numpy as np
import scipy
import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from IPython import display
import sys
import time
import pickle

from itertools import product
from types import FunctionType, BuiltinFunctionType
from functools import partial

def expand_lindblad_operator_to_qutip_std_form(*summands):
    """Expands the superoperator L(A+B+...) to sum of independent dissipator operators.
    
    Warning: only real coefficients are supported.
    
    The input `[matA, coefA], [matB, coefB], ...` represents the Lindblad dissipator
    `L(matA*coefA + matB*coefB + ...)`. `mat*` is an operator, `coef*` is a time
    dependent coefficient (in any of the qutip supported formats).
    
    If a summand is constant just input it on its own,
    e.g. `[matA, coefA], matConst, [matB, coefB]`.
    
    The output is a list of independent dissipator superoperators in form that can be
    used directly with `mesolve`'s `c_ops`."""
    matrix_coeff_pairs = [_ if isinstance(_, list) else [_, None] for _ in summands]

    if all(isinstance(_, np.ndarray) or _ is None for __, _ in matrix_coeff_pairs):
        coef_type = 'array'
    elif all(isinstance(_, str) or _ is None for __, _ in matrix_coeff_pairs):
        coef_type = 'str'
    elif all(isinstance(_, (FunctionType, BuiltinFunctionType, partial)) or _ is None for __, _ in matrix_coeff_pairs):
        coef_type = 'function'
    else:
        raise TypeError("Incorrect format for the coefficients")

    c_ops = sum((properly_conjugated_lindblad(*_, coef_type=coef_type) for _ in product(matrix_coeff_pairs, repeat=2)),[])
    return sorted(c_ops, key=lambda _:isinstance(_, list), reverse=True)

def properly_conjugated_lindblad(left, right, coef_type):
    '''Return the two components of the Lindblad superoperator with properly conjugated coeffs.
    
    For input `(A, ca), (B, cb)` return 
    `[[pre(A)*post(Bdag), ca*conj(cb)],[-0.5*(pre(Adag*B)+post(Adag*B)), conj(ca)*cb]]`.
    
    It supports constant operators (just set `ca` or `cb` to `None`).
    '''
    if coef_type == 'array':
        conj = lambda _: np.conj(_)
        prod = lambda a,b: a*b
    elif coef_type == 'str':
        conj = lambda _: 'conj(%s)'%_
        prod = lambda a,b: '(%s)*(%s)'%(a,b)
    elif coef_type == 'function':
        raise NotImplementedError
        
    m_left, c_left = left
    m_right, c_right = right
    
    A = qutip.spre(m_left) * qutip.spost(m_right.dag())
    ld_r = m_left.dag()*m_right
    B = - 0.5 * qutip.spre(ld_r) - 0.5 * qutip.spost(ld_r)

    if c_left is None and c_right is None:
        return [A+B]
    elif c_left is None:
        return [[A, conj(c_right)],
                [B,      c_right]]
    elif c_right is None:
        return [[A,      c_left],
                [B, conj(c_left)]]
    else:
        return [[A, prod(     c_left,  conj(c_right))],
                [B, prod(conj(c_left),      c_right)]]


# In[4]:


def phase_gate(alpha0=2, T=200, samples_factor=1):
    assert np.abs(alpha0) < 2.1
    N = 40
    radius = 0.5**0.5
    samples = int(15000*samples_factor)

    unit_interval = np.linspace(0,1,samples)
    centered_interval = np.linspace(-1,1,samples)
    one = np.ones(samples)
    zero = np.zeros(samples)
    gaussian = np.exp(-centered_interval**2)

    a_op = qutip.destroy(N)
    adag_op = qutip.create(N)
    id_op = qutip.identity(N)
    zero_op = qutip.zero_oper(N)
    num_op = qutip.num(N)

    beta0 = -alpha0
    init_state = (qutip.coherent(N,alpha0) + qutip.coherent(N,beta0)).unit()
    alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)
    betas  = beta0 * np.ones_like(alphas)

    c_ops = expand_lindblad_operator_to_qutip_std_form(
                a_op**2,
                [-a_op, alphas+betas],
                [id_op, alphas*betas]
            )
    
    ret = qutip.mesolve(zero_op,init_state,unit_interval*T,
                        c_ops=c_ops,
                        e_ops=[],
                        progress_bar=qutip.ui.progressbar.TextProgressBar())
    return ret.states[::100]


# In[5]:


with open("results_animate_phase_gate", "wb") as f:
    pickle.dump(phase_gate(), f)


# In[6]:


with open("results_animate_phase_gate", "rb") as f:
    phase_gate_results = pickle.load(f)


# In[7]:


qutip.WignerDistribution(phase_gate_results[0]).visualize()


# In[8]:


qutip.WignerDistribution(phase_gate_results[-1]).visualize()


# In[9]:


def pop_gate(alpha0=2, T=200, samples_factor=1):
    assert np.abs(alpha0) < 2.1
    N = 40
    samples = int(60000*samples_factor)
   
    unit_interval = np.linspace(0,1,samples)
    unit_interval_radius = np.linspace(0,1,int(samples/(2+2*np.pi/6)))
    unit_interval_perim = np.linspace(0,1,samples-2*int(samples/(2+2*np.pi/6)))
    one = np.ones(samples)
    zero = np.zeros(samples)

    a_op = qutip.destroy(N)
    adag_op = qutip.create(N)
    id_op = qutip.identity(N)
    zero_op = qutip.zero_oper(N)
    num_op = qutip.num(N)

    beta0 = -alpha0
    init_state = (qutip.coherent(N,alpha0)).unit()
    
    alphas_1 = unit_interval_radius[::-1]*alpha0
    alphas_2 = unit_interval_radius*alpha0*np.exp(1j*2*np.pi/8)
    alphas_3 = alpha0*np.exp(unit_interval_perim[::-1]*1j*2*np.pi/8)
    alphas = np.concatenate([alphas_1, alphas_2, alphas_3])
    betas = -alphas
    
    c_ops = expand_lindblad_operator_to_qutip_std_form(
                a_op**2,
                [-a_op, alphas+betas],
                [id_op, alphas*betas]
            )
    
    ret = qutip.mesolve(zero_op,init_state,unit_interval*T,
                        c_ops=c_ops,
                        e_ops=[],
                        progress_bar=qutip.ui.progressbar.TextProgressBar())
    return ret.states[::400]


# In[10]:


import pickle
with open("results_animate_pop_gate", "wb") as f:
    pickle.dump(pop_gate(), f)


# In[11]:


with open("results_animate_pop_gate", "rb") as f:
    pop_gate_results = pickle.load(f)


# In[12]:


qutip.WignerDistribution(pop_gate_results[0]).visualize()


# In[13]:


qutip.WignerDistribution(pop_gate_results[-1]).visualize()


# In[48]:


def plot_frame(states, alphas, betas, i):
    f = plt.figure(figsize=(5,5))
    s = f.add_subplot(1,1,1)
    
    plt.setp(s.get_xticklabels(), visible=False)
    plt.setp(s.get_yticklabels(), visible=False)

    s.set_xlabel(r'$Re(\alpha)$')
    s.set_ylabel(r'$Im(\alpha)$')

    s.set_xlim(-3.6,3.6)
    s.set_ylim(-3.6,3.6)

    s.plot(np.real(alphas[:i]),np.imag(alphas[:i]), 'k', alpha=0.25, linewidth=2)
    s.plot(np.real(betas [:i]),np.imag(betas [:i]), 'k', alpha=0.25, linewidth=2)
    s.plot(np.real(alphas[i]),np.imag(alphas[i]), 'go')
    s.plot(np.real(betas [i]),np.imag(betas [i]), 'go')

    ra, ia = np.linspace(-3.6,3.6,100),np.linspace(-3.6,3.6,100)
    w = qutip.wigner(states[i], ra, ia, g=2)
    s.contourf(ra, ia, w, 100, vmin=-0.3, vmax=0.3, cmap=mpl.cm.get_cmap('RdBu'))

    return f


# In[49]:


alpha0 =2
beta0 = -alpha0
radius = 0.5**0.5

samples = 15000
unit_interval = np.linspace(0,1,samples)
alphas = alpha0 + radius*(1-np.cos( unit_interval*2*np.pi)) - radius*1j*np.sin(unit_interval*2*np.pi)
betas  = beta0 * np.ones_like(alphas)
phase_alphas = alphas[::100]
phase_betas  = betas [::100]

samples = 60000
unit_interval = np.linspace(0,1,samples)
unit_interval_radius = np.linspace(0,1,int(samples/(2+2*np.pi/6)))
unit_interval_perim = np.linspace(0,1,samples-2*int(samples/(2+2*np.pi/6)))
alphas_1 = unit_interval_radius[::-1]*alpha0
alphas_2 = unit_interval_radius*alpha0*np.exp(1j*2*np.pi/8)
alphas_3 = alpha0*np.exp(unit_interval_perim[::-1]*1j*2*np.pi/8)
alphas = np.concatenate([alphas_1, alphas_2, alphas_3])
betas = -alphas
pop_alphas = alphas[::400]
pop_betas  = betas [::400]


# In[50]:


for i in range(150):
    f = plot_frame(phase_gate_results, phase_alphas, phase_betas, i)
    display.clear_output(wait=True)
    display.display(f)
    f.savefig('2blob_phase_gate_%05d.png'%i)
    plt.close('all')


# In[51]:


for i in range(150):
    f = plot_frame(pop_gate_results, pop_alphas, pop_betas, i)
    display.clear_output(wait=True)
    display.display(f)
    f.savefig('2blob_pop_gate_%05d.png'%i)
    plt.close('all')


# In[59]:


get_ipython().system('ffmpeg -y -r 15 -i 2blob_phase_gate_%05d.png  -c:v libx264 -r 30 -pix_fmt yuv420p 2blob_phase_gate.mp4 2>/dev/null')


# In[60]:


get_ipython().system('ffmpeg -y -r 15 -i 2blob_pop_gate_%05d.png  -c:v libx264 -r 30 -pix_fmt yuv420p 2blob_pop_gate.mp4 2>/dev/null')


# In[ ]: