from library.SM import particles as SMP
from interactions import Interaction
from interactions.four_particle import FourParticleM, FourParticleIntegral
from particles import Particle
from evolution import Universe
from common import CONST, UNITS, Params


params = Params(T=2 * UNITS.MeV,
                dx=1e-1 * UNITS.MeV)

photon = Particle(**SMP.photon)
neutrino = Particle(**SMP.leptons.neutrino_e)

neutrino_scattering = Interaction(
    particles=((neutrino, neutrino), (neutrino, neutrino)),
    antiparticles=((False, True), (False, True)),
    decoupling_temperature=0 * UNITS.MeV,
    Ms=(FourParticleM(K1=64 * CONST.G_F**2, order=(0, 1, 2, 3)),),
    integral=FourParticleIntegral
)

universe = Universe(params=params, plotting=False)
universe.PARALLELIZE = False
universe.add_particles([photon, neutrino])
universe.interactions += [neutrino_scattering]


import numpy
addition = numpy.vectorize(lambda x: 0.1 * numpy.exp(-(x/UNITS.MeV-5)**2),
                           otypes=[numpy.float_])

import matplotlib.pyplot as plt

x = numpy.linspace(neutrino.grid.MIN_MOMENTUM, neutrino.grid.MAX_MOMENTUM*2,
                   num=neutrino.grid.MOMENTUM_SAMPLES*10, endpoint=True)
y = numpy.vectorize(neutrino.distribution)(x)
z = numpy.vectorize(neutrino.equilibrium_distribution_function)(x / universe.params.aT)
w = z + addition(x)
plt.plot(x, y)
plt.plot(x, z)
plt.plot(x, w)
plt.show()

from interactions.ds import D1, D2, D3
from plotting import plot_integrand

plot_integrand(lambda (p1, p2): D1(p1, p2, 1, 1), 'D1', 0)
plot_integrand(lambda (p1, p2): D2(p1, p2, 1, 1), 'D2', 0)
plot_integrand(lambda (p1, p2): D3(p1, p2, 1, 1), 'D3', 0)

raw_input("...")