# %%
from pathlib import Path
from string import ascii_lowercase
from functools import partial

import kwant
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from scipy.optimize import root_scalar
import cmasher as cmr
import skunk
from multiterminal_invariant.common import (
    system,
    zero_params,
    plot_potentials,
    save_params,
)
# %%
# Create scattering geometry: finite wire with two normal leads

# Number of sites of the finite wire
width_finite_NSN = 600

# 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 = 0
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(Delta="Delta_left")
right_lead = system(1, (1,))
right_lead = right_lead.substituted(Delta="Delta_right")

syst_NSN.attach_lead(left_lead)  # normal lead
syst_NSN.attach_lead(right_lead)  # normal lead
sysf = syst_NSN.finalized()


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

    return shape


params = zero_params(sysf)
params.update(
    {
        "mu": 0.015,
        "tx": 1,
        "Delta": 0.01,
        "Delta_left": 0,
        "Delta_right": 0,
        "alpha": 0.1,
        "Ez": 0.02,
        "L": width_finite_NSN,
    }
)

_ = plot_potentials(syst_NSN, params, drop=("tx",), add_cells=20)
save_params(params, "fig1")

# %%
mus = np.linspace(-params["Delta"], params["Delta"], 31)
widths = np.logspace(np.log10(20), np.log10(width_finite_NSN), 10)


# %%
def compute_invariant(B, mu, width):
    smatrix = kwant.smatrix(
        sysf,
        energy=0,
        params={
            **params,
            "mu": mu,
            "Ez": B,
            "Delta": Delta_func(params["Delta"], width),
        },
    )
    r_L = smatrix.submatrix(0, 0)
    return np.linalg.det(r_L).real


# %%
# Data for Fig. 1(c)
filename = "../data/fig1_phase_diagram.npy"
if Path(filename).exists():
    sols = np.load(filename)
else:
    sols = []
    for w in widths:
        _sol = []
        for mu in mus:
            sol = root_scalar(
                partial(compute_invariant, mu=mu, width=w),
                x0=np.sqrt(params["Delta"] ** 2 + mu**2),
                bracket=[0, 3 * params["Delta"]],
            )
            _sol.append(sol.root)
        sols.append(_sol)
    sols = np.array(sols)
    np.save(filename, sols)

# %%
# Data for Fig. 1(d)
frac_min_B, frac_max_B = 0.5, 1.5
Bs = np.linspace(params["Delta"] * frac_min_B, params["Delta"] * frac_max_B, 100)
filename = "../data/fig1_detr.npy"
if Path(filename).exists():
    dets = np.load(filename)
else:
    dets = []
    for w in widths:
        for B in Bs:
            dets.append(compute_invariant(B, 0, w))
    dets = np.array(dets).reshape(len(widths), len(Bs))
    np.save(filename, dets)
# %%
# Plot Fig. 1
figwidth = matplotlib.rcParams["figure.figsize"][0]

fig = plt.figure(figsize=(figwidth, figwidth / 2.5), layout="constrained")
axs = fig.subplot_mosaic(
    [
        ["ns", "phase_diagram"],
        ["nsn", "detr"],
    ],
    width_ratios=[2, 1],
)


lso = params["tx"] / params["alpha"]
ls = np.array(widths) / lso
l_min = ls.min()
l_max = ls.max()

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["phase_diagram"], axs["detr"]],
    location="right",
    label="$L/l_{so}$",
    pad=0.02,
)
cb.ax.yaxis.set_ticks([2, 20, 40, 60])

colors = cmap((ls - l_min) / (l_max - l_min))
for i, _ in enumerate(widths):
    axs["phase_diagram"].plot(
        sols[i] / params["Delta"], mus / params["Delta"], c=colors[i]
    )


axs["phase_diagram"].hlines(0, frac_min_B, frac_max_B, ls="--", color="red", alpha=0.5)
axs["phase_diagram"].set_ylabel(r"$\mu / \Delta$")
axs["phase_diagram"].set_xticks([])
axs["phase_diagram"].spines["right"].set_visible(False)
axs["phase_diagram"].spines["top"].set_visible(False)

for i, _ in enumerate(widths):
    axs["detr"].plot(Bs / params["Delta"], dets[i], c=colors[i])
axs["detr"].set_xlabel(r"$E_z/\Delta$")
axs["detr"].set_ylabel(r"$\det r$")
axs["detr"].set_xticks([1 / 2, 1, 3 / 2])
axs["detr"].set_xticklabels(["1/2", "1", "3/2"])
axs["detr"].set_xlim([frac_min_B, frac_max_B])
axs["detr"].axhline(0, ls="--", c="grey", alpha=0.5)
axs["detr"].spines["right"].set_visible(False)
axs["detr"].spines["top"].set_visible(False)

axs["phase_diagram"].sharex(axs["detr"])

axs["ns"].axis("off")
axs["nsn"].axis("off")


text_positions = [[0, 1], [0, 1], [-0.3, 1], [-0.3, 1]]
axs_text = [axs["ns"], axs["nsn"], axs["phase_diagram"], axs["detr"]]
for position, letter, ax in zip(text_positions, ascii_lowercase, axs_text):
    ax.text(
        position[0],
        position[1],
        f"({letter})",
        transform=ax.transAxes,
        color="black",
    )

skunk.connect(axs["ns"], "sk_ns")
skunk.connect(axs["nsn"], "sk_nsn")
svg = skunk.insert(
    {
        "sk_ns": "../src_figures/normal-superconductor-interface.svg",
        "sk_nsn": "../src_figures/normal-superconductor-normal-interface.svg",
    },
    randomize_ids=True,
)
skunk.display(svg)

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

# %%