import matplotlib.pyplot as plt
import numpy as np

# Fermi-Dirac Statistics

kB = 8.617e-5  # Boltzmann constant in eV/K


def fermi_dirac_prob(E, E_F, T):
    """
    The probability that an energy state $E$ is occupied by an electron.
    - $E_F$: The Fermi Level (the energy where the probability of occupancy is exactly 50%).
    - $T$: Absolute temperature in Kelvin.
    - $k_B$: Boltzmann constant ($8.617 \times 10^{-5} \text{ eV/K}$).
    """
    N = (E - E_F) / (kB * T)
    return 1 / (np.exp(N) + 1)


def plot_fermi_dirac(
    id,
    E_F,
    E_c=-4.05,
    E_v=-5.17,
    T=300,
):
    if E_c is None:
        E = np.linspace(E_F - 1, E_F + 1, 300)
    else:
        E = np.linspace(E_v - 0.2, E_c + 0.2, 300)
    prob = fermi_dirac_prob(E, E_F, T)

    plt.axhline(E_F, color="black", linestyle="--", label="Fermi Level $E_F$")
    plt.plot(prob, E, color="blue", label="prob(electron in energy state E)")

    def band(where, c, name):
        plt.fill_betweenx(E, 0, 1, where, color=c, alpha=0.2, label=name)

    if E_c is None:
        band(None, "red", "Continuous Band")
    else:
        band(E >= E_c, "red", "Conduction Band")
        band(E <= E_v, "blue", "Valence Band")
    plt.xlabel("prob(electron in energy state E)")
    plt.ylabel("E [eV]")
    plt.title(id)
    plt.legend()


def plot_Cu_vs_Si():
    plt.figure(figsize=(5 * 2, 4 * 2))
    plt.subplot(2, 2, 1)
    plot_fermi_dirac("Cu", -4.7, None)
    plt.subplot(2, 2, 2)
    plot_fermi_dirac("Si", -4.61)
    plt.subplot(2, 2, 3)
    plot_fermi_dirac("Si-N", -4.15)
    plt.subplot(2, 2, 4)
    plot_fermi_dirac("Si-P", -5.07)
    plt.suptitle("Fermi-Dirac Statistics: Cu vs Si")
    plt.tight_layout()
    plt.savefig("Fermi-Dirac_Cu_vs_Si")
