import qutip as qt
import numpy as np
from qutip import QobjEvo
%load_ext cython

N = 5
destroy, create, Id = qt.destroy(N), qt.create(N), qt.qeye(N)
def exp_i(t,args):
    return np.exp(-1j*t)
def cos_w(t,args):
    return np.cos(args["w"]*t)
tlist = np.linspace(0,10,10000)
tlistlog = np.logspace(-3,1,10000)

# state vector as np array
vec = np.arange(N)*.5+.5j
vec_super = np.arange(N**2)*.5+.5j
mat_c = (np.arange(N**2)*.5+.5j).reshape((5,5))
mat_f = np.asfortranarray(mat_c*1.)

# Construct QobjEvo of all type
td_cte1 = QobjEvo(Id)
td_cte2 = QobjEvo([Id])

td_func = QobjEvo([Id,[create,exp_i],[destroy,cos_w]],args={"w":2})
td_str = QobjEvo([Id,[create,"exp(-1j*t)"],[destroy,"cos(w*t)"]],args={'w':2.})
td_array = QobjEvo([Id,[create,np.exp(-1j*tlist)],[destroy,np.cos(2*tlist)]],tlist=tlist)
td_array_log = QobjEvo([Id,[create,np.exp(-1j*tlistlog)],[destroy,np.cos(2*tlistlog)]],tlist=tlistlog)

td_super = qt.liouvillian(td_func, c_ops=td_cte1)

"""
Explanation: QobjEvo "advanced" methods
Test and present diverse feature of QobjEvo.
Made by Eric Giguere
End of explanation
"""


from qutip.cy.spmatfuncs import spmv, cy_expect_rho_vec, cy_expect_psi

spmv(td_func(2).data, vec) == td_func.mul_vec(2,vec)

print(td_func(2).data * mat_c == td_func.mul_mat(2,mat_c))
mat_c.flags

print(td_func(2).data * mat_f == td_func.mul_mat(2,mat_f))
mat_f.flags

cy_expect_psi(td_str(2).data, vec, 0) == td_str.expect(2, vec, 0)

cy_expect_rho_vec(td_super(2).data, vec_super, 0) == td_super.expect(2, vec_super, 0)

"""
Explanation: Product and expectation value
QobjEvo.mul_vec(t,state) = spmv(QobjEvo(t), state)
QobjEvo.expect(t, state, real) = cy_expect_psi/cy_expect_rho_vec  (QobjEvo(t), state, real)
End of explanation
"""


def coeff_state(t, args):
    return np.max(args["psi"])*args["w"]
td_state = QobjEvo([Id, [destroy, coeff_state]],args={"w":1, "psi=vec":qt.basis(N,0)})

td_state(2,state=vec)

"""
Explanation: Function with state
For the solvers which have the option "rhs_with_state"
the QobjEvo can take coefficient function with the signature (for backward compatibility):
def coeff(t, psi. args)
or use advanced args: args={"psi=vec":psi0}
End of explanation
"""


td_str.compiled = False
print("Before compilation")
%timeit td_str(2, data=True)
%timeit td_str.mul_vec(2,vec)
%timeit td_str.mul_mat(2,mat_c)
%timeit td_str.mul_mat(2,mat_f)
%timeit td_str.expect(2,vec,0)
td_str.compile()
print("After compilation")
%timeit td_str(2, data=True)
%timeit td_str.mul_vec(2,vec)
%timeit td_str.mul_mat(2,mat_c)
%timeit td_str.mul_mat(2,mat_f)
%timeit td_str.expect(2,vec,0)

"""
Explanation: Cython object and "Compiling"
There is a cython version of the qobjevo for fast call:
- qutip.cy.cqobjevo.CQobjEvo
The cython is created when the "compile" method is created.
The cython object contain fast version of the call, expect and rhs (spmv) methods.
End of explanation
"""


def multiply(qobj, b, factor = 3.):
    return qobj*b*factor

print(td_func.apply(multiply,2)(2) == td_func(2)*6)
print(td_func.apply(multiply,2,factor=2)(2) == td_func(2)*4)

"""
Explanation: apply
Pass a function ( Qobj, *args, **kwargs) -> Qobj to act on each component of the qobjevo.
Will only be mathematicaly valid if the transformation is linear.
End of explanation
"""


def rescale_time_and_scale(f_original, time_scale, factor=2.):
    def f(t, *args, **kwargs):
        return f_original(time_scale*t, *args, **kwargs)*factor
    return f

print(td_func.apply_decorator(rescale_time_and_scale, 2)(2) == td_func(4)*2-Id)
print(td_func.apply_decorator(rescale_time_and_scale, 3, factor=3)(2) == 
      td_func(6)*3.0 - 2*Id)

"""
Explanation: apply_decorator
Transform the functions containing the time dependence using a decorator.
The decorator must return a function of (t, **kwargs).
Do not modify the constant part (the contant part do not have a function f(t) = 1).
End of explanation
"""


td_func_1 = QobjEvo([[create,exp_i]],args={"w":2})
td_str_1 = QobjEvo([[create,"exp(-1j*t)"]],args={'w':2.})
td_array_1 = QobjEvo([[create,np.exp(-1j*tlist)]],tlist=tlist)

def square_qobj(qobj):
    return qobj*qobj

def square_f(f_original):
    def f(t, *args, **kwargs):
        return f_original(t, *args, **kwargs)**2
    return f

t1 = td_func_1.apply(square_qobj).apply_decorator(square_f)
print(t1(2) == td_func_1(2)*td_func_1(2))
print((t1.ops[0][2]))

t1 = td_str_1.apply(square_qobj).apply_decorator(square_f)
print(t1(2) == td_str_1(2)*td_str_1(2))
print("str not updated:", (t1.ops[0][2]))

t1 = td_str_1.apply(square_qobj).apply_decorator(square_f, str_mod=["(",")**2"])
print(t1(2) == td_str_1(2)*td_str_1(2))
print("str  updated:",(t1.ops[0][2]))

t1 = td_array_1.apply(square_qobj).apply_decorator(square_f)
print(t1(2) == td_array_1(2)*td_array_1(2))
print("array not updated:",(t1.ops[0][2]))

t1 = td_array_1.apply(square_qobj).apply_decorator(square_f, inplace_np=1)
print(t1(2) == td_array_1(2)*td_array_1(2))
print("array  updated:",(t1.ops[0][2]))

"""
Explanation: QobjEvo based on string and np.array are changed to a function then the decorator is applied. There are option so that the type of coefficient stay unchanged: 
str_mod : change the string -&gt; str_mod[0] + str + str_mod[1]

inplace_np : modify the array  (array[i] = decorator(lambda v: v)(array[i]))
     *any modification that rescale the time will not work properly

Decorator can cause problem when used in parallel. (function cannot be pickled error)
End of explanation
"""


small = qt.destroy(2)
def f1(t,args):
    return np.sin(t)
def f2(t,args):
    return np.cos(args["w"]*t)
def f3(t,args):
    return np.sin(args["w"]*t)
def f4(t,args):
    return np.cos(t)

td_redoundance = QobjEvo([qt.qeye(2),[small,"sin(t)"],[small,"cos(w*t)"],[small,"sin(w*t)"],
                                  [small,"cos(t)"]],args={'w':2.})
td_redoundance1 = QobjEvo([qt.qeye(2),[small,"sin(t)"],[small,"cos(w*t)"],[small,"sin(w*t)"],
                                   [small,"cos(t)"]],args={'w':2.})
td_redoundance2 = QobjEvo([qt.qeye(2),[small,f1],[small,f2],[small,f3],[small,f4]],args={'w':2.})
td_redoundance3 = QobjEvo([qt.qeye(2),[small,np.sin(tlist)],[small,np.cos(2*tlist)],
                                   [small,np.sin(2*tlist)],[small,np.cos(tlist)]],tlist=tlist)
td_redoundance4 = QobjEvo([qt.qeye(2),[small,f1],[small,"cos(w*t)"],
                                   [small,np.sin(2*tlist)],[small,"cos(t)"]],args={'w':2.},tlist=tlist)

td_redoundance1.compress()
print(td_redoundance1.to_list())
print(len(td_redoundance1.ops))
print(td_redoundance(1.) == td_redoundance1(1.))
td_redoundance2.compress()
print(len(td_redoundance2.ops))
print(td_redoundance(1.) == td_redoundance2(1.))
td_redoundance3.compress()
print(len(td_redoundance3.ops))
print(td_redoundance(1.) == td_redoundance3(1.))
td_redoundance4.compress()
print(len(td_redoundance4.ops))
print(td_redoundance(1.) == td_redoundance4(1.))

td_redoundance_v2 = QobjEvo([qt.qeye(2),[qt.qeye(2),"sin(t)"],[small,"sin(t)"],[qt.create(2),"sin(t)"]])
td_redoundance_v2.compress()
td_redoundance_v2.to_list()

"""
Explanation: Removing redundant Qobj
When multiple components of the QobjEvo are made from the same Qobj, you can unite them with the "compress" method.
It is only done with the same form of time dependance:
End of explanation
"""


%%cython
from qutip.cy.cqobjevo cimport CQobjEvo
cimport numpy as np


def rhs_call_from_cy(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, np.ndarray[complex, ndim=1] out):
    qobj._mul_vec(t,&vec[0],&out[0])
    
    
def expect_call_from_cy(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, int isherm):
    return qobj._expect(t,&vec[0])
    
    
def rhs_cdef_timing(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, np.ndarray[complex, ndim=1] out):
    cdef int i
    for i in range(10000):
        qobj._mul_vec(t,&vec[0],&out[0])

        
def expect_cdef_timing(CQobjEvo qobj, double t, np.ndarray[complex, ndim=1] vec, int isherm):
    cdef complex aa = 0.
    cdef int i
    for i in range(10000):
        aa = qobj._expect(t, &vec[0])
    return aa

def rhs_def_timing(qobj, double t, np.ndarray[complex, ndim=1] vec, complex[::1] out):
    cdef int i
    for i in range(10000):
        out = qobj.mul_vec(t,vec)
        
def expect_def_timing(qobj, double t, np.ndarray[complex, ndim=1] vec, int isherm):
    cdef complex aa = 0.
    cdef int i
    for i in range(10000):
        aa = qobj.expect(t, vec)
    return aa


td_str.compile()
print(expect_call_from_cy(td_str.compiled_qobjevo, 2, vec, 0) - td_str.expect(2,vec,0))
%timeit expect_def_timing(td_str.compiled_qobjevo, 2, vec, 0)
%timeit expect_cdef_timing(td_str.compiled_qobjevo, 2, vec, 0)

out = np.zeros(N,dtype=np.complex128)
rhs_call_from_cy(td_str.compiled_qobjevo, 2, vec, out)
print( [a - b for a,b in zip(out, td_str.mul_vec(2,vec))])
%timeit rhs_def_timing(td_str.compiled_qobjevo, 2, vec, out)
%timeit rhs_cdef_timing(td_str.compiled_qobjevo, 2, vec, out)
# Most of the time gained is from allocating the out vector, not the 

td_cte = QobjEvo([Id])
td_cte.compile()
out = np.zeros(N,dtype=np.complex128)
rhs_call_from_cy(td_cte.compiled_qobjevo, 2, vec, out)
print( [a - b for a,b in zip(out, td_cte.mul_vec(2,vec))])
%timeit rhs_def_timing(td_cte.compiled_qobjevo, 2, vec, out)
%timeit rhs_cdef_timing(td_cte.compiled_qobjevo, 2, vec, out)

"""
Explanation: cimport cython object
The cython object can be 'cimported'. 
cdef class CQobjEvo:
    cdef void _mul_vec(self, double t, complex* vec, complex* out)
    cdef void _mul_matf(self, double t, complex* mat, complex* out,int nrow, int ncols)
    cdef void _mul_matc(self, double t, complex* mat, complex* out,int nrow, int ncols)
    cdef complex _expect(self, double t, complex* vec, int isherm)
    cdef complex _expect_super(self, double t, complex* rho, int isherm)
End of explanation
"""


td_str.compiled = False
print(td_str.compile(code=True))

qt.about()

"""
Explanation: Compiled string code
End of explanation
"""


