#!/usr/bin env python3
from formulary import *

# Define the functions
def psi_squared(x, xi):
    return np.tanh(x/(np.sqrt(2)*xi))**2

def B_z(x, B0, lam):
    return B0 * np.exp(-x/lam)


def n_s_boundary():
    xs = np.linspace(0, 6, 400)
    xn = np.linspace(-1, 0, 10)
    B0 = 1.0
    fig, ax = plt.subplots(figsize=size_formula_fill_default)
    ax.axvline(x=0, color='gray', linestyle='--', linewidth=0.8)
    ax.axhline(y=1, color='gray', linestyle='--', linewidth=0.8)
    ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.8)
    ax.fill_between(xn, -2, 2 , color=COLORSCHEME["bg-yellow"], alpha=0.5)
    ax.fill_between(xs, -2, 2 , color=COLORSCHEME["bg-blue"], alpha=0.5)
    ax.text(-0.5, 0.9, 'N', color=COLORSCHEME["fg-yellow"], fontsize=14, ha="center", va="center")
    ax.text(3, 0.9, 'S', color=COLORSCHEME["fg-blue"], fontsize=14, ha="center", va="center")
    ax.set_xlabel("$x$")
    ax.set_ylabel(r"$|\Psi|^2$, $B_z(x)/B_\text{ext}$")
    ax.set_ylim(-0.1, 1.1)
    ax.set_xlim(-1, 6)
    ax.grid()
    lines = []
    for i, (xi, lam, color) in enumerate([(0.5, 2, "blue"), (2, 0.5, "red")]):
        psi = psi_squared(xs, xi)
        B = B_z(xs, B0, lam)
        line, = ax.plot(xs, psi, color=color, linestyle="solid", label=f"$\\xi_\\text{{GL}}={xi}$, $\\lambda_\\text{{GL}}={lam}$")
        lines.append(line)
        ax.plot(xs, B, color=color, linestyle="dashed")
        if i == 1:
            ylam = 1/np.exp(1)
            ax.plot([0, lam], [ylam, ylam], linestyle="dashed", color=COLORSCHEME["fg2"])
            ax.text(lam/2, ylam, r'$\lambda_\text{GL}$', color=color, ha="center", va="bottom")
            yxi = psi_squared(xi, xi)
            ax.plot([0, xi], [yxi, yxi], linestyle="dotted", color=COLORSCHEME["fg2"])
            ax.text(xi/2, yxi, r'$\xi_\text{GL}$', color=color, ha="center", va="bottom")
    lines.append(mpl.lines.Line2D([], [], color="black", label=r"$\lvert\Psi\rvert^2$"))
    lines.append(mpl.lines.Line2D([], [], color="black", linestyle="dashed", label=r"$B_z(x)/B_\text{ext}$"))
    ax.legend(loc='center right', handles=lines)
    return fig


if __name__ == "__main__":
    export(n_s_boundary(), "cm_sc_n_s_boundary")