%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

from wavefunction.wavefunction2d import *

args = {'Ec': 1/35.0, 'Ej': 1.0, 'alpha': 0.8, 'f': 0.50, 'beta_L': 0.0}
globals().update(args)

# pick a truncation of the fourier series: n1 = [-L1, ..., L1], n2 = [-L2, ..., L2]
L1 = 10
L2 = 10

# pick periods for the coordinates
Tx1 = Tx2 = 2 * np.pi

#
k11, k12, k22 = 2 * Ec / Ej, 0.0, 2 / (1 + 2 * alpha) * Ec / Ej

K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)

def assemble_u_flux_qubit(L1, L2, args, sparse=False):

    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    u = np.zeros((L1n, L2n), dtype=np.complex)

    u[L1, L2] = 2 + alpha - beta_L/4

    u[+1+L1, +1+L2] = - 0.5
    u[-1+L1, +1+L2] = - 0.5
    u[+1+L1, -1+L2] = - 0.5
    u[-1+L1, -1+L2] = - 0.5

    u[L1, +2+L2] = -0.5 * alpha * np.exp(+2j * np.pi * f)
    u[L1, -2+L2] = -0.5 * alpha * np.exp(-2j * np.pi * f)

    u[+2+L1, +2+L2] = -beta_L/8.0
    u[-1+L1, -2+L2] = -beta_L/8.0

    return u

u = assemble_u_flux_qubit(L1, L2, args)

phi_p = np.linspace(-2*np.pi, 2*np.pi, 100)
phi_m = np.linspace(-2*np.pi, 2*np.pi, 100)
PHI_P, PHI_M = np.meshgrid(phi_p, phi_m)

U = evalute_fourier_series(PHI_P, PHI_M, L1, L2, u)

fig, axes = plt.subplots(2, 2, figsize=(10, 8))

for n, bL in enumerate([0.0, 1.0, 4.0, 10.0]):

    args['beta_L'] = bL
    
    u = assemble_u_flux_qubit(L1, L2, args)
    U = evalute_fourier_series(PHI_P, PHI_M, L1, L2, u)
    
    Z = np.real(U)
    p = axes[n // 2, n % 2].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n // 2, n % 2])
    axes[n // 2, n % 2].axis('tight')
    axes[n // 2, n % 2].set_xlabel(r'$\varphi_p$', fontsize=16)
    axes[n // 2, n % 2].set_ylabel(r'$\varphi_m$', fontsize=16)
    cb.set_label(r'$U(\varphi_p, \varphi_m)/E_J$', fontsize=16);

phi_p = np.linspace(-np.pi, np.pi, 100)
phi_m = np.linspace(-np.pi, np.pi, 100)
PHI_P, PHI_M = np.meshgrid(phi_p, phi_m)

args['beta_L'] = 0    
u = assemble_u_flux_qubit(L1, L2, args)

def assemble_V_flux_qubit(L1, L2, args, sparse=False):
    
    globals().update(args)
    
    L1n = 2 * L1 + 1
    L2n = 2 * L2 + 1
    
    V = np.zeros((L1n*L1n, L2n*L2n), dtype=np.complex)
    
    for n1 in np.arange(-L1, L1+1):
        for n2 in np.arange(-L1, L1+1):
            N = index_m2v(L1, n1, n2)
            for m1 in np.arange(-L2, L2+1):
                for m2 in np.arange(-L2, L2+1):
                    M = index_m2v(L2, m1, m2)

                    V[N,M] = (2 + alpha - 0.25 * beta_L) * delta(m1, n1) * delta(m2, n2) + \
                             - 0.5 * (delta(m1 + 1, n1) * delta(m2 + 1, n2) + 
                                      delta(m1 + 1, n1) * delta(m2 - 1, n2) + 
                                      delta(m1 - 1, n1) * delta(m2 + 1, n2) + 
                                      delta(m1 - 1, n1) * delta(m2 - 1, n2)) + \
                             - 0.5 * alpha * np.exp(+2j * np.pi * f) * delta(m1, n1) * delta(m2 + 2, n2) + \
                             - 0.5 * alpha * np.exp(-2j * np.pi * f) * delta(m1, n1) * delta(m2 - 2, n2) + \
                             - 0.125 * beta_L * delta(m1 + 2, n1) * delta(m2 + 2, n2) + \
                             - 0.125 * beta_L * delta(m1 - 2, n1) * delta(m2 - 2, n2)

    return V

V = assemble_V_flux_qubit(L1, L2, args)

H = K + V

vals, vecs = solve_eigenproblem(H)

Nstates = 4

fig, axes = plt.subplots(Nstates, 3, figsize=(12, 3 * Nstates))

for n in range(Nstates):

    psi = convert_v2m(L1, L2, vecs[n])
    
    Z = np.real(evalute_fourier_series(PHI_P, PHI_M, L1, L2, u))
    PSI = evalute_fourier_series(PHI_P, PHI_M, L1, L2, psi)
    
    p = axes[n, 0].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), np.real(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.real).max(), vmax=abs(PSI.real).max())
    c = axes[n, 0].contour(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 0])
    cb.set_clim(-5, 5)

    p = axes[n, 1].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), np.imag(PSI),
                          cmap=cm.RdBu, vmin=-abs(PSI.imag).max(), vmax=abs(PSI.imag).max())
    c = axes[n, 1].contour(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 1])
    cb.set_clim(-5, 5)
    
    p = axes[n, 2].pcolor(PHI_P/(2*np.pi), PHI_M/(2*np.pi), np.abs(PSI), cmap=cm.Blues, vmin=0, vmax=abs(PSI).max())
    c = axes[n, 2].contour(PHI_P/(2*np.pi), PHI_M/(2*np.pi), Z, cmap=cm.RdBu, vmin=Z.min(), vmax=Z.max())
    cb = fig.colorbar(p, ax=axes[n, 2])

    axes[n, 0].set_ylabel("%d state" % n);

axes[0, 0].set_title(r"$\mathrm{re}\;\Psi(\varphi_p, \varphi_m)$", fontsize=16);
axes[0, 1].set_title(r"$\mathrm{im}\;\Psi(\varphi_p, \varphi_m)$", fontsize=16);
axes[0, 2].set_title(r"$|\Psi(\varphi_p, \varphi_m)|^2$", fontsize=16);

f_vec = np.linspace(0.45, 0.55, 50)

e_vals = np.zeros((4, len(vals), len(f_vec)))

for n, bL in enumerate([0.0, 0.1, 0.5, 1.0]):

    args['beta_L'] = bL

    K = assemble_K(L1, L2, k11, k12, k22, Tx1, Tx2)

    for f_idx, f in enumerate(f_vec):
        args['f'] = f
        V = assemble_V_flux_qubit(L1, L2, args)
        H = K + V
        vals, vecs = solve_eigenproblem(H)
    
        e_vals[n, :, f_idx] = np.real(vals)

fig, axes = plt.subplots(2, 2, figsize=(10,10))

for n, bL in enumerate([0.0, 0.1, 0.5, 1.0]):

    idx = n // 2, n % 2
    
    axes[idx].plot(f_vec, e_vals[n, 0, :], 'b')
    axes[idx].plot(f_vec, e_vals[n, 2, :], 'r')
    for m in range(4, len(vals), 2):
        axes[idx].plot(f_vec, e_vals[n, m, :], 'k')
    
    axes[idx].axis('tight')
    axes[idx].set_ylim(e_vals[n, 0, :].min(), e_vals[n, 12, :].max())
    axes[idx].set_title(r"$\beta_L =$ %.1f" % bL)
    axes[idx].set_ylabel("Energy levels")
    axes[idx].set_xlabel(r'$f$');

%reload_ext version_information
%version_information numpy, scipy, matplotlib, wavefunction