import unittest
import numpy as np

from chroma.geometry import Solid, Geometry
from chroma.make import box
from chroma.sim import Simulation
from chroma.optics import vacuum
from chroma.event import Photons

class TestPropagation(unittest.TestCase):
    def testAbort(self):
        '''Photons that hit a triangle at normal incidence should not abort.

        Photons that hit a triangle at exactly normal incidence can sometimes produce a dot product
        that is outside the range allowed by acos().  Trigger these with axis aligned photons in a box
        '''

        # Setup geometry
        cube = Geometry()
        cube.add_solid(Solid(box(100,100,100), vacuum, vacuum))
        cube.pmtids = [0]
        sim = Simulation(cube, vacuum, bvh_bits=4, geant4_processes=0,
                              use_cache=False)

        # Create initial photons
        nphotons = 10000
        positions = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
        directions = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
        polarizations = np.zeros_like(positions)
        phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
        polarizations[:,0] = np.cos(phi)
        polarizations[:,1] = np.sin(phi)
        times = np.zeros(nphotons, dtype=np.float32)
        wavelengths = np.empty(nphotons, np.float32)
        wavelengths.fill(400.0)

        photons = Photons(positions=positions, directions=directions, polarizations=polarizations,
                          times=times, wavelengths=wavelengths)

        # First make one step to check for strangeness
        photon_stop = sim.propagate_photons(photons, max_steps=1)
        self.assertFalse(np.isnan(photon_stop.positions).any())
        self.assertFalse(np.isnan(photon_stop.directions).any())
        self.assertFalse(np.isnan(photon_stop.polarizations).any())
        self.assertFalse(np.isnan(photon_stop.times).any())
        self.assertFalse(np.isnan(photon_stop.wavelengths).any())

        # Now let it run the usual ten steps
        photon_stop = sim.propagate_photons(photons, max_steps=10)
        aborted = (photon_stop.histories & (1 << 31)) > 0
        print 'aborted photons: %1.1f' % (float(np.count_nonzero(aborted)) / nphotons)
        self.assertFalse(aborted.any())