from plot import *
import scqubits as scq
import qutip as qt

# flux = scq.FluxQubit()
fluxonium = scq.Fluxonium(EJ=9, EC=3, EL=0.5, flux=1, cutoff=100)



def _plot_transmon_n_wavefunctions(qubit: scq.Transmon, fig_ax, which=[0,1]):
    fig, ax = fig_ax
    # ax.set_ylabel(r"$|\psi_i(n)|^2$")
    ax.set_ylabel(r"$\psi_i(n)$")
    ax.set_xlabel(r"$n$")
    colors = "brgy"
    for i in which:
        wf = qubit.numberbasis_wavefunction(which=i)
        x = wf.basis_labels
        y = wf.amplitudes
        offset = (len(which)/2 - i) * 0.2
        ax.bar(x-offset, y, width=0.2, align='center', label=f"$i={i}$")
    xlim = (-4, 4)
    ax.set_xlim(*xlim)
    ax.set_xticks(np.arange(xlim[0], xlim[1]+1))

def _plot_transmon(qubit: scq.Transmon, ngs, fig, axs):
    _,_ = qubit.plot_evals_vs_paramvals("ng", ngs, fig_ax=(fig, axs[0]), evals_count=5, subtract_ground=False)
    _,_ = qubit.plot_wavefunction(fig_ax=(fig, axs[1]), which=[0, 1, 2], mode="abs_sqr")
    _plot_transmon_n_wavefunctions(qubit, (fig, axs[2]), which=[0, 1, 2])
    qubit.ng = 0.5
    _plot_transmon_n_wavefunctions(qubit, (fig, axs[3]), which=[0, 1, 2])
    qubit.ng = 0


def transmon_cpb():
    EC = 1
    qubit = scq.Transmon(EJ=30, EC=EC, ng=0, ncut=30)
    ngs = np.linspace(-2, 2, 200)
    fig, axs = plt.subplots(4, 3, squeeze=True, figsize=(full,full))
    axs = axs.T
    qubit.ng = 0
    qubit.EJ = 0.1 * EC
    title = lambda x: f"$E_J/E_C = {x}$"
    _plot_transmon(qubit, ngs, fig, axs[0])
    axs[0][0].set_title("Cooper-Pair-Box\n"+title(qubit.EJ))

    qubit.EJ = EC
    _plot_transmon(qubit, ngs,  fig, axs[1])
    axs[1][0].set_title("Quantronium\n"+title(qubit.EJ))

    qubit.EJ = 20 * EC
    _plot_transmon(qubit, ngs, fig, axs[2])
    axs[2][0].set_title("Transmon\n"+title(qubit.EJ))

    for ax in axs[1:,:].flatten(): ax.set_ylabel("")
    for ax in axs[:,0].flatten():
        ax.set_xticks([-2, -1, -0.5, 0, 0.5, 1, 2])
        ax.set_xticklabels(["-2", "-1", "", "0", "", "1", "2"])
        ylim = ax.get_ylim()
        ax.vlines([-1, -0.5], ymin=ylim[0], ymax=ylim[1], color="#aaa", linestyle="dotted")
    axs[0][2].legend()
    fig.tight_layout()
    return fig

export(transmon_cpb(), "qubit_transmon")


def flux_onium():
    fig, axs = plt.subplots(1, 3, squeeze=True, figsize=(full,full/2))
    fluxs = np.linspace(0.4, 0.6, 50)
    EJ = 35.0
    alpha = 0.3
    # todo find useful parameters
    fluxqubit = scq.FluxQubit(
                                  EJ1 = EJ,
                                  EJ2 = EJ,
                                  EJ3 = alpha*EJ,
                                  ECJ1 = 1.0,
                                  ECJ2 = 1.0,
                                  ECJ3 = 1.0/alpha,
                                  ECg1 = 50.0,
                                  ECg2 = 50.0,
                                  ng1 = 0.0,
                                  ng2 = 0.0,
                                  flux = 1,
                                  ncut = 10)
    fluxqubit.plot_evals_vs_paramvals("flux", fluxs, evals_count=5, subtract_ground=False, fig_ax=(fig, axs[0]))
    axs[0].set_title(f"Flux {texvar('alpha', alpha)}")

    alpha = 0.8
    fluxqubit.EJ3 = alpha * EJ
    fluxqubit.ECJ3 = 1.0/alpha
    fluxqubit.plot_evals_vs_paramvals("flux", fluxs, evals_count=5, subtract_ground=False, fig_ax=(fig, axs[1]))
    axs[1].set_title(f"Flux {texvar('alpha', alpha)}")
    # axs[0].set_xlim(0.4, 0.6)

    fluxs = np.linspace(-1.1, 1.1, 101)
    fluxonium = scq.Fluxonium(EJ=9, EC=3, EL=0.5, flux=1, cutoff=100)
    fluxonium.plot_evals_vs_paramvals("flux", fluxs, evals_count=5, subtract_ground=True, fig_ax=(fig, axs[2]))
    axs[2].set_title("Fluxonium")
    return fig

export(flux_onium(), "qubit_flux_onium")