import numpy as np
import matplotlib.pyplot as plt


def first_pk(y):
    th = 0.08 * np.max(y)
    for i in range(1, len(y) - 1):
        if y[i] > th and y[i] >= y[i - 1] and y[i] >= y[i + 1]:
            return i
    ix = np.flatnonzero(y > th)
    return int(ix[0]) if len(ix) else 0


def wave(z):
    ps = np.empty(z.shape, dtype=float)
    for i in range(z.shape[1]):
        re = z[:, i].real
        im = z[:, i].imag
        y  = re if np.max(np.abs(re)) >= np.max(np.abs(im)) else im
        k  = first_pk(np.abs(y))
        sg = -1 if y[k] < 0 else 1
        y  = sg * y
        ps[:, i] = y / max(np.max(np.abs(y)), 1e-12)
    return ps


def solve(ej=12.5, ec=0.25, ng=0.0, ncut=25, nlev=6, nphi=401):
    n = np.arange(-ncut, ncut + 1)
    h = np.diag(4 * ec * (n - ng) ** 2)
    for i in range(len(n) - 1):
        h[i, i + 1] = h[i + 1, i] = -ej / 2

    e, v = np.linalg.eigh(h)
    e, v = e[:nlev], v[:, :nlev]

    ph = np.linspace(-np.pi, np.pi, nphi)
    z  = np.exp(1j * np.outer(ph, n)) @ v / np.sqrt(2 * np.pi)
    ps = wave(z)

    u = -ej * np.cos(ph)
    return ph, u, e, ps


def plot(ej=12.5, ec=0.25, ng=0.0):
    ph, u, e, ps = solve(ej, ec, ng)
    y = e - e[0]
    u = u - e[0]
    amp = 2 * max(0.28, min(0.9, 0.035 * (max(u) - min(u))))

    fig, ax = plt.subplots(figsize=(6.4, 4.0), constrained_layout=True)
    ax.plot(ph, u, color="black", lw=1.6)
    for i, yy in enumerate(y):
        ax.axhline(yy, color="0.72", lw=0.8)
        ax.fill_between(ph, yy, yy + amp * ps[:, i], alpha=0.18)
        ax.plot(ph, yy + amp * ps[:, i], lw=1.4)

    ax.set_xlim(-np.pi, np.pi)
    ax.set_ylim(min(u) - 0.8, max(max(u), y[-1]) + 0.8)
    ax.set_xlabel(r"$\varphi$")
    ax.set_ylabel(r"$(E-E_0)/h$ (GHz)")
    ax.set_xticks([-np.pi, -np.pi / 2, 0, np.pi / 2, np.pi])
    ax.set_xticklabels([r"$-\pi$", r"$-\pi/2$", "0", r"$\pi/2$", r"$\pi$"])
    return fig, ax


if __name__ == "__main__":
    fig, ax = plot(ej=12.5, ec=0.25, ng=0.0)
    fig.savefig("transmon_levels.png", dpi=180)