import numpy as np
from pycuda import gpuarray as ga
import time
from uncertainties import ufloat
import sys

from chroma.gpu import GPU, to_float3
from chroma.camera import get_rays
from chroma.event import Photons
from chroma.sample import uniform_sphere

def progress(seq):
    "Print progress while iterating over `seq`."
    n = len(seq)
    print '[' + ' '*21 + ']\r[',
    sys.stdout.flush()
    for i, item in enumerate(seq):
        if i % (n//10) == 0:
            print '.',
            sys.stdout.flush()
        yield item
    print ']'
    sys.stdout.flush()

def ray_trace(gpu, number=1000):
    """
    Return the number of mean and standard deviation of the number of ray
    intersections per second as a ufloat for the geometry loaded onto `gpu`.

    .. note::
        The rays are thrown from a camera sitting *outside* of the geometry.

    Args:
        - gpu, chroma.gpu.GPU
            The GPU object with a geometry already loaded.
        - number, int
            The number of kernel calls to average.
    """
    lb, ub = gpu.geometry.mesh.get_bounds()
    scale = np.linalg.norm(ub-lb)
    point = [0,scale,(lb[2]+ub[2])/2]

    size = (800,600)
    width, height = size

    origins, directions = get_rays(point, size, 0.035, focal_length=0.018)

    origins_gpu = ga.to_gpu(to_float3(origins))
    directions_gpu = ga.to_gpu(to_float3(directions))
    pixels_gpu = ga.zeros(width*height, dtype=np.int32)

    run_times = []
    for i in progress(range(number)):
        t0 = time.time()
        gpu.kernels.ray_trace(np.int32(pixels_gpu.size), origins_gpu, directions_gpu, pixels_gpu, block=(gpu.nthreads_per_block,1,1), grid=(pixels_gpu.size//gpu.nthreads_per_block+1,1))
        gpu.context.synchronize()
        elapsed = time.time() - t0
        run_times.append(elapsed)

    return pixels_gpu.size/ufloat((np.mean(run_times),np.std(run_times)))

def propagate(gpu, number=10, nphotons=500000):
    """
    Return the mean and standard deviation of the number of photons propagated
    per second as a ufloat for the geometry loaded onto `gpu`.

    Args:
        - gpu, chroma.gpu.GPU
            The GPU object with a geometry already loaded.
        - number, int
            The number of kernel calls to average.
        - nphotons, int
            The number of photons to propagate per kernel call.
    """
    gpu.setup_propagate()

    run_times = []
    for i in progress(range(number)):
        photons = Photons(np.zeros((nphotons,3)), uniform_sphere(nphotons), np.random.uniform(400,800,size=nphotons))
        gpu.load_photons(photons)
        t0 = time.time()
        gpu.propagate()
        gpu.context.synchronize()
        elapsed = time.time() - t0
        run_times.append(elapsed)

    return nphotons/ufloat((np.mean(run_times),np.std(run_times)))

if __name__ == '__main__':
    from chroma.detectors import build_lbne_200kton, build_minilbne

    lbne = build_lbne_200kton()
    lbne.build(bits=11)

    gpu = GPU()
    gpu.load_geometry(lbne, print_usage=False)

    print '%s track steps/s' % ray_trace(gpu)
    print '%s photons/s' % propagate(gpu)