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

# In[1]:


from qutip import * 
import numpy as np
import matplotlib.pyplot as plt


# In[42]:


Nph = 10
om_c = 1.
U = 0.2

a = destroy(Nph)

H = om_c*a.dag()*a + U/2.*a.dag()*a.dag()*a*a

alpha = 1.
psi0 = coherent(Nph, alpha)
tlist = np.linspace(0, 4.*2*np.pi/om_c, 100)
psi_t = sesolve(H, psi0, tlist, [])
Navg = expect(a.dag()*a, psi_t.states)
A = expect(a+a.dag(), psi_t.states)

plt.figure(figsize=(8,8))
plt.plot( tlist/(2*np.pi), A  )
plt.grid()
#plt.axis([0, range(M), 0, 10])
font = {'family': 'serif',
        'color':  'darkred',
        'weight': 'normal',
        'size': 26,
        }
plt.xlabel('t/2pi', fontdict=font)
plt.ylabel('Photon Number', fontdict=font)
plt.show()


# In[ ]:





# In[ ]:


from matplotlib import cm
from matplotlib.ticker import LinearLocator

#fig, ax = plt.subplots(subplot_kw={"projection": "3d"})

# Make data.
Y = np.asarray(arr_omL)
X = np.asarray(arr_I)
X, Y = np.meshgrid(X, Y)
Z = np.asarray(mat_N_avg)
print(X.shape, Y.shape, Z.shape)

plt.pcolor(X, Y, Z)
plt.show()

# Plot the surface.
#surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                       #linewidth=0, antialiased=False)

# Customize the z axis.
#ax.set_zlim(0.0, 10.)
#ax.zaxis.set_major_locator(LinearLocator(10))
# A StrMethodFormatter is used automatically
#ax.zaxis.set_major_formatter('{x:.02f}')

# Add a color bar which maps values to colors.
#fig.colorbar(surf, shrink=0.5, aspect=5)

#plt.show()


# In[ ]:





# In[ ]:





# In[ ]:


"""    
       #Wigner function
       x_vec = np.arange(-5,5,.1)
       y_vec = x_vec
       Wf = wigner(rho_ph, x_vec, y_vec )
       wtitle = "wdata_%d" % countw
       p_w = open("data/wigner/%s.dat" % (wtitle), "w")
       nx = 0
       for x in x_vec:
           ny = 0
           for y in y_vec:
               p_w.write("%f %f %f\n" % (x, y, Wf[nx,ny]) )
               ny += 1
           p_w.write("\n")
           p_w.flush()
           nx += 1
       p_w.close()
       """
       
       
       #Spin Q fun
       n_sample = 40.
       phi_vec = np.arange(-np.pi, np.pi+2*np.pi/100., 2*np.pi/n_sample )
       theta_vec = np.arange(0, np.pi+np.pi/100., np.pi/n_sample )
       stitle = "sdata_%d" % countw
       p_s = open("data/wigner/%s.dat" % (stitle), "w")
       for theta in theta_vec:
           for phi in phi_vec:
               state_n = spin_coherent(1./2., theta, phi, type='ket')
               Qfun = expect( rho_s, state_n ) 
               p_s.write("%f %f %f\n" % (phi, theta, Qfun.real))
               p_s.flush()
           p_s.write("\n")
           p_s.flush()
           
"""
plt.figure(figsize=(8,8))
plt.plot( arr_I, N_avg  )
plt.grid()
#plt.axis([0, range(M), 0, 10])
font = {'family': 'serif',
       'color':  'darkred',
       'weight': 'normal',
       'size': 26,
       }
plt.xlabel('I', fontdict=font)
plt.ylabel('Photon Number', fontdict=font)
plt.show()
"""