#!/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

from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import griddata

def critical_type2():
    Jc0 = 100
    Bc2_0 = 30
    Tc = 90

    T = np.linspace(0, Tc, 100)
    Jc_T = Jc0 * (1 - (T / Tc)**2)
    Bc2_T = Bc2_0 * (1 - (T / Tc)**2)
    B = np.linspace(0, Bc2_0, 100)
    Jc_B = Jc0 * (1 - B / Bc2_0)

    fig = plt.figure(figsize=size_formula_normal_default)
    ax = fig.add_subplot(111, projection='3d')

    ax.plot(T, np.zeros_like(Jc_T), Jc_T, label='$J_c(T)$', color='r')
    ax.plot(T, Bc2_T, np.zeros_like(Bc2_T), label='$B_{c2}(T)$', color='g')
    ax.plot(np.zeros_like(Jc_B), B, Jc_B, label='$J_c(B)$', color='b')

    ax.set_xlim(0, Tc)
    ax.set_ylim(0, Bc2_0)
    ax.set_zlim(0, Jc0)

    # surface
    # T_grid, B_grid = np.meshgrid(T, B)
    # Jc_grid = Jc0 * (1 - (T_grid / Tc)**2) * (1 - B_grid / Bc2_0)
    # surf = ax.plot_surface(T_grid, B_grid, Jc_grid, color='cyan', alpha=0.5)
    ax.set_xlabel('$T$')
    ax.set_ylabel('$B_{c2}$')
    ax.set_zlabel('$J_c$')
    # ax.legend()
    ax.grid(True)

    ax.view_init(elev=30., azim=45)
    ax.set_box_aspect(None, zoom=0.85)
    return fig



def heat_capacity():
    fig, ax = plt.subplots(1, 1, figsize=size_formula_small_quadratic)

    T_max  = 1.7
    Cn_max = 3
    f_Cn = lambda T: T * Cn_max/T_max
    Delta_C = f_Cn(1.0) * 1.43  # BCS prediction
    CsTc = f_Cn(1.0) * (1+1.43)  # BCS prediction
    # exp decay from there
    f_Cs = lambda T:  np.exp(-1 / T + 1) * CsTc

    Tns = np.linspace(0.0, T_max, 100)
    Tss = np.linspace(0.0, 1.0, 100)
    Cns = f_Cn(Tns)
    Css = f_Cs(Tss)
    ax.plot(Tns, Cns, label=r"$c_\text{n}$")
    ax.plot(Tss, Css, label=r"$c_\text{s}$")
    ax.vlines([1.0], ymin=f_Cn(1.0), ymax=(CsTc), color=COLORSCHEME["fg1"], linestyles="dashed")
    ax.text(1.05, CsTc - Delta_C/2, "$\\Delta c$",  color=COLORSCHEME["fg1"])
    ax.set_xlabel(r"$T/T_\text{c}$")
    ax.set_ylabel(r"$c$ [a.u.]")
    ax.legend()
    return fig



if __name__ == "__main__":
    export(n_s_boundary(), "cm_super_n_s_boundary")
    export(critical_type2(), "cm_super_critical_type2")
    export(heat_capacity(), "cm_super_heat_capacity")