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

# # The Dissipator of a Two-Level-System as an Explicit Superoperator

# We derive the analytical dissipation superoperator for decay and dephasing in the two-level system

# In[1]:


from qnet.algebra.operator_algebra import *
from qnet.algebra.hilbert_space_algebra import *
from qnet.algebra.state_algebra import *
from qnet.convert.to_sympy_matrix import convert_to_sympy_matrix
import sympy
import numpy as np
from sympy import sqrt, symbols


# In[2]:


sympy.init_printing()


# We define our symbols as follows:

# In[3]:


hs = LocalSpace('TLS', basis=(0,1))


# In[4]:


γ1, γ2 = symbols('gamma_1 gamma_2', positive=True)
ρ00, ρ11 = symbols('rho_00, rho_11', positive=True)
ρ01, ρ10 = symbols('rho_01, rho_10')
half = sympy.S(1) / 2


# In[5]:


L1 = sqrt(γ1) * LocalSigma(0, 1, hs=hs)
L2 = sqrt(2*γ2) * LocalSigma(1, 1, hs=hs)
L = [L1, L2]


# In[6]:


ket0, ket1 = BasisKet(0, hs=hs), BasisKet(1, hs=hs)
bra0, bra1 = Bra(ket0), Bra(ket1)


# In[7]:


def ketbra(i, j):
    k = {0: ket0, 1: ket1}
    return k[i] * Bra(k[j])


# In[8]:


basis_rho = [ # concatenation of columns!
    ketbra(0, 0),
    ketbra(1, 0),
    ketbra(0, 1),
    ketbra(1, 1)
]
rho = ρ00 * ketbra(0, 0) + ρ01 * ketbra(0, 1) + ρ10 * ketbra(1, 0) + ρ11 * ketbra(1, 1)


# The dissipator in Lindblad form is defined as

# In[9]:


def D(Ls, rho):
    res = 0
    for L in Ls:
        L_dag_L = L.dag() * L
        res += (L * rho * L.dag() - half * (L_dag_L * rho + rho * L_dag_L)).expand()
    return res


# Applied to a general density matrix this is:

# In[10]:


convert_to_sympy_matrix(D(L, rho))


# We now calculate the explicit form of the dissipator as a superoperator

# In[11]:


def hs_prod(A, B):
    """Hilbert-Schmidt product"""
    return tr((A.dag() * B).expand(), over_space=hs)


# In[12]:


def get_DSup():
    DSup = sympy.Matrix(np.zeros((4,4)))
    for i in range(4):
        for j in range(4):
            d = hs_prod(basis_rho[i], D(L, basis_rho[j]))
            if d != 0:
                assert d.term == IdentityOperator
                DSup[i, j] = d.coeff
            else:
                DSup[i, j] = sympy.S(0)
    return DSup


# In[13]:


DSup = get_DSup()


# We now have the desired superoperator

# In[14]:


DSup


# Let's test this by applying the above matrix to a vectorized form of a general density matrix (concatenation of columns):

# In[15]:


rho_vec = sympy.Matrix([
    [ρ00,],
    [ρ10,],
    [ρ01,],
    [ρ11,]])


# In[16]:


(DSup * rho_vec).expand()


# This is equal to the vectorized form of the previous result

# In[17]:


convert_to_sympy_matrix(D(L, rho)).transpose().reshape(4,1)


# We can also use this to solve the equation of motion:

# In[18]:


t = symbols('t', positive=True)


# In[19]:


(DSup * t).exp() * rho_vec