# Set up the cluster for parallel computation
from IPython.parallel import Client
client = Client(profile='lsf')

dview = client[:]
dview.use_dill()
lview = client.load_balanced_view();

%%px --local

%matplotlib inline
import qutip
import numpy as np
import scipy
import matplotlib as mpl
import matplotlib.pyplot as plt
from IPython import display
import sys
import time

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)]]

%%px --local
def do_all_for_given_circle_size(radius):
    N = 40
    samples = 1600
    T = 200

    title = 'two_blobs_two_circles_%.3f'%radius

    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)

    alpha0 =  2 + 0j
    beta0 = -2 + 0j
    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  - radius*(1-np.cos(-unit_interval*2*np.pi)) + radius*1j*np.sin(-unit_interval*2*np.pi)

    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())
    
    def animate_states(states, paths=[], every=5, extent=[[-4,4],[-4,4]], title=''):
        plt.close('all')
        for i, s in enumerate(states[::every]):
            f, a = qutip.WignerDistribution(s, extent=np.array(extent)*np.sqrt(2)).visualize()
            for p in paths:
                pp = p*np.sqrt(2)
                a.plot(np.real(pp),np.imag(pp),'gx-')
                a.plot(np.real(pp[i*every]),np.imag(pp[i*every]),'ro')
            f.savefig(title+'_ani_%06d.png'%(i*every))
            print('%d/%d'%(i+1,len(states[::every])))
            sys.stdout.flush()
        plt.close('all')

    def plot_xyz(ret, alphas, betas, title=''):
        plt.close('all')
        proj_pop_coh = [(
                         qutip.expect(s,qutip.coherent_dm(N,a)-qutip.coherent_dm(N,b)),
                         qutip.expect(s,qutip.coherent(N,a)*qutip.coherent(N,b).dag())
                        )
                        for s, a, b in zip(ret.states, alphas, betas)]
        proj_pop_coh = np.array(proj_pop_coh)
        z = proj_pop_coh[:,0]
        x = 2*np.real(proj_pop_coh[:,1])
        y = 2*np.imag(proj_pop_coh[:,1])
        f = plt.figure()
        plt.plot(ret.times,np.real(z), label='z')
        plt.plot(ret.times,x, label='x')
        plt.plot(ret.times,y, label='y')
        plt.plot(ret.times,np.sqrt(np.abs(z)**2+x**2+y**2), label='norm(x,y,z)')
        plt.xlabel('Time')
        plt.title('x y z components of the logical qubit')
        plt.legend()
        f.savefig(title+'_xyz.png')
        return f

    def plot_purity(ret, title=''):
        plt.close('all')
        f = plt.figure()
        plt.plot(ret.times,[(s**2).tr() for s in ret.states])
        plt.title('Purity')
        plt.xlabel('Time')
        f.savefig(title+'_purity.png')
        return f

    animate_states(ret.states, paths=[alphas,betas], every=10, extent=[[-3,3],[-2,2]], title=title)
    plot_xyz(ret, alphas, betas, title=title)
    plot_purity(ret, title=title)
    return ret

radii = np.linspace(-0.5,1.,num=40)
results = lview.map(do_all_for_given_circle_size, radii)

result_vals = results.result

last_states = [r.states[-1] for r in result_vals]

last_states_pop_coh =  [(
                         qutip.expect(s,qutip.coherent_dm(40,2)-qutip.coherent_dm(40,-2)),
                         qutip.expect(s,qutip.coherent(40,2)*qutip.coherent(40,-2).dag())
                        )
                        for s in last_states]

last_states_pop_coh = np.array(last_states_pop_coh)
z = last_states_pop_coh[:,0]
x = 2*np.real(last_states_pop_coh[:,1])
y = 2*np.imag(last_states_pop_coh[:,1])

last_purity = [(r.states[-1]**2).tr() for r in result_vals]