from numpy import fmax
from formulary import *
import itertools


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

# 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

# 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

# GAMMA
@np.vectorize
def fgamma(x, alpha, lam):
    return lam**alpha / scp.special.gamma(alpha) * x**(alpha-1) * np.exp(-lam*x)
def gamma():
    fig, ax = get_fig()
    x = np.linspace(0, 20, 300)
    for (alpha, lam) in itertools.product([1, 2, 5], [1, 2]):
        y = fgamma(x, alpha, lam)
        label = f"$\\alpha = {alpha}, \\lambda = {lam}$"
        ax.plot(x, y, label=label)
    ax.set_ylim(0, 1.1)
    ax.set_xlim(0, 10)
    ax.legend()
    return fig

# BETA
@np.vectorize
def fbeta(x, alpha, beta):
    return x**(alpha-1) * (1-x)**(beta-1) / scp.special.beta(alpha, beta)
def beta():
    fig, ax = get_fig()
    x = np.linspace(0, 20, 300)
    for (alpha, lam) in itertools.product([1, 2, 5], [1, 2]):
        y = fgamma(x, alpha, lam)
        label = f"$\\alpha = {alpha}, \\beta = {lam}$"
        ax.plot(x, y, label=label)
    ax.set_ylim(0, 1.1)
    ax.set_xlim(0, 10)
    ax.legend()
    return fig


# 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

# 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


if __name__ == '__main__':
    export(gauss(), "distribution_gauss")
    export(cauchy(), "distribution_cauchy")
    export(maxwell(), "distribution_maxwell-boltzmann")
    export(gamma(), "distribution_gamma")
    export(beta(), "distribution_beta")
    export(poisson(), "distribution_poisson")
    export(binomial(), "distribution_binomial")


# FERMI-DIRAC
# BOSE-EINSTEIN
# see stat-mech