#!/usr/local/bin/env python

# ==============================================================================
# MODULE DOCSTRING
# ==============================================================================

"""
ParallelTemperingSampler
========================

Derived multi-thermodynamic state multistate class with exchanging configurations between replicas of different
temperatures. This is a special case which accepts a single thermodynamic_state and different temperatures to sample.
If you want different temperatures and Hamiltonians, use ReplicaExchangeSampler with temperatures pre-set.

COPYRIGHT

Current version by Andrea Rizzi <andrea.rizzi@choderalab.org>, Levi N. Naden <levi.naden@choderalab.org> and
John D. Chodera <john.chodera@choderalab.org> while at Memorial Sloan Kettering Cancer Center.

Original version by John D. Chodera <jchodera@gmail.com> while at the University of
California Berkeley.

LICENSE

This code is licensed under the latest available version of the MIT License.

"""

import copy
import math
import logging
import numpy as np
import openmmtools as mmtools

from .replicaexchange import ReplicaExchangeSampler, ReplicaExchangeAnalyzer
from .multistatereporter import MultiStateReporter

logger = logging.getLogger(__name__)

# ==============================================================================
# PARALLEL TEMPERING
# ==============================================================================

class ParallelTemperingSampler(ReplicaExchangeSampler):
    """Parallel tempering simulation facility.

    This class provides a facility for parallel tempering simulations. It
    is a subclass of :class:`ReplicaExchange`, but provides efficiency improvements
    for parallel tempering simulations, so should be preferred for this type
    of simulation. In particular, this makes use of the fact that the reduced
    potentials are linear in inverse temperature.

    Examples
    --------

    Create the system.

    >>> from simtk import unit
    >>> from openmmtools import testsystems, states, mcmc
    >>> import tempfile
    >>> testsystem = testsystems.AlanineDipeptideImplicit()

    Create thermodynamic states for parallel tempering with exponentially-spaced schedule.

    >>> n_replicas = 3  # Number of temperature replicas.
    >>> T_min = 298.0 * unit.kelvin  # Minimum temperature.
    >>> T_max = 600.0 * unit.kelvin  # Maximum temperature.
    >>> reference_state = states.ThermodynamicState(system=testsystem.system, temperature=T_min)

    Initialize simulation object with options. Run with a GHMC integrator.

    >>> move = mcmc.GHMCMove(timestep=2.0*unit.femtoseconds, n_steps=50)
    >>> simulation = ParallelTemperingSampler(mcmc_moves=move, number_of_iterations=2)

    Create simulation with its storage file (in a temporary directory) and run.

    >>> storage_path = tempfile.NamedTemporaryFile(delete=False).name + '.nc'
    >>> reporter = MultiStateReporter(storage_path, checkpoint_interval=10)
    >>> simulation.create(reference_state,
    ...                   states.SamplerState(testsystem.positions),
    ...                   reporter, min_temperature=T_min,
    ...                   max_temperature=T_max, n_temperatures=n_replicas)
    >>> simulation.run(n_iterations=1)

    Clean up.

    >>> os.remove(storage_path)

    See Also
    --------
    MultiStateSampler
    ReplicaExchangeSampler

    """

    _TITLE_TEMPLATE = ('Parallel tempering simulation created using ParallelTempering '
                       'class of yank.multistate on {}')

    def create(self, thermodynamic_state, sampler_states: list, storage,
               min_temperature=None, max_temperature=None, n_temperatures=None,
               temperatures=None, **kwargs):
        """Initialize a parallel tempering simulation object.

        Parameters
        ----------
        thermodynamic_state : openmmtools.states.ThermodynamicState
            Reference thermodynamic state that will be simulated at the given
            temperatures.

            WARNING: This is a SINGLE state, not a list of states!

        sampler_states : openmmtools.states.SamplerState or list
            One or more sets of initial sampler states. If a list of SamplerStates,
            they will be assigned to replicas in a round-robin fashion.
        storage : str or Reporter
            If str: path to the storage file, checkpoint options are default
            If Reporter: Instanced :class:`Reporter` class, checkpoint information is read from
            In the future this will be able to take a Storage class as well.
        min_temperature : simtk.unit.Quantity, optional
           Minimum temperature (units of temperature, default is None).
        max_temperature : simtk.unit.Quantity, optional
           Maximum temperature (units of temperature, default is None).
        n_temperatures : int, optional
           Number of exponentially-spaced temperatures between ``min_temperature``
           and ``max_temperature`` (default is None).
        temperatures : list of simtk.unit.Quantity, optional
           If specified, this list of temperatures will be used instead of
           ``min_temperature``, ``max_temperature``, and ``n_temperatures`` (units of temperature,
           default is None).
        metadata : dict, optional
           Simulation metadata to be stored in the file.

        Notes
        -----
        Either (``min_temperature``, ``max_temperature``, ``n_temperatures``) must all be
        specified or the list of '`temperatures`' must be specified.

        """
        # Create thermodynamic states from temperatures.
        if not isinstance(thermodynamic_state, mmtools.states.ThermodynamicState):
            raise ValueError("ParallelTempering only accepts a single ThermodynamicState!\n"
                             "If you have already set temperatures in your list of states, please use the "
                             "standard ReplicaExchange class with your list of states.")
        if temperatures is not None:
            logger.debug("Using provided temperatures")
        elif min_temperature is not None and max_temperature is not None and n_temperatures is not None:
            temperatures = [min_temperature + (max_temperature - min_temperature) *
                            (math.exp(i / n_temperatures-1) - 1.0) / (math.e - 1.0)
                            for i in range(n_temperatures)]  # Python 3 uses true division for /
            logger.debug('using temperatures {}'.format(temperatures))
        else:
            raise ValueError("Either 'temperatures' or ('min_temperature', 'max_temperature', "
                             "and 'n_temperatures') must be provided.")

        thermodynamic_states = [copy.deepcopy(thermodynamic_state) for _ in range(n_temperatures)]
        for state, temperature in zip(thermodynamic_states, temperatures):
            state.temperature = temperature

        # Initialize replica-exchange simulation.
        super(ParallelTemperingSampler, self).create(thermodynamic_states, sampler_states, storage=storage, **kwargs)

    def _compute_replica_energies(self, replica_id):
        """Compute the energy for the replica at every temperature.

        Because only the temperatures differ among replicas, we replace the generic O(N^2)
        replica-exchange implementation with an O(N) implementation.

        """
        # Initialize replica energies for each thermodynamic state.
        energy_thermodynamic_states = np.zeros(self.n_states)
        energy_unsampled_states = np.zeros(len(self._unsampled_states))

        # Determine neighborhood
        state_index = self._replica_thermodynamic_states[replica_id]
        neighborhood = self._neighborhood(state_index)
        # Only compute energies over neighborhoods
        energy_neighborhood_states = energy_thermodynamic_states[neighborhood]  # Array, can be indexed like this
        neighborhood_thermodynamic_states = [self._thermodynamic_states[n] for n in neighborhood]  # List

        # Retrieve sampler states associated to this replica.
        sampler_state = self._sampler_states[replica_id]

        # Thermodynamic state differ only by temperatures.
        reference_thermodynamic_state = self._thermodynamic_states[0]

        # Get the context, any Integrator works.
        context, integrator = mmtools.cache.global_context_cache.get_context(reference_thermodynamic_state)

        # Update positions and box vectors.
        sampler_state.apply_to_context(context)

        # Compute energy.
        reference_reduced_potential = reference_thermodynamic_state.reduced_potential(context)

        # Strip reference potential of reference state's beta.
        reference_beta = 1.0 / (mmtools.constants.kB * reference_thermodynamic_state.temperature)
        reference_reduced_potential /= reference_beta

        # Update potential energy by temperature.
        for thermodynamic_state_id, thermodynamic_state in enumerate(neighborhood_thermodynamic_states):
            beta = 1.0 / (mmtools.constants.kB * thermodynamic_state.temperature)
            energy_neighborhood_states[thermodynamic_state_id] = beta * reference_reduced_potential

        # Since no assumptions can be made about the unsampled thermodynamic states, do it the hard way
        for unsampled_id, state in enumerate(self._unsampled_states):
            if unsampled_id == 0 or not state.is_state_compatible(context_state):
                context_state = state

                # Get the context, any Integrator works.
                context, integrator = mmtools.cache.global_context_cache.get_context(state)

                # Update positions and box vectors. We don't need
                # to set Context velocities for the potential.
                sampler_state.apply_to_context(context, ignore_velocities=True)
            else:
                # If this state is compatible with the context, just fix the
                # thermodynamic state as positions/box vectors are the same.
                state.apply_to_context(context)

            # Compute energy.
            energy_unsampled_states[unsampled_id] = state.reduced_potential(context)

        # Return the new energies.
        return energy_neighborhood_states, energy_unsampled_states

class ParallelTemperingAnalyzer(ReplicaExchangeAnalyzer):
    """
    The ParallelTemperingAnalyzer is the analyzer for a simulation generated from a Parallel Tempering sampler
    simulation, implemented as an instance of the :class:`ReplicaExchangeAnalyzer` as the sampler is a subclass of
    the :class:`yank.multistate.ReplicaExchangeSampler`

    See Also
    --------
    PhaseAnalyzer
    ReplicaExchangeAnalyzer

    """
    pass

# ==============================================================================
# MAIN AND TESTS
# ==============================================================================

if __name__ == "__main__":
    import doctest
    doctest.testmod()


# -*- coding: utf-8 -*-

import cv2
import numpy as np
import os
import time
import gestureCNN as myNN

minValue = 70

x0 = 400
y0 = 200
height = 200
width = 200

img_flag = "CAM"
cap = cv2.VideoCapture(0)

saveImg = False
guessGesture = True
visualize = False

lastgesture = -1

kernel = np.ones((15, 15), np.uint8)
kernel2 = np.ones((1, 1), np.uint8)
skinkernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))

# Which mask mode to use BinaryMask or SkinMask (True|False)
binaryMode = True
counter = 0

# This parameter controls number of image samples to be taken PER gesture
numOfSamples = 301
gestname = ""
path = ""
mod = 0

banner =  '''\nWhat would you like to do ?
    1- Use pretrained model for gesture recognition & layer visualization
    2- Train the model (you will require image samples for training under .\imgfolder)
    3- Visualize feature maps of different layers of trained model
    '''

def saveROIImg(img):
    global counter, gestname, path, saveImg
    if counter > (numOfSamples - 1):
        # Reset the parameters
        saveImg = False
        gestname = ''
        counter = 0
        return

    counter = counter + 1
    name = gestname + str(counter)
    print("Saving img:", name)
    cv2.imwrite(path+name + ".png", img)
    time.sleep(0.04 )

def skinMask(frame, x0, y0, width, height ):
    global guessGesture, visualize, mod, lastgesture, saveImg
    # HSV values
    low_range = np.array([0, 50, 80])
    upper_range = np.array([30, 200, 255])

    cv2.rectangle(frame, (x0,y0),(x0+width,y0+height),(0,255,0),1)
    roi = frame[y0:y0+height, x0:x0+width]

    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)

    #Apply skin color range
    mask = cv2.inRange(hsv, low_range, upper_range)

    mask = cv2.erode(mask, skinkernel, iterations = 1)
    mask = cv2.dilate(mask, skinkernel, iterations = 1)

    #blur
    mask = cv2.GaussianBlur(mask, (15,15), 1)

    #bitwise and mask original frame
    res = cv2.bitwise_and(roi, roi, mask = mask)

    # color to grayscale
    res = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY)

    guess_res = "^______^Skin"

    if saveImg == True:
        saveROIImg(res)

    elif guessGesture == True:
        retgesture, guess_res  = myNN.guessGesture(mod, res)

        if lastgesture != retgesture:
            lastgesture = retgesture
            print myNN.output[lastgesture]
            time.sleep(0.01 )
            #guessGesture = False

    elif visualize == True:
        layer = int(raw_input("Enter which layer to visualize "))
        cv2.waitKey(0)
        myNN.visualizeLayers(mod, res, layer)
        visualize = False

    return res, guess_res

def binaryMask(frame, x0, y0, width, height):

    global guessGesture, visualize, mod, lastgesture, saveImg

    cv2.rectangle(frame, (x0,y0),(x0+width,y0+height),(0,255,0),1)
    roi = frame[y0:y0+height, x0:x0+width]

    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray,(5,5),2)

    th3 = cv2.adaptiveThreshold(blur,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV,11,2)
    ret, res = cv2.threshold(th3, minValue, 255, cv2.THRESH_BINARY_INV+cv2.THRESH_OTSU)

    guess_res = "Bin"

    if saveImg == True:
        saveROIImg(res)

    elif guessGesture == True:
        retgesture, guess_res = myNN.guessGesture(mod, res)

        if lastgesture != retgesture:
            lastgesture = retgesture

    elif visualize == True:
        layer = int(raw_input("Enter which layer to visualize "))
        cv2.waitKey(1)
        myNN.visualizeLayers(mod, res, layer)
        visualize = False

    return res, guess_res

def Main():
    global guessGesture, visualize, mod, binaryMode, x0, y0, width, height, saveImg, gestname, path, img_flag, cap

    font = cv2.FONT_HERSHEY_SIMPLEX
    size = 0.7
    fx = 10
    fy = 355
    fh = 18

    #Call CNN model loading callback
    while True:
        # ans = int(raw_input(banner))
        ans = 1
        if ans == 2:
            mod = myNN.loadCNN(-1)
            myNN.trainModel(mod)
            raw_input("Press any key to continue")
            break
        elif ans == 1:
            print "Will load default weight file"
            mod = myNN.loadCNN(0)
            break
        elif ans == 3:
            if not mod:
                w = int(raw_input("Which weight file to load (0 or 1)"))
                mod = myNN.loadCNN(w)
            else:
                print "Will load default weight file"

            img = int(raw_input("Image number "))
            layer = int(raw_input("Enter which layer to visualize "))
            myNN.visualizeLayers(mod, img, layer)
            raw_input("Press any key to continue")
            continue

        else:
            print "Get out of here!!!"
            return 0

    import os
    print os.getpid()

    ## Grab camera input
    cv2.namedWindow('Original', cv2.WND_PROP_FULLSCREEN)  # Set Camera Size, cv2.WINDOW_NORMAL
    cv2.setWindowProperty('Original', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

    import signal

    # MAC : SIGPROF 27
    def signalCameraOnHandler(signum, frame):
        print signum
        global img_flag
        global cap

        cap = cv2.VideoCapture(0)
        cv2.namedWindow('Original', cv2.WINDOW_NORMAL)
        img_flag = "CAM"

    # MAC : SIGINFO 29
    def signalMapHandler(signum, frame):
        print signum
        global img_flag
        global cap

        cap = cv2.VideoCapture('map.gif')
        cv2.namedWindow('Original', cv2.WINDOW_NORMAL)
        img_flag = "MAP"

    # MAC : SIGUSR1 30
    def signalSmileFaceHandler(signum, frame):
        print signum
        global img_flag
        global cap

        cap = cv2.VideoCapture('smile_glow.gif')
        cv2.namedWindow('Original', cv2.WINDOW_NORMAL)
        img_flag = "SMILE"

    # MAC : SIGUSR2 31
    def signalDefaultFaceHandler(signum, frame):
        print signum
        global img_flag
        global cap

        cap = cv2.VideoCapture('normal_glow.gif')
        cv2.namedWindow('Original', cv2.WINDOW_NORMAL)
        img_flag = "NORMAL"

    # Set Signal
    signal.signal(signal.SIGPROF, signalCameraOnHandler)
    signal.signal(signal.SIGINFO, signalMapHandler)
    signal.signal(signal.SIGUSR1, signalSmileFaceHandler)
    signal.signal(signal.SIGUSR2, signalDefaultFaceHandler)

    reset_flag = False

    thumbs_up_guess_stack = 0
    smile_face_stack = 0
    camera_page_stack = 0
    map_page_stack = 0

    while True:
        ret, frame = cap.read()

        # for end of gif
        if not ret:
            if img_flag == "NORMAL":
                cap = cv2.VideoCapture('normal_glow.gif')
                cv2.namedWindow('Original', cv2.WINDOW_NORMAL)

            elif img_flag == "SMILE":
                cap = cv2.VideoCapture('smile_glow.gif')
                cv2.namedWindow('Original', cv2.WINDOW_NORMAL)

            elif img_flag == "MAP":
                cap = cv2.VideoCapture('map.gif')
                cv2.namedWindow('Original', cv2.WINDOW_NORMAL)

            ret, frame = cap.read()

        if img_flag == "CAM":
            camera_page_stack += 1

            frame = cv2.flip(frame, 3)

            guess_res = ":)"

            if ret:
                if binaryMode: # on
                    roi, guess_res = binaryMask(frame, x0, y0, width, height)
                else:
                    roi, guess_res = skinMask(frame, x0, y0, width, height)

            # detact thumbs-up frame
            if guess_res.split()[0] == "PEACE" or guess_res.split()[0] == "OK":
                camera_page_stack -= 1
                thumbs_up_guess_stack += 1
            else:
                thumbs_up_guess_stack = 0

            # print GUESS
            cv2.putText(frame, str(thumbs_up_guess_stack) + " " + guess_res, (fx, fy), font, 1, (0, 0, 255), 2, 1)
            cv2.putText(frame, "Show Thumbs-Up on TONY's hat camera!", (fx, fy + fh), font, size, (0, 0, 255), 1, 1)

        elif img_flag == "SMILE":
            smile_face_stack += 1

            cap = cv2.VideoCapture('smile_glow.gif')

        elif img_flag == "NORMAL":
            cap = cv2.VideoCapture('normal_glow.gif')

        elif img_flag == "MAP":
            map_page_stack += 1
            cap = cv2.VideoCapture('map.gif')

        # for signal call
        if thumbs_up_guess_stack >= 6:
            # paging to smile
            os.system("kill -30 "+str(os.getpid()))
            reset_flag = True

        if smile_face_stack >= 30:
            # paging to map
            os.system("kill -29 "+str(os.getpid()))
            reset_flag = True

        if camera_page_stack >= 150:
            # paging to map
            os.system("kill -29 "+str(os.getpid()))
            reset_flag = True

        # now threshold == 300
        if map_page_stack >= 200:
            # paging to camera
            os.system("kill -27 " + str(os.getpid()))
            reset_flag = True

        if reset_flag:
            thumbs_up_guess_stack = 0
            smile_face_stack = 0
            camera_page_stack = 0
            map_page_stack = 0
            reset_flag = False

        cv2.imshow('Original', frame)

        key = cv2.waitKey(10) & 0xff
        if key == 27:
            break

    #Realse & destroy
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    Main()


# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Define geometry for ACIS radiator and sunshade and perform raytrace
calculation of Earth illumination on the radiator surface.
"""
import hashlib
import functools
from contextlib import ExitStack

from Quaternion import Quat
import numpy as np
from numpy.random import RandomState

prng = RandomState()

_RANDOM_SALT = 1

def set_random_salt(salt):
    """
    Set the internal ``_RANDOM_SALT`` module attribute to ``salt``.

    If set to ``None`` then the system random seed will be used everywhere and
    the results of ``calc_earth_vis`` will not be deterministic (i.e. they will
    vary slightly from one run to the next).

    If set otherwise then the ``calc_earth_vis`` function will return precisely
    the same results for the same input arguments and same value of ``salt``.
    Different ``salt`` values will generate different random variations of the
    results.  This is done by making an MD5 digest of random salt and the
    function arguments and using this to generate an integer which sets the
    random seed.

    Any object type is allowed for ``salt``, the only thing that matters is
    is that it has a unique string representation.

    :param salt: any object with a unique str representation, repr(salt)
    """
    global _RANDOM_SALT
    _RANDOM_SALT = salt

def _encode_data(val):
    with ExitStack() as stack:
        # Numpy 1.14 improved printing of arrays but this changes the encoded
        # array data here. For compatibility with the 1.12 formatting that forms
        # the basis of regression testing, use the old style here. np.printoptions
        # was added in numpy 1.15, so if that is found then we must use the legacy
        # print method. Note that legacy="1.12" does not seem to work.
        if hasattr(np, 'printoptions'):
            stack.enter_context(np.printoptions(legacy="1.13"))
        return np.array_str(np.array(val, ndmin=1)).encode("utf8")

def _make_reproducible(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        if _RANDOM_SALT is not None:
            md5 = hashlib.md5()
            md5.update(repr(_RANDOM_SALT).encode("utf8"))
            for val in args:
                md5.update(_encode_data(val))
            for val in kwargs.values():
                md5.update(_encode_data(val))
            digest = md5.hexdigest()
            seed = int(digest[:7], 16)
            prng.seed(seed)

        out = func(*args, **kwargs)
        return out
    return wrapper

def make_taco():
    """Define geometry for ACIS radiator sun-shade (aka the space taco) in mm."""
    y_pnts = np.array([-689.0, -333, -63, 553, 689]) + TACO_Y_OFF
    z_pnts = np.array([0.0, 777, 777, 421, 0]) - RAD_Z_OFF

    y_edge = np.arange(max(y_pnts))
    z_edge = interpolate(z_pnts, y_pnts, y_edge)

    return z_edge

@_make_reproducible
def calc_earth_vis(p_chandra_eci=None,
                   chandra_att=None,
                   max_reflect=10,
                   p_earth_body=None):
    """Calculate the relative Earth visibility for the ACIS radiator given
    the Chandra orbit position ``p_chandra_eci`` and attitude ``chandra_att``.
    Alternately one can directly supply ``p_earth_body`` as a numpy array
    representing the position of the Earth in Chandra body coordinates.

    The relative visibility is normalized so that 1.0 represents the entire
    radiator having visibility toward the surface point and being exactly
    normal toward the surface.

    Total illumination gives the effective solid angle (steradians) of the
    visible Earth from the radiator perspective.  It is averaged over the
    different points on the radiator.

    :param p_chandra_eci: Chandra orbital position [x, y, z] (meters)
    :param chandra_att: Chandra attitude [ra, dec, roll] (deg)
    :param max_reflect: number of reflections to compute (default=10)
    :param p_earth_body: optional 3-vector [x, y, z] (meters) of Earth
        in Chandra body coordinates (instead of p_chandra_eci and
        chandra_att inputs)

    :returns: relative visibility, total illumination, projected rays
    """

    if p_earth_body is None:
        # Calculate position of earth in ECI and Chandra body coords.
        q_att = Quat(chandra_att)

        # For T = attitude transformation matrix then p_body = T^-1 p_eci
        p_earth_body = np.dot(q_att.transform.transpose(), -np.array(p_chandra_eci))
    else:
        p_earth_body = np.asarray(p_earth_body)  # no-op if already an array

    illum = np.zeros(max_reflect+1)
    out_rays = []

    # If Earth disk is entirely below the "horizon" of the ACIS radiator then illum=0
    if p_earth_body[2] < -RAD_EARTH:
        return [], illum, out_rays

    # Quaternion to transform x-axis to the body-earth vector
    q_earth = quat_x_v2(p_earth_body)
    open_angle = np.arcsin(RAD_EARTH / np.sqrt(np.sum(p_earth_body**2)))

    rays, earth_solid_angle = sphere_rand(open_angle)
    n_rays = len(rays)
    rays_to_earth = np.dot(q_earth.transform, rays.transpose()).transpose()  # shape (n_rays, 3)

    # Accept only rays with a positive Z component and make sure no X component is < 1e-6
    rays_to_earth = rays_to_earth[rays_to_earth[:, 2] > 0.0, :]
    n_rays_to_earth = len(rays_to_earth)
    if n_rays_to_earth % 2 == 1:  # Ugh, make sure this list has even length
        rays_to_earth = rays_to_earth[:-1, :]
        n_rays_to_earth -= 1

    vis = np.zeros((max_reflect+1, n_rays))
    if len(rays_to_earth) == 0:
        return vis, illum, out_rays

    rays_to_earth_x = np.abs(rays_to_earth[:, 0])
    rays_to_earth_x[rays_to_earth_x < 1e-6] = 1e-6

    # Main ray-trace loop.  Calculate ray visibility for an increasing number
    # of reflections.  Rays that get blocked after N reflections are candidates
    # for getting out after N+1 reflections.

    # Radiator size: 17 inches wide (X) by 19 inches long (Y), centered within
    # SIM shaded structure.

    rad_x = prng.uniform(low=0.0, high=215.9, size=n_rays_to_earth // 2)
    rad_y = prng.uniform(low=-241.3, high=241.3, size=n_rays_to_earth // 2) + TACO_Y_OFF
    rad_x = np.append(rad_x, -rad_x)
    rad_y = np.append(rad_y, rad_y)
    rad_z = np.zeros(n_rays_to_earth)
    rays_x = rays_to_earth_x
    rays_y = rays_to_earth[:, 1]
    rays_z = rays_to_earth[:, 2]

    for refl in range(max_reflect + 1):
        if len(rays_x) == 0:
            break
        taco_x = TACO_X_OFF * (refl + 1)
        dx = (taco_x - rad_x) / rays_x

        # Find rays that intersect shade within Y limits of the shade (0, N_TACO mm)
        ray_taco_iys = np.array(rad_y + rays_y * dx, dtype=int)
        y_ok = (ray_taco_iys >= 0) & (ray_taco_iys < N_TACO)
        y_not_ok = ~y_ok
        ray_taco_iys[y_not_ok] = 0

        # From those rays find ones below the TACO_Z_EDGES curve
        ray_taco_zs = rays_z * dx
        z_ok = ray_taco_zs > TACO_Z_EDGES[ray_taco_iys]
        y_z_ok = (y_ok & z_ok) | y_not_ok

        # Calculate the delta-visibility for good rays (cos(incidence_angle) = Z component)
        d_vis = rays_z[y_z_ok] * REFLECT_ATTEN**refl
        # vis[refl, rays_i[y_z_ok]] += d_vis
        illum[refl] += earth_solid_angle * np.sum(d_vis) / n_rays

        blocked = ~y_z_ok
        rays_x = rays_x[blocked]
        rays_y = rays_y[blocked]
        rays_z = rays_z[blocked]
        # rays_i = rays_i[blocked]
        rad_x = rad_x[blocked]
        rad_y = rad_y[blocked]

    return vis, illum, out_rays

def interpolate(yin, xin, xout):
    """
    Interpolate the curve defined by (xin, yin) at points xout.  The array
    xin must be monotonically increasing.  The output has the same data type as
    the input yin.

    :param yin: y values of input curve
    :param xin: x values of input curve
    :param xout: x values of output interpolated curve

    @:rtype: numpy array with interpolated curve
    """
    lenxin = len(xin)

    i1 = np.searchsorted(xin, xout)
    i1[ i1==0 ] = 1
    i1[ i1==lenxin ] = lenxin-1

    x0 = xin[i1-1]
    x1 = xin[i1]
    y0 = yin[i1-1]
    y1 = yin[i1]

    return (xout - x0) / (x1 - x0) * (y1 - y0) + y0

def norm(vec):
    return vec / np.sqrt(np.sum(vec**2))

def quat_x_v2(v2):
    """Generate quaternion that rotates X-axis into v2 by the shortest path"""
    x = np.array([1.,0,0])
    v2 = norm(np.array(v2))
    dot = np.dot(x, v2)
    if abs(dot) > 1-1e-8:
        x = norm(np.array([1., 0., 1e-5]))
        dot = np.dot(v2, x)
    angle = np.arccos(dot)
    axis = norm(np.cross(x, v2))
    sin_a = np.sin(angle / 2)
    cos_a = np.cos(angle / 2)
    return Quat([axis[0] * sin_a,
                            axis[1] * sin_a,
                            axis[2] * sin_a,
                            cos_a])

def sphere_grid(ngrid, open_angle):
    """Calculate approximately uniform spherical grid of rays containing
    ``ngrid`` points and extending over the opening angle ``open_angle``
    (radians).

    :returns: numpy array of unit length rays, grid area (steradians)
    """
    from math import sin, cos, radians, pi, sqrt

    grid_area = 2*pi*(1-cos(open_angle))
    if ngrid <= 1:
        return np.array([[1., 0., 0.]]), grid_area

    gridsize = sqrt(grid_area / ngrid)

    grid = []
    theta0 = pi/2-open_angle
    n_d = int(round(open_angle / gridsize))
    d_d = open_angle / n_d

    for i_d in range(0, n_d+1):
        dec = i_d * d_d + pi/2 - open_angle
        if abs(i_d) != n_d:
            n_r = int(round( 2*pi * cos(dec) / d_d))
            d_r = 2*pi / n_r
        else:
            n_r = 1
            d_r = 1
        for i_r in range(0, n_r):
            ra = i_r * d_r
            # This has x <=> z (switched) from normal formulation to make the
            # grid centered about the x-axis
            grid.append((sin(dec), sin(ra) * cos(dec), cos(ra) * cos(dec)))

    # (x, y, z) = zip(*grid)  (python magic)
    return np.array(grid), grid_area

@_make_reproducible
def sphere_rand(open_angle, min_ngrid=100, max_ngrid=10000):
    """Calculate approximately uniform spherical grid of rays containing
    ``ngrid`` points and extending over the opening angle ``open_angle``
    (radians).

    :returns: numpy array of unit length rays, grid area (steradians)
    """
    from math import sin, cos, radians, pi, sqrt

    xmin = cos(open_angle)
    grid_area = 2*pi*(1-xmin)

    idx_xmin = np.searchsorted(SPHERE_X, xmin)
    if N_SPHERE - idx_xmin < min_ngrid:
        idx_xmin = N_SPHERE - min_ngrid

    ngrid = int(min_ngrid + (max_ngrid-min_ngrid) * (1-xmin))
    idx_sphere = prng.randint(idx_xmin, N_SPHERE, ngrid)

    return SPHERE_XYZ[idx_sphere, :], grid_area

@_make_reproducible
def random_hemisphere(nsample):
    x = prng.uniform(low=0.3, high=1.0, size=nsample)
    x.sort()                    # x is not random
    t = 2*np.pi * prng.uniform(size=nsample)
    r = np.sqrt(1-x**2)
    z = r * np.cos(t)
    y = r * np.sin(t)
    return np.array([x, y, z]).transpose()

RAD_EARTH = 6371e3
RAD_Z_OFF = 36
TACO_X_OFF = 230
TACO_Y_OFF = 689

REFLECT_ATTEN = 0.9
TACO_Z_EDGES = make_taco()
N_TACO = len(TACO_Z_EDGES)

N_SPHERE = 1500000
SPHERE_XYZ = random_hemisphere(N_SPHERE)
SPHERE_X = SPHERE_XYZ[:, 0].copy()


import numpy as np
import chroma.api as gpuapi
from chroma.gpu.tools import get_module, api_options, \
    chunk_iterator, to_uint3, to_float3, copy_to_uint3, copy_to_float3
from chroma.gpu.gpufuncs import GPUFuncs
if gpuapi.is_gpu_api_cuda():
    import pycuda.driver as cuda
    from pycuda import gpuarray as ga
    from pycuda import characterize
    #from chroma.gpu.cutools import Mapped
    import chroma.gpu.cutools as cutools
elif gpuapi.is_gpu_api_opencl():
    import pyopencl as cl
    import pyopencl.array as ga
    import chroma.gpu.cltools as cltools

from chroma.bvh.bvh import WorldCoords, node_areas, NCHILD_MASK, CHILD_BITS

def round_up_to_multiple(x, multiple):
    remainder = x % multiple
    if remainder == 0:
        return x
    else:
        return x + multiple - remainder

def create_leaf_nodes(mesh, morton_bits=16, round_to_multiple=1, nthreads_per_block=32, max_blocks=16):
    '''Compute the leaf nodes surrounding a triangle mesh.

      ``mesh``: chroma.geometry.Mesh
        Triangles to box
      ``morton_bits``: int
        Number of bits to use per dimension when computing Morton code.
      ``round_to_multiple``: int
        Round the number of nodes created up to multiple of this number
        Extra nodes will be all zero.

    Returns (world_coords, nodes, morton_codes), where
      ``world_coords``: chroma.bvh.WorldCoords
        Defines the fixed point coordinate system
      ``nodes``: ndarray(shape=len(mesh.triangles), dtype=uint4)
        List of leaf nodes.  Child IDs will be set to triangle offsets.
      ``morton_codes``: ndarray(shape=len(mesh.triangles), dtype=np.uint64)
        Morton codes for each triangle, using ``morton_bits`` per axis.
        Must be <= 16 bits.
    '''
    # it would be nice not to duplicate code, make functions transparent...
    context = None
    queue = None
    if gpuapi.is_gpu_api_opencl():
        context = cltools.get_last_context()
        #print context
        queue = cl.CommandQueue( context )

    # Load GPU functions
    if gpuapi.is_gpu_api_cuda():
        bvh_module = get_module('bvh.cu', options=api_options, include_source_directory=True)
    elif gpuapi.is_gpu_api_opencl():
        # don't like the last context method. trouble. trouble.
        bvh_module = get_module('bvh.cl', cltools.get_last_context(), options=api_options, include_source_directory=True)
    bvh_funcs = GPUFuncs(bvh_module)

    # compute world coordinates
    world_origin_np = mesh.vertices.min(axis=0)
    world_scale = np.max((mesh.vertices.max(axis=0) - world_origin_np)) / (2**16 - 2)
    world_coords = WorldCoords(world_origin=world_origin_np, world_scale=world_scale)

    # Put triangles and vertices into host and device memory
    # unfortunately, opencl and cuda has different methods for managing memory here
    # we have to write divergent code
    if gpuapi.is_gpu_api_cuda():
        # here cuda supports a nice feature where we allocate host and device memory that are mapped onto one another.
        # no explicit requests for transfers here
        triangles = cutools.mapped_empty(shape=len(mesh.triangles), dtype=ga.vec.uint3, write_combined=True)
        triangles[:] = to_uint3(mesh.triangles)
        vertices = cutools.mapped_empty(shape=len(mesh.vertices), dtype=ga.vec.float3, write_combined=True)
        vertices[:] = to_float3(mesh.vertices)
        #print triangles[0:10]
        #print vertices[0:10]

        # Call GPU to compute nodes
        nodes = ga.zeros(shape=round_up_to_multiple(len(triangles), round_to_multiple), dtype=ga.vec.uint4)
        morton_codes = ga.empty(shape=len(triangles), dtype=np.uint64)

        # Convert world coords to GPU-friendly types
        world_origin = ga.vec.make_float3(*world_origin_np)
        world_scale = np.float32(world_scale)

        # generate morton codes on GPU
        for first_index, elements_this_iter, nblocks_this_iter in \
                chunk_iterator(len(triangles), nthreads_per_block,
                               max_blocks=30000):
            bvh_funcs.make_leaves(np.uint32(first_index),
                                  np.uint32(elements_this_iter),
                                  cutools.Mapped(triangles), cutools.Mapped(vertices),
                                  world_origin, world_scale,
                                  nodes, morton_codes,
                                  block=(nthreads_per_block,1,1),
                                  grid=(nblocks_this_iter,1))

        morton_codes_host = morton_codes.get() >> (16 - morton_bits)

    elif gpuapi.is_gpu_api_opencl():
        # here we need to allocate a buffer on the host and on the device
        triangles = np.empty( len(mesh.triangles), dtype=ga.vec.uint3 )
        copy_to_uint3(mesh.triangles, triangles)
        vertices = np.empty( len(mesh.vertices), dtype=ga.vec.float3 )
        copy_to_float3(mesh.vertices, vertices)
        # now create a buffer object on the device and push data to it
        triangles_dev = ga.to_device( queue, triangles )
        vertices_dev = ga.to_device( queue, vertices )

        # Call GPU to compute nodes
        nodes = ga.zeros(queue, shape=round_up_to_multiple(len(triangles), round_to_multiple), dtype=ga.vec.uint4)
        morton_codes = ga.empty(queue, shape=len(triangles), dtype=np.uint64)

        # Convert world coords to GPU-friendly types
        #world_origin = np.array(world_origin_np,dtype=np.float32)
        world_origin = np.empty(1,dtype=ga.vec.float3)
        world_origin['x'] = world_origin_np[0]
        world_origin['y'] = world_origin_np[1]
        world_origin['z'] = world_origin_np[2]
        world_scale = np.float32(world_scale)
        #print world_origin, world_scale

        # generate morton codes on GPU
        for first_index, elements_this_iter, nblocks_this_iter in \
                chunk_iterator(len(triangles), nthreads_per_block, max_blocks):
            print first_index, elements_this_iter, nblocks_this_iter
            bvh_funcs.make_leaves( queue, (nblocks_this_iter,1,1), (nthreads_per_block,1,1),
            #bvh_funcs.make_leaves( queue, (elements_this_iter,1,1), None,
                                   np.uint32(first_index),
                                   np.uint32(elements_this_iter),
                                   triangles_dev.data, vertices_dev.data,
                                   world_origin, world_scale,
                                   nodes.data, morton_codes.data, g_times_l=True ).wait()

        morton_codes_host = morton_codes.get() >> (16 - morton_bits)

    return world_coords, nodes.get(), morton_codes_host

def merge_nodes_detailed(nodes, first_child, nchild):
    '''Merges nodes into len(first_child) parent nodes, using
    the provided arrays to determine the index of the first
    child of each parent, and how many children there are.'''
    nthreads_per_block = 256
    context = None
    queue = None
    if gpuapi.is_gpu_api_opencl():
        context = cltools.get_last_context()
        #print context
        queue = cl.CommandQueue( context )

    # Load GPU functions
    if gpuapi.is_gpu_api_cuda():
        bvh_module = get_module('bvh.cu', options=api_options, include_source_directory=True)
    elif gpuapi.is_gpu_api_opencl():
        # don't like the last context method. trouble. trouble.
        bvh_module = get_module('bvh.cl', context, options=api_options, include_source_directory=True)
    else:
        raise RuntimeError('API is neither CUDA nor OpenCL?!')
    bvh_funcs = GPUFuncs(bvh_module)

    # Load Memory
    if gpuapi.is_gpu_api_cuda():
        gpu_nodes = ga.to_gpu(nodes)
        gpu_first_child = ga.to_gpu(first_child.astype(np.int32))
        gpu_nchild = ga.to_gpu(nchild.astype(np.int32))

        nparent = len(first_child)
        gpu_parent_nodes = ga.empty(shape=nparent, dtype=ga.vec.uint4)
    elif gpuapi.is_gpu_api_opencl():
        gpu_nodes = ga.to_device(queue, nodes)
        gpu_first_child = ga.to_device(queue, first_child.astype(np.int32))
        gpu_nchild = ga.to_device(queue, nchild.astype(np.int32))
        nparent = len(first_child)
        parent_nodes_np = np.zeros(shape=nparent, dtype=ga.vec.uint4)
        gpu_parent_nodes = ga.to_device( queue, parent_nodes_np )
    else:
        raise RuntimeError('API is neither CUDA nor OpenCL?!')

    # Run Kernel
    for first_index, elements_this_iter, nblocks_this_iter in \
            chunk_iterator(nparent, nthreads_per_block, max_blocks=10000):
        if gpuapi.is_gpu_api_cuda():
            bvh_funcs.make_parents_detailed(np.uint32(first_index),
                                            np.uint32(elements_this_iter),
                                            gpu_nodes,
                                            gpu_parent_nodes,
                                            gpu_first_child,
                                            gpu_nchild,
                                            block=(nthreads_per_block,1,1),
                                            grid=(nblocks_this_iter,1))
        elif gpuapi.is_gpu_api_opencl():
            bvh_funcs.make_parents_detailed( queue, (elements_this_iter,1,1), None,
                                             np.uint32(first_index),
                                             np.uint32(elements_this_iter),
                                             gpu_nodes.data,
                                             gpu_parent_nodes.data,
                                             gpu_first_child.data,
                                             gpu_nchild.data ).wait()
        else:
            raise RuntimeError('API is neither CUDA nor OpenCL?!')

    return gpu_parent_nodes.get()

def collapse_chains(nodes, layer_bounds):
    if gpuapi.is_gpu_api_cuda():
        bvh_module = get_module('bvh.cu', options=api_options,include_source_directory=True)
    elif gpuapi.is_gpu_api_opencl():
        context = cltools.get_last_context()
        queue = cl.CommandQueue( context )
        bvh_module = get_module('bvh.cl', context, options=api_options, include_source_directory=True)
    else:
        raise RuntimeError('API neither CUDA or OpenCL')

    bvh_funcs = GPUFuncs(bvh_module)

    if gpuapi.is_gpu_api_cuda():
        gpu_nodes = ga.to_gpu(nodes)
    elif gpuapi.is_gpu_api_opencl():
        gpu_nodes = ga.to_device( queue, nodes )

    bounds = zip(layer_bounds[:-1], layer_bounds[1:])[:-1]
    bounds.reverse()
    nthreads_per_block = 256
    for start, end in bounds:
        if gpuapi.is_gpu_api_cuda():
            bvh_funcs.collapse_child(np.uint32(start),
                                     np.uint32(end),
                                     gpu_nodes,
                                     block=(nthreads_per_block,1,1),
                                     grid=(120,1))
        elif gpuapi.is_gpu_api_opencl():
            bvh_funcs.collapse_child( queue, ( end-start, 1, 1 ), None,
                                      np.uint32(start),
                                      np.uint32(end),
                                      gpu_nodes.data ).wait()

    return gpu_nodes.get()

def area_sort_nodes(gpu_geometry, layer_bounds):
    bvh_module = get_cu_module('bvh.cu', options=cuda_options,
                               include_source_directory=True)
    bvh_funcs = GPUFuncs(bvh_module)

    bounds = zip(layer_bounds[:-1], layer_bounds[1:])[:-1]
    bounds.reverse()
    nthreads_per_block = 256
    for start, end in bounds:
        bvh_funcs.area_sort_child(np.uint32(start),
                                  np.uint32(end),
                                  gpu_geometry,
                                  block=(nthreads_per_block,1,1),
                                  grid=(120,1))
    return gpu_geometry.nodes.get()

def merge_nodes(nodes, degree, max_ratio=None):
    nthreads_per_block = 256
    context = None
    queue = None
    if gpuapi.is_gpu_api_opencl():
        context = cltools.get_last_context()
        queue = cl.CommandQueue( context )

    # Load GPU functions
    if gpuapi.is_gpu_api_cuda():
        bvh_module = get_module('bvh.cu', options=api_options, include_source_directory=True)
    elif gpuapi.is_gpu_api_opencl():
        # don't like the last context method. trouble. trouble.
        bvh_module = get_module('bvh.cl', context, options=api_options, include_source_directory=True)
    else:
        raise RuntimeError('API is neither CUDA nor OpenCL?!')
    bvh_funcs = GPUFuncs(bvh_module)

    # determine number of parents
    nparent = len(nodes) / degree
    if len(nodes) % degree != 0:
        nparent += 1

    if nparent == 1:
        nparent_pad = nparent
    else:
        nparent_pad = round_up_to_multiple(nparent, 1) #degree

    # allocate memory
    if gpuapi.is_gpu_api_cuda():
        gpu_parent_nodes = ga.zeros(shape=nparent_pad, dtype=ga.vec.uint4)
    elif gpuapi.is_gpu_api_opencl():
        parent_nodes_np = np.zeros(shape=nparent, dtype=ga.vec.uint4)
        gpu_parent_nodes = ga.to_device( queue, parent_nodes_np )
        gpu_nodes = ga.to_device(queue,nodes)
    else:
        raise RuntimeError('API is neither CUDA nor OpenCL?!')

    # run kernel
    if gpuapi.is_gpu_api_cuda():
        for first_index, elements_this_iter, nblocks_this_iter in \
                chunk_iterator(nparent, nthreads_per_block, max_blocks=10000):
            bvh_funcs.make_parents(np.uint32(first_index),
                                   np.uint32(elements_this_iter),
                                   np.uint32(degree),
                                   gpu_parent_nodes,
                                   cuda.In(nodes),
                                   np.uint32(0),
                                   np.uint32(len(nodes)),
                                   block=(nthreads_per_block,1,1),
                                   grid=(nblocks_this_iter,1))
    elif gpuapi.is_gpu_api_opencl():
        for first_index, elements_this_iter, nblocks_this_iter in \
                chunk_iterator(nparent, nthreads_per_block, max_blocks=1):
            bvh_funcs.make_parents(queue, (elements_this_iter,1,1), None,
                                   np.uint32(first_index),
                                   np.uint32(elements_this_iter),
                                   np.uint32(degree),
                                   gpu_parent_nodes.data,
                                   gpu_nodes.data,
                                   np.uint32(0),
                                   np.uint32(len(nodes))).wait()
    else:
        raise RuntimeError('API is neither CUDA nor OpenCL?!')

    parent_nodes = gpu_parent_nodes.get()

    if max_ratio is not None:
        areas = node_areas(parent_nodes)
        child_areas = node_areas(nodes)

        excessive_area = np.zeros(shape=len(areas), dtype=bool)
        for i, parent_area in enumerate(areas):
            nchild = parent_nodes['w'][i] >> CHILD_BITS
            child_index = parent_nodes['w'][i] & ~NCHILD_MASK
            child_area = child_areas[child_index:child_index+nchild].sum()
            #if parent_area > 1e9:
            #    print i, 'Children: %e, Parent: %e' % (child_area, parent_area)
            if child_area/parent_area < 0.3:
                excessive_area[i] = True
                #print i, 'Children: %e, Parent: %e' % (child_area, parent_area)

        extra_slots = round_up_to_multiple((degree - 1) * np.count_nonzero(excessive_area), 1)
        print 'Extra slots:', extra_slots
        new_parent_nodes = np.zeros(shape=len(parent_nodes) + extra_slots,
                                    dtype=parent_nodes.dtype)
        new_parent_nodes[:len(parent_nodes)] = parent_nodes

        offset = 0
        for count, index in enumerate(np.argwhere(excessive_area)):
            index = index[0] + offset
            nchild = new_parent_nodes['w'][index] >> CHILD_BITS
            child_index = new_parent_nodes['w'][index] & ~NCHILD_MASK
            new_parent_nodes[index] = nodes[child_index]
            #new_parent_nodes['w'][index] = 1 << CHILD_BITS | child_index
            tmp_nchild = new_parent_nodes['w'][index] >> CHILD_BITS
            tmp_child_index = new_parent_nodes['w'][index] & ~NCHILD_MASK
            new_parent_nodes['w'][index] = tmp_nchild << CHILD_BITS | (tmp_child_index + len(nodes))

            if nchild == 1:
                continue

            # slide everyone over
            #print index, nchild, len(new_parent_nodes)
            new_parent_nodes[index+nchild:] = new_parent_nodes[index+1:-nchild+1]
            offset += nchild - 1
            for sibling in xrange(nchild - 1):
                new_parent_index = index + 1 + sibling
                new_parent_nodes[new_parent_index] = nodes[child_index + sibling + 1]
                if new_parent_nodes['x'][new_parent_index] != 0:
                    tmp_nchild = new_parent_nodes['w'][new_parent_index] >> CHILD_BITS
                    tmp_child_index = new_parent_nodes['w'][new_parent_index] & ~NCHILD_MASK
                    new_parent_nodes['w'][new_parent_index] = tmp_nchild << CHILD_BITS | (tmp_child_index + len(nodes))

                    #new_parent_nodes['w'][new_parent_index] = 1 << CHILD_BITS | (child_index + sibling + 1)

            #print 'intermediate: %e' % node_areas(new_parent_nodes).max()
        print 'old: %e' % node_areas(parent_nodes).max()
        print 'new: %e' % node_areas(new_parent_nodes).max()
        if len(new_parent_nodes) < len(nodes):
            # Only adopt new set of parent nodes if it actually reduces the
            # total number of nodes at this level by 1.
            parent_nodes = new_parent_nodes

    return parent_nodes

def concatenate_layers(layers):
    nthreads_per_block = 1024
    context = None
    queue = None
    if gpuapi.is_gpu_api_opencl():
        context = cltools.get_last_context()
        #print context
        queue = cl.CommandQueue( context )

    # Load GPU functions
    if gpuapi.is_gpu_api_cuda():
        bvh_module = get_module('bvh.cu', options=api_options, include_source_directory=True)
    elif gpuapi.is_gpu_api_opencl():
        # don't like the last context method. trouble. trouble.
        bvh_module = get_module('bvh.cl', cltools.get_last_context(), options=api_options, include_source_directory=True)
    else:
        raise RuntimeError('API neither CUDA nor OpenCL?!')
    bvh_funcs = GPUFuncs(bvh_module)

    # Put 0 at beginning of list
    layer_bounds = np.insert(np.cumsum(map(len, layers)), 0, 0)

    # allocate memory
    if gpuapi.is_gpu_api_cuda():
        nodes = ga.empty(shape=int(layer_bounds[-1]), dtype=ga.vec.uint4)
    elif gpuapi.is_gpu_api_opencl():
        totsize = 0
        layer_pos = []
        print layer_bounds[-1]
        for n,layer in enumerate(layers):
            layer_pos.append( totsize )
            print "LAYER ",n," size=",len(layer),"start=",totsize
            totsize += len(layer)
        print "totsize: ",totsize
        nodes_iter_np = np.empty(totsize,dtype=ga.vec.uint4)
        nodes_iter_gpu = ga.to_device(queue,nodes_iter_np)
        nodeset_np = []
    else:
        raise RuntimeError('API neither CUDA nor OpenCL?!')

    ilayer = 0
    for layer_start, layer_end, layer in zip(layer_bounds[:-1],
                                             layer_bounds[1:],
                                             layers):
        if layer_end == layer_bounds[-1]:
            # leaf nodes need no offset
            child_offset = 0
        else:
            child_offset = layer_end

        #print "ilayer,start,end,child_offset: ",ilayer,layer_start, layer_end, child_offset
        nmax_blocks = 10000
        if gpuapi.is_gpu_api_opencl():
            nthreads_per_block = 256
            nmax_blocks = 1
        for first_index, elements_this_iter, nblocks_this_iter in \
                chunk_iterator(layer_end-layer_start, nthreads_per_block,max_blocks=nmax_blocks):
            #print "   ",ilayer,first_index, elements_this_iter, nblocks_this_iter, layer_start
            if gpuapi.is_gpu_api_cuda():
                bvh_funcs.copy_and_offset(np.uint32(first_index),
                                          np.uint32(elements_this_iter),
                                          np.uint32(child_offset),
                                          cuda.In(layer),
                                          nodes[layer_start:],
                                          block=(nthreads_per_block,1,1),
                                          grid=(nblocks_this_iter,1))
            elif gpuapi.is_gpu_api_opencl():
                layer_gpu = ga.to_device(queue,layer)
                bvh_funcs.copy_and_offset( queue, (elements_this_iter,1,1), (1,1,1),
                                           np.uint32(first_index),
                                           np.uint32(elements_this_iter),
                                           np.uint32(child_offset),
                                           np.uint32(layer_start),
                                           layer_gpu.data,
                                           nodes_iter_gpu.data, g_times_l=True ).wait()
            else:
                raise RuntimeError('API neither CUDA nor OpenCL?!')
        ilayer += 1

    if gpuapi.is_gpu_api_cuda():
        return nodes.get(), layer_bounds
    elif gpuapi.is_gpu_api_opencl():
        return nodes_iter_gpu.get(), layer_bounds

def optimize_layer(orig_nodes):
    bvh_module = get_cu_module('bvh.cu', options=cuda_options,
                               include_source_directory=True)
    bvh_funcs = GPUFuncs(bvh_module)

    nodes = ga.to_gpu(orig_nodes)
    n = len(nodes)
    areas = ga.empty(shape=n/2, dtype=np.uint64)
    nthreads_per_block = 128

    min_areas = ga.empty(shape=int(np.ceil(n/float(nthreads_per_block))), dtype=np.uint64)
    min_index = ga.empty(shape=min_areas.shape, dtype=np.uint32)

    update = 10000

    skip_size = 1
    flag = cutools.mapped_empty(shape=skip_size, dtype=np.uint32)

    i = 0
    skips = 0
    swaps = 0
    while i < n/2 - 1:
        # How are we doing?
        if i % update == 0:
            for first_index, elements_this_iter, nblocks_this_iter in \
                    chunk_iterator(n/2, nthreads_per_block, max_blocks=10000):

                bvh_funcs.pair_area(np.uint32(first_index),
                                    np.uint32(elements_this_iter),
                                    nodes,
                                    areas,
                                    block=(nthreads_per_block,1,1),
                                    grid=(nblocks_this_iter,1))

            areas_host = areas.get()
            #print nodes.get(), areas_host.astype(float)
            print 'Area of parent layer so far (%d): %1.12e' % (i*2, areas_host.astype(float).sum())
            print 'Skips: %d, Swaps: %d' % (skips, swaps)

        test_index = i * 2

        blocks = 0
        look_forward = min(8192*50, n - test_index - 2)
        skip_this_round = min(skip_size, n - test_index - 1)
        flag[:] = 0
        for first_index, elements_this_iter, nblocks_this_iter in \
                chunk_iterator(look_forward, nthreads_per_block, max_blocks=10000):
            bvh_funcs.min_distance_to(np.uint32(first_index + test_index + 2),
                                      np.uint32(elements_this_iter),
                                      np.uint32(test_index),
                                      nodes,
                                      np.uint32(blocks),
                                      min_areas,
                                      min_index,
                                      cutools.Mapped(flag),
                                      block=(nthreads_per_block,1,1),
                                      grid=(nblocks_this_iter, skip_this_round))
            blocks += nblocks_this_iter
            #print i, first_index, nblocks_this_iter, look_forward
        cuda.Context.get_current().synchronize()

        if flag[0] == 0:
            flag_nonzero = flag.nonzero()[0]
            if len(flag_nonzero) == 0:
                no_swap_required = skip_size
            else:
                no_swap_required = flag_nonzero[0]
            i += no_swap_required
            skips += no_swap_required
            continue

        min_areas_host = min_areas[:blocks].get()
        min_index_host = min_index[:blocks].get()
        best_block = min_areas_host.argmin()
        better_i = min_index_host[best_block]

        swaps += 1
        #print 'swap', test_index+1, better_i
        assert 0 < better_i < len(nodes)
        assert 0 < test_index + 1 < len(nodes)
        bvh_funcs.swap(np.uint32(test_index+1), np.uint32(better_i),
                       nodes, block=(1,1,1), grid=(1,1))
        cuda.Context.get_current().synchronize()
        i += 1

    for first_index, elements_this_iter, nblocks_this_iter in \
            chunk_iterator(n/2, nthreads_per_block, max_blocks=10000):

        bvh_funcs.pair_area(np.uint32(first_index),
                            np.uint32(elements_this_iter),
                            nodes,
                            areas,
                            block=(nthreads_per_block,1,1),
                            grid=(nblocks_this_iter,1))

    areas_host = areas.get()

    print 'Final area of parent layer: %1.12e' % areas_host.sum()
    print 'Skips: %d, Swaps: %d' % (skips, swaps)

    return nodes.get()


import sys
import os
import argparse
import param
import pickle
import numpy as np
from threading import Thread

def Run(args):
    # create a Clairvoyante
    if args.v2 == True:
        import utils_v2 as utils
        if args.slim == True:
            import clairvoyante_v2_slim as cv
        else:
            import clairvoyante_v2 as cv
    elif args.v3 == True:
        import utils_v2 as utils # v3 network is using v2 utils
        if args.slim == True:
            import clairvoyante_v3_slim as cv
        else:
            import clairvoyante_v3 as cv
    utils.SetupEnv()
    m = cv.Clairvoyante()
    m.init()

    CalcAll(args, m, utils)

def CalcAll(args, m, utils):
    if args.bin_fn != None:
        with open(args.bin_fn, "rb") as fh:
            total = pickle.load(fh)
            XArrayCompressed = pickle.load(fh)
            YArrayCompressed = pickle.load(fh)
            posArrayCompressed = pickle.load(fh)
    else:
        total, XArrayCompressed, YArrayCompressed, posArrayCompressed = \
        utils.GetTrainingArray(args.tensor_fn,
                               args.var_fn,
                               args.bed_fn)

    predictBatchSize = param.predictBatchSize
    trainingTotal = int(total*param.trainingDatasetPercentage)
    validationStart = trainingTotal + 1
    numValItems = total - validationStart

    for n in args.chkpnt_fn:
        m.restoreParameters(os.path.abspath(n))
        datasetPtr = 0
        trainingLost = 0
        validationLost = 0
        i = 1
        XBatch, XNum, XEndFlag = utils.DecompressArray(XArrayCompressed, datasetPtr, predictBatchSize, total)
        YBatch, YNum, YEndFlag = utils.DecompressArray(YArrayCompressed, datasetPtr, predictBatchSize, total)
        datasetPtr += XNum
        while True:
            threadPool = []
            threadPool.append(Thread(target=m.getLossNoRT, args=(XBatch, YBatch, )))
            for t in threadPool: t.start()
            predictBatchSize = param.predictBatchSize
            if datasetPtr < validationStart and (validationStart - datasetPtr) < predictBatchSize:
                predictBatchSize = validationStart - datasetPtr
            elif datasetPtr >= validationStart and (datasetPtr % predictBatchSize) != 0:
                predictBatchSize = predictBatchSize - (datasetPtr % predictBatchSize)
            #print >> sys.stderr, "%d\t%d\t%d\t%d" % (datasetPtr, predictBatchSize, validationStart, total)
            XBatch2, XNum2, XEndFlag2 = utils.DecompressArray(XArrayCompressed, datasetPtr, predictBatchSize, total)
            YBatch2, YNum2, YEndFlag2 = utils.DecompressArray(YArrayCompressed, datasetPtr, predictBatchSize, total)
            if XNum != YNum or XEndFlag != YEndFlag:
                sys.exit("Inconsistency between decompressed arrays: %d/%d" % (XNum, YNum))
            for t in threadPool: t.join()
            XBatch = XBatch2; YBatch = YBatch2
            if datasetPtr >= validationStart: validationLost += m.getLossLossRTVal
            else: trainingLost += m.getLossLossRTVal
            if XEndFlag2!= 0:
                m.getLossNoRT( XBatch, YBatch )
                if datasetPtr >= validationStart: validationLost += m.getLossLossRTVal
                else: trainingLost += m.getLossLossRTVal
                print >> sys.stderr, "%s\t%.10f\t%.10f" % (n, trainingLost/trainingTotal, validationLost/numValItems)
                break;
            i += 1
            datasetPtr += XNum2

def main():
    parser = argparse.ArgumentParser(
            description="Calculate the loss different between training dataset and validation dataset" )

    parser.add_argument('--bin_fn', type=str, default = None,
            help="Binary tensor input generated by tensor2Bin.py, tensor_fn, var_fn and bed_fn will be ignored")

    parser.add_argument('--tensor_fn', type=str, default = "vartensors",
            help="Tensor input")

    parser.add_argument('--var_fn', type=str, default = "truthvars",
            help="Truth variants list input")

    parser.add_argument('--bed_fn', type=str, default = None,
            help="High confident genome regions input in the BED format")

    parser.add_argument('--chkpnt_fn', nargs='+', type=str, default = None,
            help="Input a list of checkpoint for calculation")

    parser.add_argument('--v3', type=param.str2bool, nargs='?', const=True, default = True,
            help="Use Clairvoyante version 3")

    parser.add_argument('--v2', type=param.str2bool, nargs='?', const=True, default = False,
            help="Use Clairvoyante version 2")

    parser.add_argument('--slim', type=param.str2bool, nargs='?', const=True, default = False,
            help="Train using the slim version of Clairvoyante, optional")

    args = parser.parse_args()

    if len(sys.argv[1:]) == 0:
        parser.print_help()
        sys.exit(1)

    Run(args)

if __name__ == "__main__":
    main()