#!/usr/bin env python3
from formulary import *
from scipy.constants import Boltzmann as kB, hbar

hbar = 1
kB = 1

def fone_atom_basis(q, a, M, C1, C2):
    return np.sqrt(4*C1/M * (np.sin(q*a/2)**2 + C2/C1 * np.sin(q*a)**2))

def one_atom_basis():
    a = 1.
    C1 = 0.25
    C2 = 0
    M = 1.
    qs = np.linspace(-2*np.pi/a, 2*np.pi/a, 300)
    omega = fone_atom_basis(qs, a, M, C1, C2)
    fig, ax = plt.subplots(figsize=size_formula_normal_default)
    ax.set_xlabel(r"$q$")
    ax.set_xticks([i * np.pi/a for i in range(-2, 3)])
    ax.set_xticklabels([f"${i}\\pi/a$" if i != 0 else "0" for i in range(-2, 3)])
    ax.set_ylabel(r"$\omega$ in $\left[4C_1/M\right]$")
    yunit = np.sqrt(4*C1/M)
    ax.set_ylim(0, yunit+0.1)
    ax.set_yticks([0,yunit])
    ax.set_yticklabels(["0","1"])
    ax.plot(qs, omega)
    ax.text(-1.8*np.pi/a, 0.8, "NN\n$C_2=0$", ha='center')
    ax.text(0, 0.8, "1. BZ", ha='center')
    ax.vlines([-np.pi/a, np.pi/a], ymin=-2, ymax=2, color="black")
    ax.grid()
    return fig

def ftwo_atom_basis_acoustic(q, a, M1, M2, C):
    return np.sqrt(C*(1/M1+1/M2) - C * np.sqrt((1/M1+1/M2)**2 - 4/(M1*M2) * np.sin(q*a/2)**2))

def ftwo_atom_basis_optical(q, a, M1, M2, C):
    return np.sqrt(C*(1/M1+1/M2) + C * np.sqrt((1/M1+1/M2)**2 - 4/(M1*M2) * np.sin(q*a/2)**2))

def two_atom_basis():
    a = 1.
    C = 0.25
    M1 = 1.
    M2 = 0.7
    qs = np.linspace(-2*np.pi/a, 2*np.pi/a, 300)
    omega_a = ftwo_atom_basis_acoustic(qs, a, M1, M2, C)
    omega_o = ftwo_atom_basis_optical(qs, a, M1, M2, C)
    fig, ax = plt.subplots(figsize=size_formula_normal_default)
    ax.plot(qs, omega_a, label="acoustic")
    ax.plot(qs, omega_o, label="optical")
    ax.text(0, 0.8, "1. BZ", ha='center')
    ax.vlines([-np.pi/a, np.pi/a], ymin=-2, ymax=2, color="black")
    ax.set_ylim(-0.03, 1.03)
    ax.set_ylabel(r"$\omega$ in $\left[\sqrt{2C\mu^{-1}}\right]$")
    yunit = np.sqrt(2*C*(1/M1+1/M2))
    ax.set_ylim(0, yunit+0.1)
    ax.set_yticks([0,yunit])
    ax.set_yticklabels(["0","1"])
    ax.set_xlabel(r"$q$")
    ax.set_xticks([i * np.pi/a for i in range(-2, 3)])
    ax.set_xticklabels([f"${i}\\pi/a$" if i != 0 else "0" for i in range(-2, 3)])
    ax.legend()
    ax.grid()
    return fig


def fcv_einstein(T, N, omegaE):
    ThetaT = hbar * omegaE / (kB * T)
    return 3 * N * kB * ThetaT**2 * np.exp(ThetaT) / (np.exp(ThetaT) - 1)**2

def fcv_debye_integral(x):
    print(np.exp(x), (np.exp(x) - 1)**2)
    return x**4 * np.exp(x) / ((np.exp(x) - 1)**2)

def heat_capacity_einstein_debye():
    Ts = np.linspace(0, 10, 500)
    omegaD = 1e1
    omegaE = 1
    # N = 10**23
    N = 1
    cvs_einstein = fcv_einstein(Ts, N, omegaE)
    cvs_debye = np.zeros(Ts.shape, dtype=float)
    integral = np.zeros(Ts.shape, dtype=float)
    # cvs_debye = [0.0 for _ in range(Ts.shape[0])] # np.zeros(Ts.shape, dtype=float)
    # integral = [0.0 for _ in range(Ts.shape[0])] # np.zeros(Ts.shape, dtype=float)
    dT = Ts[1] - Ts[0]
    dThetaT = kB*dT/(hbar*omegaD)
    for i, T in enumerate(Ts):
        if i == 0: continue
        ThetaT = kB*T/(hbar*omegaD)
        dIntegral = fcv_debye_integral(ThetaT) * dThetaT
        integral[i] = dIntegral
        # print(integral)
        integral[i] += integral[i-1]
        C_debye = 9 * N * kB * ThetaT**3 * integral[i]
        cvs_debye[i] = C_debye
        print(i, T, ThetaT, dIntegral, C_debye, integral[i])
    fig, ax = plt.subplots(1, 1, figsize=size_formula_normal_default)
    ax.set_xlabel("$T$")
    ax.set_ylabel("$c_V$")
    ax.plot(Ts, cvs_einstein, label="Einstein")
    ax.plot(Ts, cvs_debye, label="Debye")
    ax.plot(Ts, integral, label="integral")
    ax.hlines([3*N*kB], xmin=0, xmax=Ts[-1], colors=COLORSCHEME["fg1"], linestyles="dashed")
    # print(cvs_debye)
    ax.legend()
    return fig




if __name__ == '__main__':
    export(one_atom_basis(), "cm_vib_dispersion_one_atom_basis")
    export(two_atom_basis(), "cm_vib_dispersion_two_atom_basis")
    export(heat_capacity_einstein_debye(), "cm_vib_heat_capacity_einstein_debye")
    print(kB)