# %%
from functools import partial
from string import ascii_lowercase

import kwant
import numpy as np
import matplotlib.pyplot as plt
import skunk

from multiterminal_invariant.common import (
    system,
    zero_params,
    save_params,
)

# %% [markdown]
# Scattering setup
# %%
# 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 = 50
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 V_func(V, x0, sigma):
    def shape(x):
        return V * np.exp(-((x - x0) ** 2) / (2 * sigma**2))

    return shape


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

    return shape


def eta_func(eta, x_lead):
    def shape(x):
        return eta if np.abs(x) < x_lead else 0

    return shape


# %%
@np.vectorize
def compute_dets(param, value, param_func, params):
    smatrix = kwant.smatrix(
        sysf_NSN,
        energy=0,
        params={
            **params,
            param: param_func(value),
        },
        check_hermiticity=False,
    )
    r_L, r_R = smatrix.submatrix(0, 0), smatrix.submatrix(1, 1)
    return np.linalg.det(r_L), np.linalg.det(r_R), np.linalg.det(smatrix.data)


# %%
# Fig 4(a)
# mu >> Ez >> Delta for quasi-Majoranas
Delta_value = 0.02
params_NSN = {
    **zero_params(sysf_NSN),
    "mu": 0.12,
    "tx": 1,
    "Delta": Delta_value,
    "alpha": 0.2,
    "Ez": 0.1,
    "V_L": 0.1,
    "V_R": 0.1,
    "sigma_L": 10,
    "sigma_R": 10,
    "eta": eta_func(-0.1, x_lead_NSN),
}

save_params(params_NSN, "fig4_symmetric")

params_NSN.update(
    V=lambda x: V_func(params_NSN["V_L"], -x_lead_NSN, params_NSN["sigma_L"])(x)
    + V_func(params_NSN["V_R"], x_lead_NSN, params_NSN["sigma_R"])(x),
    Delta=Delta_func(params_NSN["Delta"], x_lead_NSN),
)
# Sample logarithmic values for eta
eta_values = np.logspace(-12, 1, 40)
det_r1_sym, det_r2_sym, det_s_sym = compute_dets(
    param="eta",
    value=-eta_values,
    param_func=partial(eta_func, x_lead=x_lead_NSN),
    params=params_NSN,
)
# %%
# Fig 4(b)
params_NSN["sigma_R"] = 2 * params_NSN["sigma_R"]
params_NSN["Delta"] = Delta_value

save_params(params_NSN, "fig4_asymmetric")


params_NSN.update(
    V=lambda x: V_func(params_NSN["V_L"], -x_lead_NSN, params_NSN["sigma_L"])(x)
    + V_func(params_NSN["V_R"], x_lead_NSN, params_NSN["sigma_R"])(x),
    Delta=Delta_func(params_NSN["Delta"], x_lead_NSN),
)


det_r1_asym, det_r2_asym, det_s_asym = compute_dets(
    param="eta",
    value=-eta_values,
    param_func=partial(eta_func, x_lead=x_lead_NSN),
    params=params_NSN,
)

# %%
# Make plot
figwidth = plt.rcParams["figure.figsize"][0]

fig, axs = plt.subplot_mosaic(
    [
        ["scheme", "sym", "asym"],
    ],
    figsize=(figwidth, figwidth / 3.5),
    constrained_layout=True,
    width_ratios=[2, 1, 1],
)

(line1,) = axs["sym"].plot(
    eta_values / Delta_value, np.real(det_r1_sym), label=r"$\det r_L$", ls="--", c="C0"
)
(line2,) = axs["sym"].plot(
    eta_values / Delta_value, np.real(det_r2_sym), label=r"$\det r_R$", ls="-.", c="C2"
)
(line3,) = axs["sym"].plot(
    eta_values / Delta_value, np.real(det_s_sym), label=r"$\det S$", ls="-", c="C3"
)

axs["sym"].set_yticks([-1, 0, 1])
axs["sym"].set_xscale("log")
axs["sym"].set_xlabel(r"$\eta / \Delta$")
axs["sym"].spines["right"].set_visible(False)
axs["sym"].spines["top"].set_visible(False)

axs["asym"].plot(
    eta_values / Delta_value, np.real(det_r1_asym), label=r"$\det r_L$", ls="--", c="C0"
)
axs["asym"].plot(
    eta_values / Delta_value, np.real(det_r2_asym), label=r"$\det r_R$", ls="-.", c="C2"
)
axs["asym"].plot(
    eta_values / Delta_value, np.real(det_s_asym), label=r"$\det S$", ls="-", c="C3"
)
axs["asym"].set_yticks([-1, 0, 1])
axs["asym"].set_xscale("log")
axs["asym"].set_xlabel(r"$\eta / \Delta$")
axs["asym"].spines["right"].set_visible(False)
axs["asym"].spines["top"].set_visible(False)

for letter, ax in zip(ascii_lowercase, axs):
    axs[ax].text(
        -0.05,
        1.1,
        f"({letter})",
        transform=axs[ax].transAxes,
        color="black",
    )
fig.legend(
    frameon=False,
    handles=[line1, line2, line3],
    ncols=3,
    loc="outside lower right",
    handlelength=1.7,
    handletextpad=1,
    columnspacing=1,
)

axs["scheme"].axis("off")
skunk.connect(axs["scheme"], "scheme")

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

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

# %%