import itertools
import numpy as np
from math import sin, cos, sqrt
import copy

import chroma.pi0 as pi0
from chroma.event import Event, Subtrack

# generator parts for use with gun()

def constant(obj):
    while True:
        yield obj

def isotropic():
    while True:
        cos_theta = np.random.uniform(-1.0, 1.0)
        sin_theta = sqrt(1.0 - cos_theta**2)
        phi = np.random.uniform(0.0, 2*np.pi)
        direction = np.array([sin_theta * cos(phi),
                              sin_theta * sin(phi),
                              cos_theta])
        yield direction
        
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, position, direction, total_energy, start_id=0):
    for i, ev_particle, ev_position, ev_direction, ev_total_energy in \
            itertools.izip(itertools.count(start_id), 
                           particle_name, position, direction, total_energy):
        ev = Event(event_id=i, particle_name=ev_particle, 
                   gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_total_energy)
        yield ev

def pi0_gun(pi0_position, pi0_direction, pi0_total_energy, start_id=0):
    for i, ev_position, ev_direction, ev_total_energy in \
            itertools.izip(itertools.count(start_id), pi0_position, pi0_direction, pi0_total_energy):

        ev = Event(event_id=i, particle_name='pi0',
                   gen_position=ev_position, gen_direction=ev_direction, gen_total_energy=ev_total_energy)

        ev_direction = ev_direction / np.linalg.norm(ev_direction)
        
        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.pi0_decay(energy=ev_total_energy,
                                                                       direction=ev_direction,
                                                                       theta=theta_rest,
                                                                       phi=phi_rest)

        ev.subtracks = [Subtrack(particle_name='gamma', position=ev_position, direction=gamma1_dir,
                                 start_time=0.0, total_energy=gamma1_e),
                        Subtrack(particle_name='gamma', position=ev_position, direction=gamma2_dir,
                                 start_time=0.0, total_energy=gamma2_e)]
        yield ev