# coding: utf-8

# In[1]:

#get_ipython().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

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[1]:


# In[2]:

def three_circles_2pop_gate_all_dance_one_blob_init_state(alpha0=2, T=200, samples_factor=1):
    assert np.abs(alpha0) < 3.1, 'alpha > 4 at any point in time goes out of the cutoff for the hilbert space'
    N = 70
    samples = int(800000*samples_factor)

    unit_interval = np.linspace(0,1,samples)
    unit_interval_radius = np.linspace(0,1,int(samples/(2+2*np.pi/12)))
    unit_interval_perim = np.linspace(0,1,samples-2*int(samples/(2+2*np.pi/12)))
    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*(-1/2 + 1j*3**0.5/2)
    gamma0 = alpha0*(-1/2 - 1j*3**0.5/2)
    init_state = (qutip.coherent(N,alpha0) + 2*qutip.coherent(N,beta0)).unit()
    alphas = alpha0*one
    betas_1  = 1j * unit_interval_radius[::-1] * beta0 .imag + beta0 .real
    gammas_1 = 1j * unit_interval_radius[::-1] * gamma0.imag + gamma0.real
    betas_2  = 1j * unit_interval_radius * beta0 .imag * np.exp(1j*2*np.pi/12) + beta0 .real
    gammas_2 = 1j * unit_interval_radius * gamma0.imag * np.exp(1j*2*np.pi/12) + gamma0.real
    betas_3  = 1j * beta0 .imag * np.exp(unit_interval_perim[::-1]*1j*2*np.pi/12) + beta0 .real
    gammas_3 = 1j * gamma0.imag * np.exp(unit_interval_perim[::-1]*1j*2*np.pi/12) + gamma0.real
    betas = np.concatenate([betas_1, betas_2, betas_3])
    gammas = np.concatenate([gammas_1, gammas_2, gammas_3])

    c_ops = expand_lindblad_operator_to_qutip_std_form(
                a_op**3,
                [-a_op**2, alphas+betas+gammas],
                [a_op, alphas*betas+betas*gammas+gammas*alphas],
                [-id_op, alphas*betas*gammas]
            )
    
    # XXX It is using too much RAM so I have to divide it
    # XXX Be sure that samples is divisible by 1000
    ret = []
    init_states = [init_state]
    for sep in range(samples)[::samples//10]:
        start = sep
        stop = sep+samples//10
        res = qutip.mesolve(zero_op,init_states[-1],unit_interval[:samples//10]*T,
                            c_ops=[c if not isinstance(c, list) else [c[0], c[1][start:stop]] for c in c_ops],
                            e_ops=[],
                            progress_bar=qutip.ui.progressbar.TextProgressBar())
        ret.extend(res.states[::200])
        init_states.append(res.states[-1])
        del res
    #return ret.states[-1], len(ret.states)
    return ret


# In[3]:


# In[ ]:

import pickle

def calculate_and_save(args):
    print(args)
    sys.stdout.flush()
    alpha0, T = args
    with open("results_three_circles_2pop_gate_all_dance_one_blob_init_state_%0.3f_%d"%(alpha0,T), "wb") as f:
        pickle.dump(three_circles_2pop_gate_all_dance_one_blob_init_state(alpha0,T), f)

parfor_args = [(2, 200), (2,220), (2,240), (2,260), (2,280), (2,300), (3.,200), (2.2,200), (2.4,200), (2.6,200), (2.8,200), (2.2,220), (2.4,240), (2.6,260), (2.8,280), (3,300)]
qutip.parfor(calculate_and_save, parfor_args)