import numpy as np
from itertools import izip, count

from chroma.pi0 import pi0_decay
from chroma import event
from chroma.sample import uniform_sphere
from chroma.itertoolset import repeat_func

# generator parts for use with gun()

def from_histogram(h):
    "Yield values drawn randomly from the histogram `h` interpreted as a pdf."
    pdf = h.hist/h.hist.sum()
    cdf = np.cumsum(pdf)

    for x in repeat_func(np.random.random_sample):
        yield h.bincenters[np.searchsorted(cdf, x)]

def constant(obj):
    while True:
        yield obj

def isotropic():
    while True:
        yield uniform_sphere()
        
def line_segment(point1, point2):
    while True:
        frac = np.random.uniform(0.0, 1.0)
        yield frac * point1 + (1.0 - frac) * point2

def fill_shell(center, radius):
    for direction in isotropic():
        r = radius * np.random.uniform(0.0, 1.0)**(1.0/3.0)
        yield r * direction

def flat(e_lo, e_hi):
    while True:
        yield np.random.uniform(e_lo, e_hi)

# vertex generators

def particle_gun(particle_name_iter, pos_iter, dir_iter, ke_iter, start_id=0):
    for i, particle_name, pos, dir, ke in izip(count(start_id), particle_name_iter, pos_iter, dir_iter, ke_iter):
        dir /= np.linalg.norm(dir)
        vertex = event.Vertex(particle_name, pos, dir, None, ke)
        ev_vertex = event.Event(i, vertex, [vertex])
        yield ev_vertex

def pi0_gun(pos_iter, dir_iter, ke_iter, start_id=0):
    for i, pos, dir, ke in izip(count(start_id), pos_iter, dir_iter, ke_iter):
        dir /= np.linalg.norm(dir)
        primary_vertex = event.Vertex('pi0', pos, dir, None, ke)

        cos_theta_rest = np.random.random_sample() * 2 - 1
        theta_rest = np.arccos(cos_theta_rest)
        phi_rest = np.random.random_sample() * 2 * np.pi

        (gamma1_e, gamma1_dir), (gamma2_e, gamma2_dir) = \
            pi0_decay(ke+134.9766, dir, theta_rest, phi_rest)

        gamma1_vertex = event.Vertex('gamma', pos, gamma1_dir, None, gamma1_e)
        gamma2_vertex = event.Vertex('gamma', pos, gamma2_dir, None, gamma2_e)

        ev_vertex = event.Event(i, primary_vertex, [gamma1_vertex, gamma2_vertex])

        yield ev_vertex


def constant_particle_gun(particle_name, pos, dir, ke, start_id=0):
    '''Convenience wrapper around particle gun that assumes all
    arguments are constants, rather than generators.'''
    return particle_gun(constant(particle_name), constant(pos), constant(dir), constant(ke), start_id=start_id)