import numpy as np

dt = 0.001  
T = 0.01  #solve p_dot from time 0 to T
t = np.linspace(0, T, int(T/dt)+1)   #discretized time
 #time step


D=0.005
J=30
g1=1.6
g2=1.6
a1=60
a2=60
b1=78.125*1j
b2=78.125*1j
c1=120
c2=120
eps=3

Id = np.array([[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]])
L = np.array([[0,-J,J,0],[-J,eps,0,J],[J,0,-eps,-J],[0,J,-J, 0]])
Phi1 = np.array([[0,0,0,0],[0,1j,0,0],[0,0,-1j,0],[0,0,0,0]])
Phi2 = np.array([[0,0,0,0],[0,-1j,0,0],[0,0,1j,0],[0,0,0,0]])
Theta1 = np.array([[c1,0,0,0],[0,1j*(a1-b1),0,0],[0,0,1j*(a1-b1),0],[0,0,0,0]])
Theta2 = np.array([[0,0,0,0],[0,1j*(-a2-b2),0,0],[0,0,1j*(a2-b2),0],[0,0,0,c2]])
g1Id = np.array([[g1,0,0,0],[0,g1,0,0],[0,0,g1,0],[0,0,0,g1]])
g2Id = np.array([[g2,0,0,0],[0,g2,0,0],[0,0,g2,0],[0,0,0,g2]])

def M(r,c,n):
    res=np.zeros((n,n))
    res[r-1,c-1]=1
    return res

M24=np.kron(M(1,1,6),Id-1j*D*L)+np.kron(M(1,2,6),Phi1*D)+np.kron(M(1,3,6),Phi2*D)+np.kron(M(2,1,6),Theta1*D)+np.kron(M(2,4,6),Phi2*D)+np.kron(M(2,6,6),Phi1*D)+np.kron(M(3,1,6),Theta2*D)+np.kron(M(3,3,6),Id-1j*D*L-g2*D)+np.kron(M(3,4,6),Phi1*D)+np.kron(M(3,5,6),Phi2*D)+np.kron(M(4,4,6),Id-1j*D*L)+np.kron(M(5,5,6),Id-1j*D*L)+np.kron(M(6,6,6),Id-1j*D*L)
sigm0 = np.array([1,0,0,0])

def f(sigm):
    M24=np.kron(M(1,1,6),Id-1j*D*L)+np.kron(M(1,2,6),Phi1*D)+np.kron(M(1,3,6),Phi2*D)+np.kron(M(2,1,6),Theta1*D)+np.kron(M(2,4,6),Phi2*D)+np.kron(M(2,6,6),Phi1*D)+np.kron(M(3,1,6),Theta2*D)+np.kron(M(3,3,6),Id-1j*D*L-g2*D)+np.kron(M(3,4,6),Phi1*D)+np.kron(M(3,5,6),Phi2*D)+np.kron(M(4,4,6),Id-1j*D*L)+np.kron(M(5,5,6),Id-1j*D*L)+np.kron(M(6,6,6),Id-1j*D*L)
    return np.dot(M24,sigm)

sigm = []  #an array storing the solution
for j in range(len(t)):
    if j !=0:
        s=(24,24)
        s_i=np.zeros(24,dtype=complex)
        sigm.append(s_i)
    else:
        sigm.append(sigm0)


sigm = _____  #Create an array of the first 4 elements of sigm and store them 
    

for i in range (1,len(t)):
    sigm[i]= _______ #Array of the first 4 elements of sigm

sigm11=[]
sigm22=[]
sigmtr=[]   
for k in range(len(sigm)):
    sigm11.append(p[k][0])#appending first element of matrix into a list called sigm11
    sigm22.append(p[k][3])
    sigmtr.append(p11+p22)
  
plt.plot(t,sigm11,'*')
plt.xlabel('time (fs)')
plt.ylabel('population')
plt.title('For $n_{1}+n_{2}\leq2$')