from numpy import fmax
from plot import *


def get_fig():
    fig, ax = plt.subplots(figsize=size_half_half)
    ax.grid()
    ax.set_xlabel(f"$x$")
    ax.set_ylabel("PDF")
    return fig, ax

# GAUSS / NORMAL
def fgauss(x, mu, sigma_sqr):
    return 1 / (np.sqrt(2 * np.pi * sigma_sqr)) * np.exp(-(x - mu)**2 / (2 * sigma_sqr))

def gauss():
    fig, ax = get_fig()
    x = np.linspace(-5, 5, 300)
    for mu, sigma_sqr in [(0, 1), (0, 0.2), (0, 5), (-2, 2)]:
        y = fgauss(x, mu, sigma_sqr)
        label = texvar("mu", mu) + ", " + texvar("sigma^2", sigma_sqr)
        ax.plot(x, y, label=label)
    ax.legend()
    return fig
export(gauss(), "distribution_gauss")

# CAUCHY / LORENTZ
def fcauchy(x, x_0, gamma):
    return 1 / (np.pi * gamma * (1 + ((x - x_0)/gamma)**2))

def cauchy():
    fig, ax = get_fig()
    x = np.linspace(-5, 5, 300)
    for x_0, gamma in [(0, 1), (0, 0.5), (0, 2), (-2, 1)]:
        y = fcauchy(x, x_0, gamma)
        label = f"$x_0 = {x_0}$ , {texvar('gamma', gamma)}"
        ax.plot(x, y, label=label)
    ax.legend()
    return fig
export(cauchy(), "distribution_cauchy")

# MAXWELL-BOLTZMANN
def fmaxwell(x, a):
    return np.sqrt(2/np.pi) * x**2 / a**3 * np.exp(-x**2 /(2*a**2))

def maxwell():
    fig, ax = get_fig()
    x = np.linspace(0, 20, 300)
    for a in [1, 2, 5]:
        y = fmaxwell(x, a)
        label = f"$a = {a}$"
        ax.plot(x, y, label=label)
    ax.legend()
    return fig

export(maxwell(), "distribution_maxwell-boltzmann")


# POISSON
def fpoisson(k, l):
    return l**k * np.exp(-l) / scp.special.factorial(k)

def poisson():
    fig, ax = get_fig()
    k = np.arange(0, 21, dtype=int)
    for l in [1, 4, 10]:
        y = fpoisson(k, l)
        label = texvar("lambda", l)
        ax.plot(k, y, color="#555")
        ax.scatter(k, y, label=label)
    ax.set_xlabel(f"$k$")
    ax.set_ylabel(f"PMF")
    ax.legend()
    return fig

export(poisson(), "distribution_poisson")

# BINOMIAL
def binom(n, k):
    return scp.special.factorial(n) / (
        scp.special.factorial(k) *
        scp.special.factorial((n-k))
    )
def fbinomial(k, n, p):
    return binom(n, k) * p**k * (1-p)**(n-k)

def binomial():
    fig, ax = get_fig()
    n = 20
    k = np.arange(0, n+1, dtype=int)
    for p in [0.3, 0.5, 0.7]:
        y = fbinomial(k, n, p)
        label = f"$n={n}$, $p={p}$"
        ax.plot(k, y, color="#555")
        ax.scatter(k, y, label=label)
    ax.set_xlabel(f"$k$")
    ax.set_ylabel(f"PMF")
    ax.legend()
    return fig

export(binomial(), "distribution_binomial")

# FERMI-DIRAC

# BOSE-EINSTEIN