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

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_half_third)
    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_half_third)
    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
if __name__ == '__main__':
    export(one_atom_basis(), "cm_phonon_dispersion_one_atom_basis")
    export(two_atom_basis(), "cm_phonon_dispersion_two_atom_basis")