"""
Bibliotecas
"""
#importa a função special da biblioteca scipy
from scipy import special
#importa funções de geração de gráficos
import matplotlib.pyplot as plt
#importa uma biblioteca de funções matemáticas
import numpy as np
#importa a funcao integral
from scipy.integrate import quad as integrate

"""
Geradores
"""

def gen_z(linspace):
    """
    Gera valores para z, com eficiencia de memoria, afim de calcular o valor da negentropia para esses valores
    """
    z = np.linspace(0,125.7/100,linspace)
    for number in z:
        yield number

"""
Definição de Constantes
"""

def get_par():
    """
    Dá parâmetros do feixe de airy para funções de integrais
    """
    x0 = 100*10**(-6)
    l0 = 500*10**(-9)
    k = 2*np.pi/l0
    q = z/(k*(x0**2))
    a = 0.1
    return x0, q, a

"""
Definição de Limites
"""
def int_limit():
    """
    Limites para as integrais
    """
    return 5*10**(-3)


"""
Integradores
"""

def campo_abs_sqr(x):
    """
    Módulo quadrado do campo
    """
    x0, q, a = get_par()
    return abs(special.airy(x/x0 -(q/2)**2 + (a*q)*1j)[0]*np.exp(a*x/x0 + -a*(q**2)/2 + (q**3)/12*-1j + (a**2)*q/2*1j + (x/x0*q/2)*1j))**2

def xp(x):
    """
    Função x.p(x)
    """
    return (1/kn)*x*campo_abs_sqr(x)

def desvio_padrao(x):
    return ((x-mi)**2/kn)*campo_abs_sqr(x)

def entropia_funcao(x):
    return (1/kn)*campo_abs_sqr(x)*np.log((1/kn)*campo_abs_sqr(x))


"""
Negentropia em Função de Z
"""
linspace = 100

def neg_vs_z():
    """
    Calcula a Negentropia em função de Z
    """
    neg = []
    for value in gen_z(linspace):
        global z
        z = value
        global kn
        kn = integrate(campo_abs_sqr,-int_limit(),int_limit())[0]
        global mi
        mi = integrate(xp, -int_limit(), int_limit())[0]
        dp = integrate(desvio_padrao, -int_limit(), int_limit())[0]
        dp = np.sqrt(dp)
        entG = np.log(dp*np.sqrt(2*np.pi*np.exp(1)))
        entF = -integrate(entropia_funcao,-int_limit(), int_limit())[0]
        neg.append(entG-entF)
    return neg

plt.plot(np.linspace(-int_limit(),int_limit(),linspace),neg_vs_z())
plt.ylabel('Negentropia')
plt.xlabel('z(m)')
plt.show()