from plot import *

def flennard_jones(r, epsilon, sigma):
    return 4 * epsilon * ((sigma/r)**12 - (sigma/r)**6)

def lennard_jones():
    fig, ax = plt.subplots(figsize=size_half_half)
    ax.grid()
    ax.set_xlabel(r"$r$")
    ax.set_ylabel(r"$V(r)$")

    r = np.linspace(0.5, 4, 300)
    for epsilon, sigma in [(1, 1), (1, 2)]:
        y = flennard_jones(r, epsilon, sigma)
        label = texvar("epsilon", epsilon) + ", " + texvar("sigma", sigma)
        ax.plot(r, y, label=label)
    ax.legend()
    ax.set_ylim(-1.1, 1.1)
    return fig
export(lennard_jones(), "potential_lennard_jones")

# BOLTZMANN / FERMI-DIRAC / BOSE-EINSTEN DISTRIBUTIONS
def fboltzmann(x):
    return  1 / np.exp(x)
def fbose_einstein(x):
    return  1 / (np.exp(x) -1)
def ffermi_dirac(x):
    return  1 / (np.exp(x) + 1)


def id_qgas():
    fig, ax = plt.subplots(figsize=size_half_half)
    ax.grid()
    ax.set_xlabel(r"$\beta(\epsilon-\mu)$")
    ax.set_ylabel(r"$\langle n(\epsilon)\rangle$")
    x = np.linspace(-4, 4, 300)
    xbe = np.linspace(0.01, 4, 300)
    yb = fboltzmann(x)
    ybe = fbose_einstein(xbe)
    yfd = ffermi_dirac(x)
    ax.vlines([0], ymin=-10, ymax=10, color="black", linestyle="dashed")
    ax.plot(xbe, ybe, label="Bose-Einstein")
    ax.plot(x, yb, label="Boltzmann")
    ax.plot(x, yfd, label="Fermi-Dirac")
    ax.legend()
    ax.set_ylim(-0.1, 4)
    return fig
export(id_qgas(), "td_id_qgas_distributions")

@np.vectorize
def fstep(x):
    return 1 if x >= 0 else 0

def fermi_occupation():
    fig, ax = plt.subplots(figsize=size_half_third)
    # ax.grid()
    # ax.set_xlabel(r"$\beta(\epsilon-\mu)$")
    ax.set_xticks([0])
    ax.set_xticklabels([r"$\epsilon=\mu$"])
    ax.set_ylabel(r"$\langle n(\epsilon)\rangle$")
    x = np.linspace(-4, 4, 300)
    ystep = fstep(-x)
    yfd = ffermi_dirac(x)
    ax.vlines([0], ymin=-10, ymax=10, color="black", linestyle="dashed")
    ax.plot(x, ystep, label="$T = 0$")
    ax.plot(x, yfd, label=r"$T > 0$")
    ax.legend()
    ax.set_ylim(-0.1, 1.1)
    return fig
export(fermi_occupation(), "td_fermi_occupation")

def fermi_heat_capacity():
    fig, ax = plt.subplots(figsize=size_half_third)
    # ax.grid()
    # ax.set_xlabel(r"$\beta(\epsilon-\mu)$")
    x = np.linspace(0, 4, 100)
    T_F = 1
    a = np.pi**2 / (2 * T_F)

    def linear(x):
        return x * a
    X_32 = 3/2 /(np.pi**2 / (2 * T_F))


    low_temp_Cv = linear(x)
    ax.plot(x, low_temp_Cv, color="orange", linestyle="dashed", label=r"${\pi^2}/{2}\,{T}/{T_\text{F}}$")
    ax.hlines([3/2], xmin=0, xmax=10, color="blue", linestyle="dashed", label="Petit-Dulong")
    @np.vectorize
    def unphysical_f(x):
        # exponential 
        # find ṕoint where unshifted exponential has slope of the linear
        # f = 3/2 - exp(-A*x)
        # f' = A exp(-A*x) = a
        # x = -log(a/A) / A
        A = 5
        x_unshifted = -np.log(a/A) / A
        # shift so that this the exponential intersects the linear at this point
        y_intersect = 3/2 - np.exp(-A*x_unshifted)
        # find intersect x value of linear
        # x = y/a
        x_intersect = y_intersect / a
        # shift exp so that x_intersect becomes x_unshifted
        if x > x_intersect: return 3/2 - np.exp(-A * (x-(x_intersect - x_unshifted)))
        else: return a * x
    # ax.plot(x, smoothing, label="smooth")
    y = unphysical_f(x)
    ax.plot(x, y, color="black")
    ax.legend(loc="lower right")


    ax.set_xticks([0, T_F])
    ax.set_xticklabels(["0", r"$T_F$"])
    ax.set_yticks([0, 3/2])
    ax.set_yticklabels(["0", r"$3/2$"])
    ax.set_ylabel(r"${C_V}/{N k_\text{B}}$")
    ax.set_xlim(0, 1.4 * T_F)
    ax.set_ylim(0, 2)
    return fig
export(fermi_heat_capacity(), "td_fermi_heat_capacity")