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

# ## An extra short example:
# time evolution of a system of two coupled oscillators. We consider a few different Hamiltonian that describes a system of coupled oscillators. We initialize the system in a specific state (for instance coherent state for osc1 and vacuum for osc2) and we run the time evolution. A good (but more abstract) way to visualize what is going on is to make a plot of the time evolution directly un the harmonic oscillator basis (Fock states basis, instead of real space basis as before). In the simple case of two oscillators we can still visualize it as a 2-dimensional color map, and to visualize the time evolution we set a nice animation.

# In[1]:


from qutip import *
import numpy as np


N = 20

a1 = tensor(destroy(N), qeye(N))
a2 = tensor(qeye(N), destroy(N))

#H = a1.dag()*a1 + a2.dag()*a2 + .2*( a1*a2.dag() + (a1*a2.dag()).dag() )
H = a1.dag()*a1 + a2.dag()*a2 + .2*( a1*a2.dag() + (a1*a2.dag()).dag() ) + .04*a1*a1*a1.dag()*a1.dag()

#H = a1.dag()*a1 + a2.dag()*a2 + .2*(a1*a1*a2.dag() + (a1*a1*a2.dag()).dag() )
#H = a1.dag()*a1 + a2.dag()*a2 + .2*(a1*a1*a2.dag()*a2.dag() + (a1*a1*a2.dag()*a2.dag()).dag() )



psi0 = tensor(coherent(N, 2.), fock(N,0))
tlist = np.linspace(0,100, 50)
e_ops = [ a1.dag()*a1 ]
sol = sesolve(H, psi0, tlist)


# In[2]:


from matplotlib import cm
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import matplotlib.pyplot as plt
data = []
for nt in range(tlist.size):
    M = np.zeros(shape=(N-1,N-1))
    for n1 in range(N-1):
        for n2 in range(N-1):
            M[n1][n2] = np.abs( tensor(fock(N, n1), fock(N, n2) ).overlap(sol.states[nt]) )
    data.append( M )

    
fig, ax = plt.subplots()
plt.close(fig)
# Define the update function for the animation
def update(frame):
    ax.clear()  # Clear the previous frame
    ax.imshow(data[frame], cmap='viridis')  # Plot the current matrix
    ax.set_title(f'Frame {frame}')  # Set title for the current frame
    ax.axis('off')  # Turn off axis

# Create the animation
ani = FuncAnimation(fig, update, frames=tlist.size, interval=100)

# Display the animation as HTML
# HTML(ani.to_jshtml())
HTML(ani.to_html5_video())


# In[ ]: