# %%
from pathlib import Path
from functools import partial

import kwant
import matplotlib.pyplot as plt
import numpy as np
import skunk
import cmasher as cmr
import adaptive
import matplotlib

from multiterminal_invariant.common import system, zero_params, save_params
# %%
# Create scattering geometry: finite wire with two normal leads

# Number of sites of the finite wire
width_finite_NSN = 1000

# Number of sites of the leads before they come translational invariant
# This is so that we can introduce a potential barrier in the leads
width_leads_NSN = 500

width_NSN = width_finite_NSN + 2 * width_leads_NSN
x_lead_NSN = width_finite_NSN / 2

syst_NSN = system(width_NSN, (0,))

left_lead = system(1, (-1,))
left_lead = left_lead.substituted(V="V_left", Delta="Delta_left", eta="eta_left")
right_lead = system(1, (1,))
right_lead = right_lead.substituted(V="V_right", Delta="Delta_right", eta="eta_right")
syst_NSN.attach_lead(left_lead)  # normal lead
syst_NSN.attach_lead(right_lead)  # normal lead
sysf_NSN = syst_NSN.finalized()


# %%
def Delta_func(Delta, x_lead):
    def shape(x):
        return Delta if np.abs(x) < x_lead else 0

    return shape


# Gaussian potential barrier
def V_func(V, x0, sigma):
    def shape(x):
        return V * np.exp(-((x - x0) ** 2) / (2 * sigma**2))

    return shape


# %%
# Fig 2(a)
# Parameters for Majorana regime
params_MZM = {
    **zero_params(sysf_NSN),
    "mu": 0.1,
    "tx": 1,
    "Delta": 0.05,
    "alpha": 0.02,
    "Ez": 0.2,
    "L": width_finite_NSN,
    "sigma": 10,
}
save_params(params_MZM, "fig2_mzm")


# %%
# Compute invariant for different values of V and wire width
def invariant_Vs(V_power, width):
    V = 10**V_power
    smatrix = kwant.smatrix(
        sysf_NSN,
        energy=0,
        params={
            **params_MZM,
            "Delta": Delta_func(params_MZM["Delta"], width),
            "V": lambda x: V_func(V, -width, params_MZM["sigma"])(x)
            + V_func(V, width, params_MZM["sigma"])(x),
        },
    )
    r = smatrix.submatrix(0, 0)
    return np.linalg.det(r).real


# %%
widths = np.linspace(300, 1000, 10)
# %%
filename = "../data/fig2_mzm.npz"
if Path(filename).exists():
    data_MZM = np.load(filename)
else:
    Vs = []
    det_r1_mzm = []
    n_points = 50
    for width in widths:
        learner = adaptive.Learner1D(
            partial(invariant_Vs, width=width),
            bounds=(np.log10(0.07), np.log10(0.7)),
            loss_per_interval=adaptive.learner.learner1D.curvature_loss_function(),
        )
        runner = adaptive.runner.simple(learner, npoints_goal=n_points)
        data = learner.to_numpy()
        Vs.append(10 ** data[:, 0])
        det_r1_mzm.append(data[:, 1])
    Vs = np.array(Vs)
    det_r1_mzm = np.array(det_r1_mzm)
    np.savez(filename, det_r1_mzm=det_r1_mzm, Vs=Vs, widths=widths, **params_MZM)
    data_MZM = np.load(filename)

# %%
# Fig 2(b)
# mu >> Ez >> Delta for quasi-Majoranas
params_qMZM = {
    **zero_params(sysf_NSN),
    "mu": 0.12,
    "tx": 1,
    "Delta": 0.05,
    "alpha": 0.02,
    "Ez": 0.1,
    "V": 0.15,
    "sigma": 10,
    "L": width_finite_NSN,
}
save_params(params_qMZM, "fig2_qmzm")


# %%
# Compute invariant for different values of sigma and wire width
def invariant_sigmas(sigma_power, width):
    sigma = 10**sigma_power
    smatrix = kwant.smatrix(
        sysf_NSN,
        energy=0,
        params={
            **params_qMZM,
            "Delta": Delta_func(params_qMZM["Delta"], width),
            "V": lambda x: V_func(params_qMZM["V"], -width, sigma)(x)
            + V_func(params_qMZM["V"], width, sigma)(x),
        },
        check_hermiticity=False,
    )
    r = smatrix.submatrix(0, 0)
    return np.linalg.det(r).real


# %%
filename = "../data/fig2_qmzm.npz"
if Path(filename).exists():
    data_qMZM = np.load(filename)
else:
    det_r1_qmzm = []
    sigmas = []
    n_points = 50
    for width in widths:
        learner = adaptive.Learner1D(
            partial(invariant_sigmas, width=width),
            bounds=(np.log10(3e-2), np.log10(30)),
            loss_per_interval=adaptive.learner.learner1D.curvature_loss_function(),
        )
        runner = adaptive.runner.simple(learner, npoints_goal=n_points)
        data = learner.to_numpy()
        sigmas.append(10 ** data[:, 0])
        det_r1_qmzm.append(data[:, 1])
    sigmas = np.array(sigmas)
    det_r1_qmzm = np.array(det_r1_qmzm)
    np.savez(
        filename, det_r1_qmzm=det_r1_qmzm, sigmas=sigmas, widths=widths, **params_qMZM
    )
    data_qMZM = np.load(filename)
# %%
# Plot results
figwidth = plt.rcParams["figure.figsize"][0]

fig, axs = plt.subplot_mosaic(
    [["mzm", "qmzm"]],
    sharey=True,
    figsize=(figwidth, figwidth / 4.5),
)

cmap = cmr.bubblegum_r


np.testing.assert_allclose(params_qMZM["alpha"], params_MZM["alpha"])
lso = params_qMZM["tx"] / params_qMZM["alpha"]
ls = widths / lso
l_min = ls.min()
l_max = ls.max()

colors = cmap((ls - l_min) / (l_max - l_min))
for i, _ in enumerate(widths):
    axs["mzm"].plot(
        data_MZM["Vs"][i] / data_MZM["Delta"],
        data_MZM["det_r1_mzm"][i].real,
        c=colors[i],
    )

axs["mzm"].set_xscale("log")
axs["mzm"].set_xlabel(r"$V_0 / \Delta$")
axs["mzm"].set_ylabel(r"$\det r_L$")
axs["mzm"].spines["top"].set_visible(False)
axs["mzm"].spines["right"].set_visible(False)

colors = cmap(10 ** np.log10((ls - l_min) / (l_max - l_min)))
for i, _ in enumerate(widths):
    axs["qmzm"].plot(
        data_qMZM["sigmas"][i] / lso, data_qMZM["det_r1_qmzm"][i].real, c=colors[i]
    )

axs["qmzm"].set_xscale("log")
axs["qmzm"].set_xlabel(r"$\sigma / l_{so}$")
axs["qmzm"].spines["top"].set_visible(False)
axs["qmzm"].spines["right"].set_visible(False)

norm = matplotlib.colors.Normalize(l_min, l_max)
cmap = cmr.bubblegum_r
mappable = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)
cb = plt.colorbar(
    mappable=mappable,
    ax=[axs["mzm"], axs["qmzm"]],
    label="$L/l_{so}$",
    pad=0.02,
    location="right",
)

ax_mzm = axs["mzm"].inset_axes([0.01, 0.4, 0.6, 0.6])
ax_mzm.axis("off")

ax_qmzm = axs["qmzm"].inset_axes([0.0, 0.05, 0.6, 0.8])
ax_qmzm.axis("off")

labels = ["(a)", "(b)"]
for ax, label in zip(axs.values(), labels):
    ax.text(0.0, 1.1, label, transform=ax.transAxes)

skunk.connect(ax_mzm, "mzm")
skunk.connect(ax_qmzm, "qmzm")

svg = skunk.insert(
    {
        "mzm": "../src_figures/coupled-mzm.svg",
        "qmzm": "../src_figures/coupled-qmzm.svg",
    },
    randomize_ids=True,
)

with open("../publication/figures/fig2.svg", "w") as f:
    f.write(svg)

skunk.display(svg)
# %%