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

# OPTICS
def eps_r(omega, omega0, gamma, chi, N):
    return 1 + chi + N * elementary_charge**2 /(epsilon_0 * electron_mass) * (1/(omega0**2 - omega**2 - 1j*gamma*omega))

def eps_r_lim0(omega0, gamma, chi, N):
    return 1 + chi + N * elementary_charge**2 /(epsilon_0 * electron_mass * omega0**2)

def eps_r_lim_infty(omega0, gamma, chi, N):
    return 1 + chi

def dielectric_absorption():
    # values taken from adv. sc. physics ex 10/1
    fig, axs = plt.subplots(1, 2, figsize=size_formula_normal_default, sharex=True)
    omega0 = 100e12     # 100 THz
    gamma = 5e12        # 5 THz
    omegas = np.linspace(60e12, 140e12)
    chi = 5
    N = 1e25

    eps_complex = eps_r(omegas, omega0, gamma, chi, N)
    eps_real = eps_complex.copy().real
    eps_imag = eps_complex.copy().imag
    eps_real_0 = eps_r_lim0(omega0, gamma, chi, N)
    eps_real_infty = eps_r_lim_infty(omega0, gamma, chi, N)

    omegas_THz = omegas / 1e12
    anno_color = COLORSCHEME["fg2"]
    axs[0].set_ylabel(r"$\epsReal(\omega)$")
    axs[0].hlines([eps_real_0], omegas_THz[0], omegas_THz[omegas_THz.shape[0]//3], color=anno_color, linestyle="dotted")
    axs[0].hlines([eps_real_infty], omegas_THz[omegas_THz.shape[0]*2//3], omegas_THz[-1], color=anno_color, linestyle="dotted")
    axs[0].text(omegas_THz[-1], eps_real_infty, r"$\epsilon_\txr(\infty)$", ha="right", va="bottom", color=anno_color)
    axs[0].text(omegas_THz[0], eps_real_0, r"$\epsilon_\txr(0)$", ha="left", va="top", color=anno_color)
    axs[1].set_ylabel(r"$\epsImag(\omega)$")
    vals = [eps_real, eps_imag]
    for i in range(2):
        ax = axs[i]
        val = vals[i]
        ax.set_xlabel(r"$\omega\,[\si{\tera\hertz}]$")
        ax.vlines([omega0/1e12], 0, 100, color=anno_color, linestyle="dashed")
        ax.plot(omegas_THz, val)
        vmax = val.max()
        vmin = val.min()
        margin = (vmax - vmin) * 0.05
        ax.set_ylim(vmin - margin, vmax + margin)

    return fig

# Free e-
@np.vectorize
def fn(omega: float, omega_p:float, eps_infty:float, tau:float):
    eps_real = eps_infty * (1-(omega_p**2 * tau**2)/(1+omega**2 * tau**2))
    eps_imag = eps_infty * omega_p**2 * tau/(omega*(1+omega**2 * tau**2))
    eps_complex = (eps_real + 1j*eps_imag)
    n_complex: complex = np.sqrt(eps_complex)
    # n_real = n_complex.real
    # n_imag = n_complex.imag
    # return n_real, n_imag
    return n_complex

def free_electrons_absorption():
    # values taken from adv. sc. physics ex 10/1
    fig, axs = plt.subplots(2, 2, figsize=size_formula_fill_default, sharex=True)
    omega_p = 30e12 # 30 THz
    eps_infty = 11.7
    tau = 1e-6
    omegas = omega_p * np.logspace(-1, 3, 1000, base=10, dtype=float)
    # omegas = omega_p * np.linspace(1e-1, 1e3, 1000)
    # print(omegas)
    n_complex = fn(omegas, omega_p, eps_infty, tau)
    n_real = n_complex.copy().real
    n_imag = n_complex.copy().imag
    R = np.abs((n_complex-1)/(n_complex+1))
    alpha = 2*omegas*n_imag/3e8
    for ax in axs[1,:]:
        ax.set_xlabel(r"$\omega/\omega_\text{p}$")
        ax.set_xscale("log")
        ax.set_xticks([1e-1,1e0,1e1,1e2,1e3])
        ax.set_xticklabels([0.1, 1, 10, 100, 1000])

    omegas_rel = omegas/omega_p
    axs[0,0].plot(omegas_rel, n_real)
    axs[0,0].set_yticks([0,1,2,np.sqrt(eps_infty)])
    axs[0,0].set_yticklabels([0,1,2,r"$\epsilon_\infty$"])
    axs[0,0].set_ylim(-0.2,0.2+np.sqrt(eps_infty))
    axs[0,0].set_ylabel(r"$n^\prime(\omega)$")

    axs[1,0].plot(omegas_rel, n_imag)
    axs[1,0].set_ylabel(r"$n^{\prime\prime}(\omega)$")
    axs[1,0].set_yscale("log")

    axs[0,1].plot(omegas_rel, R)
    axs[0,1].set_ylabel(r"$R(\omega)$")

    axs[1,1].plot(omegas_rel, alpha)
    axs[1,1].set_ylabel(r"$\alpha(\omega)$")
    axs[1,1].set_yscale("log")
    return fig

if __name__ == '__main__':
    export(free_electrons_absorption(), "cm_optics_absorption_free_electrons")
    export(dielectric_absorption(), "cm_optics_absorption_dielectric")