#  Copyright (C) 2012 Matt Hagy <hagy@gatech.edu>
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

from __future__ import division

from functools import wraps

import numpy as np
try:
    from scipy.optimize import fmin_cg
except ImportError:
    def fmin_cg(*args, **kwds):
        raise RuntimeError('fmin_cg not available; install scipy to use this functionality')

from base_components import (NeighborsTable, BaseForceField, LJForceField, BaseConfig,
                             BasePairCorrelationFunctionCalculator,
                             BaseMeanSquareDisplacementCalculator,
                             BaseVelocityAutocorrelationCalculator)
from util import periodic_distances

# TODO
# Separate simulation and analysis code into separate modules.
# Further separate analysis code into static vs. dynamic analysis.

__all__ = ['NeighborsTable', 'BaseForceField', 'LJForceField',
           'Config', 'NeighborsTableTracker', 'EnergyMinimzer',
           'MDSimulator',
           'StaticPairCorrelation',
           'StaticPairCorrelationCalculator',
           'StaticPairCorrelationIntegrator',
           'MeanSquareDisplacementCalculator',
           'VelocityAutocorrelationCalculator']

def create_random_state(op=None):
    if isinstance(op, np.random.RandomState):
        return op
    return np.random.RandomState(op)

def create_velocities(N, T=1.0, mass=1.0, random_state=None):
    return create_random_state(random_state).normal(scale=np.sqrt(T / mass), size=(N, 3))

def calculate_box_size(N, sigma, rho):
    '''Calculate the lateral size for a cubic box of N
       particles, each of diameter sigma, such that the particle
       density is rho = N*sigma**3/V
    '''
    V = N * sigma**3 / rho
    return V**(1/3)

def create_hcp_positions(l):
    '''Create a hexagonal close-packed (hcp) lattice of sphere positions that are
       contained within a cubic box of lateral size l (each sphere has unit diameter).
    '''
    r = 0.5
    row_height_shift = r * np.sqrt(3)
    plane_height_shift = np.sqrt(6) * r * 2.0 / 3.0
    n_row = int(np.floor(l))

    row0 = np.array([np.arange(n_row), np.zeros(n_row), np.zeros(n_row)]).T
    plane0 = np.array([row0 + np.array([r if i%2==1 else 0, i*row_height_shift, 0])
                       for i in xrange(int(np.floor(l / row_height_shift)))])
    planes = np.array([plane0 + np.array([r if i%2==1 else 0,
                                         np.sqrt(3)/3*r if i%2==1 else 0,
                                          i * plane_height_shift])
                       for i in xrange(int(np.floor(l / plane_height_shift)))])

    sites = planes.reshape((np.prod(planes.shape[:3:]), 3))
    return sites

class Config(BaseConfig):
    '''Represents the current static and dynamic state of isotropic
       particle system within a cubic periodic box.
    '''

    @classmethod
    def create(cls, N, rho, dt, sigma=1.0, T=1.0, mass=1.0, random_state=None):
        '''Create a new random config with N particle at density rho.
        '''
        random_state = create_random_state(random_state)
        box_size = calculate_box_size(N, sigma, rho)
        hcp_positions = sigma * create_hcp_positions(box_size / sigma)
        if len(hcp_positions) < N:
            raise ValueError("config cannot be initialized from hcp lattice")
        hcp_indices = np.arange(len(hcp_positions))
        random_state.shuffle(hcp_indices)
        positions = hcp_positions[hcp_indices[:N:]]
        inst = cls(positions, last_positions=None, box_size=box_size, dt=dt, sigma=sigma)
        inst.randomize_velocities(T=T, mass=mass, random_state=random_state)
        return inst

    def copy(self):
        return self.__class__(positions=self.positions.copy(),
                              last_positions=self.last_positions.copy(),
                              box_size=self.box_size,
                              dt=self.dt,
                              sigma=self.sigma)

    def with_new_positions(self, new_positions):
        cp = self.__class__(positions=new_positions, last_positions=None,
                            box_size=self.box_size, dt=self.dt, sigma=self.sigma)
        cp.set_velocities(self.calculate_velocities())
        return cp

    def randomize_velocities(self, **kwds):
        self.set_velocities(create_velocities(self.N, **kwds))

    def set_velocities(self, velocities):
        self.last_positions = self.positions - self.dt * velocities

    def calculate_velocities(self):
        return (self.positions - self.last_positions) / self.dt

    def calculate_rms_velocity(self):
        v = self.calculate_velocities()
        return (v**2).sum(axis=1).mean()**0.5

    def calculate_kinetic_energy(self, mass):
        v_rms = self.calculate_rms_velocity()
        return 0.5 * mass * v_rms**2

    def calculate_temperature(self, mass):
        KE = self.calculate_kinetic_energy(mass)
        return 2.0 / 3.0 * KE

    @property
    def N(self):
        return len(self.positions)

    @property
    def V(self):
        return self.box_size**3

    @property
    def rho(self):
        return self.N / self.V

    def rescale_boxsize(self, new_boxsize):
        velocities = self.calculate_velocities()
        self.positions %= self.box_size
        self.positions *= new_boxsize / self.box_size
        self.box_size = new_boxsize
        self.set_velocities(velocities)

    def rescale_boxsize_rho(self, new_rho):
        self.rescale_boxsize(calculate_box_size(self.N, self.sigma, new_rho))

    def normalize_positions(self):
        velocities = self.calculate_velocities()
        self.positions %= self.box_size
        self.set_velocities(velocities)

    def change_dt(self, dt):
        velocities = self.calculate_velocities()
        self.dt = dt
        self.set_velocities(velocities)

    def propagate(self, forces, mass=1.0):
        positions = -self.last_positions + 2 * self.positions + (self.dt**2 / mass) * forces
        self.last_positions = self.positions
        self.positions = positions

    def __reduce__(self):
        return (Config, (self.positions, self.last_positions, self.box_size, self.dt, self.sigma))

class NeighborsTableTracker(object):

    def __init__(self, neighbors_table, box_size):
        self.neighbors_table = neighbors_table
        self.box_size = box_size
        self.last_positions = None

    def reset(self, current_positions):
        self.neighbors_table.rebuild_neighbors(current_positions, self.box_size)
        self.last_positions = current_positions
        if self.acc_delta is None:
            self.acc_delta = np.empty_like(current_positions)
        self.acc_delta.fill(0.0)

    def moved(self, current_positions, check_were_valid=True):
        delta = periodic_distances(current_positions, self.last_positions, self.box_size)
        self.last_positions = current_positions
        self.acc_delta += delta
        dr_max = (self.acc_delta**2).sum(axis=1).max()**0.5
        r_acceptable = 0.5 * self.neighbors_table.r_skin
        if dr_max > r_acceptable:
            return self.rebuild_neighbors(current_positions, check_were_valid=check_were_valid)
        return True

    acc_delta = None

    def rebuild_neighbors(self, current_positions, check_were_valid=True):
        if check_were_valid:
            old_vital_neighbors = self.find_vital_neighbors(current_positions)
        self.neighbors_table.rebuild_neighbors(current_positions, self.box_size)
        if check_were_valid:
            new_vital_neighbors = self.find_vital_neighbors(current_positions)
            were_valid = old_vital_neighbors >= new_vital_neighbors
        self.acc_delta.fill(0.0)
        if check_were_valid:
            return were_valid

    def find_vital_neighbors(self, positions):
        return self.neighbors_table.find_set_of_neighbors_within_distance(self.neighbors_table.r_forcefield_cutoff,
                                                                          positions, self.box_size)

class EnergyMinimzer(object):

    def __init__(self, config_init, forcefield, maxiter=100, neighbors_table=None, neighbors_table_skin=1.0):
        self.config_init = config_init
        self.forcefield = forcefield
        self.maxiter = maxiter

        if neighbors_table is None:
            neighbors_table = NeighborsTable(
                r_forcefield_cutoff=forcefield.r_cutoff,
                r_skin=neighbors_table_skin)
        self.neighbors_table = neighbors_table

        self.forces = np.empty_like(self.config_init.positions)
        self.neighbors_table_tracker = NeighborsTableTracker(self.neighbors_table, self.config_init.box_size)

    def minimize(self):
        self.neighbors_table_tracker.reset(self.config_init.positions)
        [x_final, self.U_min, self.n_func_calls, self.n_grad_calls, self.warnfalgs
         ] = fmin_cg(self.evaluate_potential,
                     self.config_init.positions,
                     fprime=self.evaluate_gradient,
                     maxiter=self.maxiter,
                     callback=self.callback,
                     full_output=True, disp=False)
        self.config_final = self.config_init.with_new_positions(self.create_positions(x_final))
        return self.config_final

    def evaluate_potential(self, x):
        return self.forcefield.evaluate_potential(self.create_positions(x),
                                                  self.config_init.box_size,
                                                  self.neighbors_table)

    def evaluate_gradient(self, x):
        self.forcefield.evaluate_forces(self.forces,
                                        self.create_positions(x),
                                        self.config_init.box_size,
                                        self.neighbors_table)
        return -self.forces.reshape(3 * self.config_init.N)

    def callback(self, x):
        self.neighbors_table_tracker.moved(self.create_positions(x))

    def create_positions(self, x):
        return x.reshape((self.config_init.N, 3)) % self.config_init.box_size

class MDSimulator(object):

    kB = 1.0

    def __init__(self, config, forcefield, mass=1.0, r_skin=1.0):
        self.config = config
        self.forcefield = forcefield
        self.forces = np.empty_like(config.positions)
        self.mass = mass
        self.neighbors_table = NeighborsTable(r_forcefield_cutoff=self.forcefield.r_cutoff,
                                              r_skin=r_skin)
        self.neighbors_table_tracker = NeighborsTableTracker(self.neighbors_table, self.config.box_size)
        self.neighbors_table_tracker.reset(self.config.positions)

        self.backup_positions = np.empty_like(self.config.positions)
        self.backup_last_positions = np.empty_like(self.config.last_positions)

    normalize_positions_rate = 20

    def cycle(self, n=1):
        for i in xrange(n):
            if not i%self.normalize_positions_rate:
                self.normalize_positions()

            self.backup_positions[...] = self.config.positions
            self.backup_last_positions[...] = self.config.last_positions
            were_neighbors_valid = self.propagate_attempt()
            if not were_neighbors_valid:
                self.config.positions[...] = self.backup_positions
                self.config.last_positions[...] = self.backup_last_positions
                were_neighbors_valid = self.propagate_attempt()
                if not were_neighbors_valid:
                    raise RuntimeError("couldn't create valid neighbor list, increase r_skin or decrease dt")

    def propagate_attempt(self):
        self.forcefield.evaluate_forces(self.forces,
                                        self.config.positions % self.config.box_size,
                                        self.config.box_size, self.neighbors_table)
        self.config.propagate(self.forces, self.mass)
        return self.neighbors_table_tracker.moved(self.config.positions, check_were_valid=True)

    def normalize_positions(self):
        self.config.normalize_positions()
        self.neighbors_table_tracker.reset(self.config.positions)

    def rescale_boxsize_rho(self, rho):
        self.config.rescale_boxsize_rho(rho)
        self.neighbors_table_tracker.reset(self.config.positions)

    def compute_potential_energy(self):
        '''Compute the internal potential energy due to the force field
           interactions. Result are inversely scaled by the number of particles.
        '''
        return self.forcefield.evaluate_potential(self.config.positions % self.config.box_size,
                                                  self.config.box_size, self.neighbors_table)

    def compute_kinetic_energy(self):
        '''Compute the kinetic energy due to translational veloctiy of particles.
           Result are inversely scaled by the number of particles.
        '''
        return self.config.calculate_kinetic_energy(self.mass)

    def compute_energy(self):
        '''Compute the total energy (sum kinetic and potential) of the system.
           Result are inversely scaled by the number of particles.
        '''
        return self.compute_potential_energy() + self.compute_kinetic_energy()

    def compute_temperature(self):
        return self.config.calculate_temperature(self.mass)

    def compute_excess_pressure(self, correct_long_range):
        '''Compute the pressure due to particle pair interactions.
           Result are inversely scaled by the number of particles.
        '''
        v = 1.0 / 3.0 * self.forcefield.evaluate_virial_sum(
            self.config.positions % self.config.box_size,
            self.config.box_size, self.neighbors_table)
        if correct_long_range:
            v += self.forcefield.long_range_virial_correction()
        return (self.kB / self.config.V) * v

    def compute_ideal_pressure(self):
        '''Compute pressure due to kinetic energy alone.
           Result are inversely scaled by the number of particles.
        '''
        return self.kB * self.compute_temperature() / self.config.V

    def compute_total_pressure(self, correct_long_range=True):
        '''Compute the total (sum kinetic and potential) pressure.
           Result are inversely scaled by the number of particles.
        '''
        return self.compute_excess_pressure(correct_long_range) + self.compute_ideal_pressure()

    def compute_virial(self, correct_long_range=True):
        P = self.compute_total_pressure(correct_long_range)
        return P * self.config.V / (self.kB * self.compute_temperature())

    def compute_excess_virial(self, correct_long_range):
        v = ((3 * self.kB * self.compute_temperature())**-1 *
             self.forcefield.evaluate_virial_sum(self.config.positions % self.config.box_size,
                                                 self.config.box_size, self.neighbors_table))
        if correct_long_range:
            v += self.forcefield.long_range_virial_correction()
        return v

    # Old (deprecated) methods
    def evaluate_potential(self):
        return self.compute_potential_energy()

    def evaluate_hamiltonian(self):
        return self.compute_energy()

    def minimize(self, maxiter=10):
        minimizer = EnergyMinimzer(self.config, self.forcefield, maxiter=8,
                                   neighbors_table=self.neighbors_table)
        self.config = minimizer.minimize()
        return minimizer.U_min

    def minimize_until(self, cutoff, verbose=True):
        while True:
            U_min = self.minimize(100)
            if verbose:
                print '%.4e' % U_min
            if U_min < cutoff:
                return U_min

    def get_config(self):
        config = self.config.copy()
        config.normalize_positions()
        return config

def cached_property(name_or_func, cache_name=None):
    '''Wrapper to create a cached readonly property for a class.
    '''
    if isinstance(name_or_func, basestring):
        return lambda func: cached_property(func, cache_name=name_or_func)

    func = name_or_func
    assert callable(func)

    if cache_name is None:
        cache_name = '_' + func.func_name
    assert isinstance(cache_name, basestring)

    @property
    @wraps(func)
    def wrapper(self):
        try:
            return getattr(self, cache_name)
        except AttributeError:
            value = func(self)
            setattr(self, cache_name, value)
            return value
    return wrapper

class StaticPairCorrelation(object):
    '''Describes the static correlation of isotropic particle pairs.
          e.g. The numerical analogs of the standard g(r) and h(r) functions.

       The underlying data is a histogram of pair separation distances at
       fixed spacing dr (i.e. the i-th bin contains the total number of
       pair observed at distances of i*dr <= r < (i+1)*dr.

       Additionally, the separation distances can be shifted by an r_offset.
       This is useful in avoiding the avoid the leading zeros that correspond
       to unphysically close pair distances.
    '''

    def __init__(self, pair_distance_histogram, dr, r_offset=0.0):
        pair_distance_histogram = np.asarray(pair_distance_histogram)
        assert pair_distance_histogram.ndim == 1
        assert pair_distance_histogram.size > 0
        self.pair_distance_histogram = pair_distance_histogram
        self.dr = dr
        self.r_offset = r_offset

    @cached_property
    def r_lower(self):
        '''Radial distance corresponding to correlation distances at
           the lower bound of each bin.
        '''
        return self.r_offset + self.dr * np.arange(self.pair_distance_histogram.size)

    @cached_property
    def r_mid(self):
        '''Radial distance corresponding to correlation distances at
           center of each bin.
        '''
        return self.r_lower + 0.5 * self.dr

    @property
    def r(self):
        return self.r_mid

    @cached_property
    def g(self):
        '''Reducued density pair correlations
        '''
        N = self.pair_distance_histogram.sum()
        if not N:
            return None

        r_max = self.pair_distance_histogram.size * self.dr
        V = 4.0 / 3.0 * np.pi * r_max**3
        rho = N / V
        v = 4.0 / 3.0 * np.pi * ((self.r_lower + self.dr)**3 - self.r_lower**3)
        rhos =  self.pair_distance_histogram / v
        return rhos / rho

    @cached_property
    def h(self):
        '''Shifted reducued density pair correlations
        '''
        if self.g is None:
            return None
        return self.g - 1.0

class StaticPairCorrelationCalculator(BasePairCorrelationFunctionCalculator):
    '''Calculate the PairCorrelationData from configurations (Config) of
       particles. The calculation is performed by accumulating a histogram
       of pair separation distances one configuration at a time. The
       intermediate state of calculation can be saved by pickeling the
       calculator object.
    '''

    def accumulate_config(self, config):
        self.accumulate_positions(config.positions, config.box_size)

    def get_accumulated(self):
        return StaticPairCorrelation(self.bins.copy(), self.dr, self.r_min)

class StaticPairCorrelationIntegrator(object):
    '''Calculate thermodynamic properties of an isotropic particle system
       by integrating over the sampled pair correlation (StaticPairCorrelation object)
       for the system. Uses the potential and gradient of the force field associated
       with the system.
    '''

    def __init__(self, pair_correlation, forcefield, rho, beta):
        self.pair_correlation = pair_correlation
        self.forcefield = forcefield
        self.rho = rho
        self.beta = beta

    def integrate_g_product_over_space_ex(self, func, mask=Ellipsis):
        '''Numerically integrate
            \int g(r) * r**2 * func(r)

           The mask argument allows the specification of which
           elements of data arrays to include (i.e. allows the
           exclusion of zero elements)
        '''
        r = self.pair_correlation.r[mask]
        g = self.pair_correlation.g[mask]
        return np.trapz(r**2 * g * func(r), r)

    @cached_property
    def where_g_nonzero(self):
        return self.pair_correlation.g != 0.0

    @cached_property
    def where_g_nonzero_and_in_cutoff(self):
        return (self.pair_correlation.g != 0.0) & (self.pair_correlation.r <= self.forcefield.r_cutoff)

    def integrate_g_product_over_space(self, func):
        return self.integrate_g_product_over_space_ex(func, self.where_g_nonzero)

    def integrate_g_product_over_space_in_cutoff(self, func):
        return self.integrate_g_product_over_space_ex(func, self.where_g_nonzero_and_in_cutoff)

    def calculate_excess_internal_energy(self):
        return (2.0 * np.pi * self.rho *
                self.integrate_g_product_over_space_in_cutoff(self.forcefield.evaluate_potential_function))

    def calculate_virial(self, correct_long_range=True):
        v = 1.0 - 2.0 / 3.0 * np.pi * self.beta * self.rho * (
            self.integrate_g_product_over_space_in_cutoff(
            lambda r: r * -self.forcefield.evaluate_scalar_force_function(r)))
        if correct_long_range:
            v += self.forcefield.long_range_virial_correction(self.pair_correlation.r.max())
        return v

class BaseTimeCorelationCalculator(object):
    '''Base class for Time Correlation Function (TCF) computations
    '''

    def __init__(self, window_size, N_particles, *args, **kwds):
        assert 'analyze_rate' in kwds
        self.analyze_rate = kwds.pop('analyze_rate')
        super(BaseTimeCorelationCalculator, self).__init__(window_size, N_particles, *args, **kwds)

    @classmethod
    def create(cls, window_size, N_particles, analyze_rate=1):
        # ensure analyze_rate is passed as a keyword
        return cls(window_size, N_particles, analyze_rate=analyze_rate)

    def compute_time(self):
        return self.analyze_rate * np.arange(self.window_size)

class MeanSquareDisplacementCalculator(BaseTimeCorelationCalculator, BaseMeanSquareDisplacementCalculator):
    '''Compute the mean square displacment TCF; i.e. th self-positional TCF
    '''

    def analyze_config(self, config):
        self.analyze_positions(config.positions, config.box_size)

    def compute_msd(self):
        n_acc = self.calculate_n_accumulates()
        if not n_acc:
            return None

        return self.acc_msd_data / float(n_acc * self.N_particles)

    def __reduce__(self):
        return (create_msdc,
                (self.window_size, self.N_particles,
                 self.analyze_rate,
                 self.n_positions_seen,
                 self.displacement_window, self.last_positions,
                 self.acc_msd_data))

def create_msdc(window_size, N_particles, analyze_rate, n_positions_seen,
                displacement_window, last_positions, acc_msd_data):
    return MeanSquareDisplacementCalculator(window_size, N_particles,
                                            analyze_rate=analyze_rate,
                                            n_positions_seen=n_positions_seen,
                                            displacement_window=displacement_window,
                                            last_positions=last_positions,
                                            acc_msd_data=acc_msd_data)

class VelocityAutocorrelationCalculator(BaseTimeCorelationCalculator, BaseVelocityAutocorrelationCalculator):
    '''Compute the velocity autocorrelation TCF (VACF); i.e. the self-velocity TCF
    '''

    def analyze_config(self, config):
        self.analyze_velocities(config.calculate_velocities())

    def compute_vacf(self):
        n_acc = self.calculate_n_accumulates()
        if not n_acc:
            return None

        return self.acc_correlations / float(n_acc * self.N_particles)

    def __reduce__(self):
        return (create_vacfc,
                (self.window_size, self.N_particles,
                 self.analyze_rate,
                 self.n_velocities_seen,
                 self.velocities_windows, self.acc_correlations))

def create_vacfc(window_size, N_particles, analyze_rate, n_velocities_seen,
                 velocities_windows, acc_correlations):
    return VelocityAutocorrelationCalculator(window_size, N_particles,
                                             analyze_rate=analyze_rate,
                                             n_velocities_seen=n_velocities_seen,
                                             velocities_windows=velocities_windows,
                                             acc_correlations=acc_correlations)


# Copyright 2020 The FedLearner Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# coding: utf-8

import os
import math
import queue
import logging
import multiprocessing as mp
import numpy as np
from google.protobuf import text_format

from fedlearner.model.tree.loss import LogisticLoss
from fedlearner.model.crypto import paillier, fixed_point_number
from fedlearner.common import tree_model_pb2 as tree_pb2

BST_TYPE = np.float32
PRECISION = 1e38
EXPONENT = math.floor(
    math.log(PRECISION, fixed_point_number.FixedPointNumber.BASE))
KEY_NBITS = 1024
CIPHER_NBYTES = (KEY_NBITS * 2)//8

def _send_public_key(bridge, public_key):
    msg = tree_pb2.EncryptedNumbers()
    msg.ciphertext.append(public_key.n.to_bytes(KEY_NBITS//8, 'little'))
    bridge.send_proto(bridge.current_iter_id, 'public_key', msg)

def _receive_public_key(bridge):
    msg = tree_pb2.EncryptedNumbers()
    bridge.receive_proto(bridge.current_iter_id, 'public_key').Unpack(msg)
    return paillier.PaillierPublicKey(
        int.from_bytes(msg.ciphertext[0], 'little'))

def _encode_encrypted_numbers(numbers):
    return [
        i.ciphertext(False).to_bytes(CIPHER_NBYTES, 'little') \
        for i in numbers]

def _encrypt_numbers(public_key, numbers):
    return _encode_encrypted_numbers(
        [public_key.encrypt(i, PRECISION) for i in numbers])

def _decrypt_number(private_key, numbers):
    return [private_key.decrypt(i) for i in numbers]

def _from_ciphertext(public_key, ciphertext):
    return [
        paillier.PaillierEncryptedNumber(
            public_key, int.from_bytes(i, 'little'), EXPONENT)
        for i in ciphertext]

def _encrypt_and_send_numbers(bridge, name, public_key, numbers):
    msg = tree_pb2.EncryptedNumbers()
    msg.ciphertext.extend(_encrypt_numbers(public_key, numbers))
    bridge.send_proto(bridge.current_iter_id, name, msg)

def _receive_encrypted_numbers(bridge, name, public_key):
    msg = tree_pb2.EncryptedNumbers()
    bridge.receive_proto(bridge.current_iter_id, name).Unpack(msg)
    return _from_ciphertext(public_key, msg.ciphertext)

class BinnedFeatures(object):
    def __init__(self, features, max_bins):
        super(BinnedFeatures, self).__init__()

        self._max_bins = max_bins
        self.features = features
        self.binned, self.thresholds = self._bin_features(features)

    def _bin_features(self, features):
        thresholds = []
        binned = np.zeros_like(features, dtype=np.uint8, order='F')
        for i in range(features.shape[1]):
            x = features[:, i]
            missing_mask = np.isnan(x)
            nonmissing_x = x
            if missing_mask.any():
                nonmissing_x = x[~missing_mask]
            nonmissing_x = np.ascontiguousarray(
                nonmissing_x, dtype=BST_TYPE)
            unique_x = np.unique(nonmissing_x)
            if len(unique_x) <= self._max_bins:
                threshold = (unique_x[:-1] + unique_x[1:]) * 0.5
            else:
                percentiles = np.linspace(0, 100, num=self._max_bins + 1)
                percentiles = percentiles[1:-1]
                threshold = np.percentile(
                    nonmissing_x, percentiles, interpolation='midpoint')
                assert threshold.size == self._max_bins - 1
            thresholds.append(threshold)

            binned[:, i] = np.searchsorted(threshold, x, side='right')
            binned[missing_mask, i] = threshold.size + 1

        return binned, thresholds

def _compute_histogram_helper(args):
    values, binned_features, thresholds, zero = args
    hists = []
    for i, threshold in enumerate(thresholds):
        num_bins = threshold.size + 2
        hist = np.asarray([zero for _ in range(num_bins)])
        np.add.at(hist, binned_features[:, i], values)
        hists.append(hist)

    return hists

class HistogramBuilder(object):
    def __init__(self, binned_features, dtype=BST_TYPE,
                 num_parallel=1, pool=None):
        self._bins = binned_features
        self._dtype = dtype
        self._zero = dtype(0.0)
        self._num_parallel = num_parallel
        self._pool = pool
        if self._num_parallel > 1:
            assert pool is not None
            self._job_size = \
                (len(self._bins.binned) + num_parallel - 1)//num_parallel

    def compute_histogram(self, values, sample_ids):
        if not self._pool:
            return _compute_histogram_helper(
                (values[sample_ids], self._bins.binned[sample_ids],
                 self._bins.thresholds, self._zero))

        args = [
            (values[sample_ids],
             self._bins.binned[
                 sample_ids, self._job_size*i:self._job_size*(i+1)],
             self._bins.thresholds[self._job_size*i:self._job_size*(i+1)],
             self._zero)
            for i in range(self._num_parallel)
        ]

        rets = self._pool.map(_compute_histogram_helper, args)
        return sum(rets, [])

class GrowerNode(object):
    def __init__(self, node_id):
        self.node_id = node_id
        self.feature_id = None
        self.threshold = None
        self.is_owner = None
        self.owner_id = None
        self.weight = None

        self.parent = None
        self.left_child = None
        self.right_child = None

        self.sample_ids = None
        self.grad_hists = None
        self.hess_hists = None

    def to_proto(self):
        return tree_pb2.RegressionTreeNodeProto(
            node_id=self.node_id,
            left_child=self.left_child,
            right_child=self.right_child,
            parent=self.parent,
            is_owner=self.is_owner,
            owner_id=self.owner_id,
            feature_id=self.feature_id,
            threshold=self.threshold,
            weight=self.weight)

class BaseGrower(object):
    def __init__(self, binned, grad, hess, grow_policy='depthwise',
                 max_leaves=None, max_depth=None, learning_rate=0.3,
                 l2_regularization=1.0, dtype=BST_TYPE, num_parallel=1):
        self._binned = binned
        self._grad = grad
        self._hess = hess

        self._grow_policy = grow_policy
        if grow_policy == 'depthwise':
            self._split_candidates = queue.Queue()
            assert max_depth is not None, \
                "max_depth must be set when grow_policy is depthwise"
            self._max_depth = max_depth
            self._max_leaves = 2**max_depth
        else:
            self._split_candidates = queue.PriorityQueue()
            assert max_leaves is not None, \
                "max_leaves must be set when grow_policy is lossguided"
            self._max_leaves = max_leaves
            self._max_depth = max_depth if max_depth is not None else 2**31

        self._learning_rate = learning_rate
        self._l2_regularization = l2_regularization

        self._num_parallel = num_parallel
        self._pool = None
        if self._num_parallel > 1:
            self._pool = mp.Pool(num_parallel)
        self._hist_builder = HistogramBuilder(
            binned, dtype, num_parallel, self._pool)

        self._nodes = [GrowerNode(0)]
        self._nodes[0].sample_ids = list(range(binned.features.shape[0]))
        self._num_leaves = 1

    def _compute_histogram(self, node):
        node.grad_hists = self._hist_builder.compute_histogram(
            self._grad, node.sample_ids)
        node.hess_hists = self._hist_builder.compute_histogram(
            self._hess, node.sample_ids)

    def _compute_histogram_from_sibling(self, node, sibling):
        parent = self._nodes[node.parent]
        node.grad_hists = [
            p - l for p, l in zip(parent.grad_hists, sibling.grad_hists)]
        node.hess_hists = [
            p - l for p, l in zip(parent.hess_hists, sibling.hess_hists)]

    def _find_split_and_push(self, node):
        max_gain = -1
        max_fid = None
        split_point = None
        left_weight = None
        right_weight = None
        lam = self._l2_regularization
        for fid, (grad_hist, hess_hist) in \
                enumerate(zip(node.grad_hists, node.hess_hists)):
            sum_g = sum(grad_hist[:-1])
            sum_h = sum(hess_hist[:-1])
            left_g = 0.0
            left_h = 0.0
            for i in range(len(grad_hist[:-1]) - 1):
                left_g += grad_hist[i]
                left_h += hess_hist[i]
                right_g = sum_g - left_g
                right_h = sum_h - left_h
                gain = left_g*left_g/(left_h + lam) + \
                    right_g*right_g/(right_h + lam) - \
                    sum_g*sum_g/(sum_h + lam)
                if gain > max_gain:
                    max_gain = gain
                    max_fid = fid
                    split_point = i
                    left_weight = - left_g/(left_h + lam)
                    right_weight = - right_g/(right_h + lam)

        split_info = tree_pb2.SplitInfo(
            node_id=node.node_id, gain=max_gain, feature_id=max_fid,
            split_point=split_point,
            left_weight=left_weight * self._learning_rate,
            right_weight=right_weight * self._learning_rate)

        self._split_candidates.put((-max_gain, split_info))

        return max_gain, split_info

    def _add_node(self, parent_id):
        node_id = len(self._nodes)
        node = GrowerNode(node_id)
        node.parent = parent_id
        self._nodes.append(node)
        return node_id

    def _set_node_partition(self, node, split_info):
        node.is_owner = True
        node.feature_id = split_info.feature_id
        node.threshold = self._binned.thresholds[
            split_info.feature_id][split_info.split_point]
        left_child = self._nodes[node.left_child]
        left_child.sample_ids = [
            i for i in node.sample_ids if \
                self._binned.features[i, node.feature_id] < node.threshold]
        right_child = self._nodes[node.right_child]
        right_child.sample_ids = [
            i for i in node.sample_ids if \
                self._binned.features[i, node.feature_id] >= node.threshold]

    def _split_next(self):
        _, split_info = self._split_candidates.get()
        node = self._nodes[split_info.node_id]

        node.left_child = self._add_node(node.node_id)
        left_child = self._nodes[node.left_child]
        left_child.weight = split_info.left_weight

        node.right_child = self._add_node(node.node_id)
        right_child = self._nodes[node.right_child]
        right_child.weight = split_info.right_weight

        self._num_leaves += 1

        self._set_node_partition(node, split_info)

        return left_child, right_child, split_info

    def _log_split(self, left_child, right_child, split_info):
        parent = self._nodes[split_info.node_id]

        logging.info(
            "Split node %d at feature %d for gain=%f. " \
            "Left(w=%f, nsamples=%d), Right(w=%f, nsamples=%d)",
            split_info.node_id, split_info.feature_id, split_info.gain,
            left_child.weight, len(left_child.sample_ids),
            right_child.weight, len(right_child.sample_ids))
        assert len(left_child.sample_ids) + len(right_child.sample_ids) \
            == len(parent.sample_ids)

    def to_proto(self):
        proto = tree_pb2.RegressionTreeProto()
        for node in self._nodes:
            proto.nodes.append(node.to_proto())
        return proto

    def get_prediction(self):
        prediction = np.zeros(self._binned.features.shape[0], dtype=BST_TYPE)
        for node in self._nodes:
            if node.left_child is not None:
                continue
            prediction[node.sample_ids] = node.weight
        return prediction

    def grow(self):
        self._compute_histogram(self._nodes[0])
        self._find_split_and_push(self._nodes[0])

        while self._num_leaves < self._max_leaves:
            left_child, right_child, split_info = self._split_next()
            self._log_split(left_child, right_child, split_info)
            self._compute_histogram(left_child)
            self._find_split_and_push(left_child)
            self._compute_histogram_from_sibling(right_child, left_child)
            self._find_split_and_push(right_child)

def _decrypt_histogram_helper(args):
    public_key, private_key, hists = args
    rets = []
    for hist in hists:
        hist = _from_ciphertext(public_key, hist.ciphertext)
        rets.append(np.asarray(_decrypt_number(private_key, hist)))
    return rets

class LeaderGrower(BaseGrower):
    def __init__(self, bridge, public_key, private_key,
                 binned, grad, hess, **kwargs):
        super(LeaderGrower, self).__init__(
            binned, grad, hess, dtype=np.float32, **kwargs)
        self._bridge = bridge
        self._public_key = public_key
        self._private_key = private_key

    def _receive_and_decrypt_histogram(self, name):
        msg = tree_pb2.Histograms()
        self._bridge.receive_proto(
            self._bridge.current_iter_id, name).Unpack(msg)
        if not self._pool:
            return _decrypt_histogram_helper(
                self._public_key, self._private_key, msg.hists)

        job_size = (len(msg.hists) + self._num_parallel - 1)//self._num_parallel
        args = [
            (self._public_key, self._private_key,
             msg.hists[i*job_size:(i+1)*job_size])
            for i in range(self._num_parallel)
        ]
        hists = self._pool.map(_decrypt_histogram_helper, args)
        return sum(hists, [])

    def _compute_histogram(self, node):
        self._bridge.start(self._bridge.new_iter_id())
        grad_hists = self._hist_builder.compute_histogram(
            self._grad, node.sample_ids)
        hess_hists = self._hist_builder.compute_histogram(
            self._hess, node.sample_ids)
        follower_grad_hists = self._receive_and_decrypt_histogram('grad_hists')
        follower_hess_hists = self._receive_and_decrypt_histogram('hess_hists')
        node.grad_hists = grad_hists + follower_grad_hists
        node.hess_hists = hess_hists + follower_hess_hists
        self._bridge.commit()

    def _split_next(self):
        self._bridge.start(self._bridge.new_iter_id())

        _, split_info = self._split_candidates.get()
        node = self._nodes[split_info.node_id]

        node.left_child = self._add_node(node.node_id)
        left_child = self._nodes[node.left_child]
        left_child.weight = split_info.left_weight

        node.right_child = self._add_node(node.node_id)
        right_child = self._nodes[node.right_child]
        right_child.weight = split_info.right_weight

        self._num_leaves += 1

        if split_info.feature_id < self._binned.features.shape[1]:
            self._set_node_partition(node, split_info)
            self._bridge.send_proto(
                self._bridge.current_iter_id, 'split_info',
                tree_pb2.SplitInfo(
                    node_id=split_info.node_id, feature_id=-1,
                    left_samples=left_child.sample_ids,
                    right_samples=right_child.sample_ids))
        else:
            node.is_owner = False
            fid = split_info.feature_id - self._binned.features.shape[1]
            self._bridge.send_proto(
                self._bridge.current_iter_id, 'split_info',
                tree_pb2.SplitInfo(
                    node_id=split_info.node_id, feature_id=fid,
                    split_point=split_info.split_point))

            follower_split_info = tree_pb2.SplitInfo()
            self._bridge.receive_proto(
                self._bridge.current_iter_id, 'follower_split_info') \
                .Unpack(follower_split_info)
            left_child.sample_ids = list(follower_split_info.left_samples)
            right_child.sample_ids = list(follower_split_info.right_samples)

        self._bridge.commit()
        return left_child, right_child, split_info

class FollowerGrower(BaseGrower):
    def __init__(self, bridge, public_key, binned, grad, hess, **kwargs):
        dtype = lambda x: public_key.encrypt(x, PRECISION)
        super(FollowerGrower, self).__init__(
            binned, grad, hess, dtype=dtype, **kwargs)
        self._bridge = bridge
        self._public_key = public_key

    def _compute_histogram_from_sibling(self, node, sibling):
        pass

    def _find_split_and_push(self, node):
        pass

    def _send_histograms(self, name, hists):
        msg = tree_pb2.Histograms()
        for hist in hists:
            ciphertext = _encode_encrypted_numbers(hist)
            msg.hists.append(
                tree_pb2.EncryptedNumbers(ciphertext=ciphertext))
        self._bridge.send_proto(self._bridge.current_iter_id, name, msg)

    def _compute_histogram(self, node):
        self._bridge.start(self._bridge.new_iter_id())
        grad_hists = self._hist_builder.compute_histogram(
            self._grad, node.sample_ids)
        hess_hists = self._hist_builder.compute_histogram(
            self._hess, node.sample_ids)
        self._send_histograms('grad_hists', grad_hists)
        self._send_histograms('hess_hists', hess_hists)
        self._bridge.commit()

    def _split_next(self):
        self._bridge.start(self._bridge.new_iter_id())

        split_info = tree_pb2.SplitInfo()
        self._bridge.receive_proto(
            self._bridge.current_iter_id, 'split_info') \
            .Unpack(split_info)

        node = self._nodes[split_info.node_id]

        node.left_child = self._add_node(node.node_id)
        left_child = self._nodes[node.left_child]
        left_child.weight = float('nan')

        node.right_child = self._add_node(node.node_id)
        right_child = self._nodes[node.right_child]
        right_child.weight = float('nan')

        self._num_leaves += 1

        if split_info.feature_id >= 0:
            self._set_node_partition(node, split_info)
            self._bridge.send_proto(
                self._bridge.current_iter_id, 'follower_split_info',
                tree_pb2.SplitInfo(
                    left_samples=left_child.sample_ids,
                    right_samples=right_child.sample_ids))
        else:
            node.is_owner = False
            left_child.sample_ids = list(split_info.left_samples)
            right_child.sample_ids = list(split_info.right_samples)

        self._bridge.commit()
        return left_child, right_child, split_info

class BoostingTreeEnsamble(object):
    def __init__(self, bridge, learning_rate=0.3, max_iters=50, max_depth=6,
                 max_leaves=None, l2_regularization=1.0, max_bins=33,
                 grow_policy='depthwise', num_parallel=1):
        self._learning_rate = learning_rate
        self._max_iters = max_iters
        self._max_depth = max_depth
        self._max_leaves = max_leaves
        self._l2_regularization = l2_regularization
        self._grow_policy = grow_policy
        self._num_parallel = num_parallel

        assert max_bins < 255, "Only support max_bins < 255"
        self._max_bins = max_bins

        self._loss = LogisticLoss()
        self._trees = []

        self._bridge = bridge
        if bridge is not None:
            self._role = self._bridge.role
            self._bridge.connect()
            self._make_key_pair()

    def _make_key_pair(self):
        # make key pair
        self._bridge.start(self._bridge.new_iter_id())
        if self._role == 'leader':
            self._public_key, self._private_key = \
                paillier.PaillierKeypair.generate_keypair(KEY_NBITS)
            _send_public_key(self._bridge, self._public_key)
        else:
            self._public_key = _receive_public_key(self._bridge)
            self._private_key = None
        self._bridge.commit()

    def save_model(self, path):
        with open(path, 'w') as fout:
            model = tree_pb2.BoostingTreeEnsambleProto()
            model.trees.extend(self._trees)
            fout.write(text_format.MessageToString(model))

    def load_saved_model(self, path):
        with open(path, 'r') as fin:
            model = tree_pb2.BoostingTreeEnsambleProto()
            text_format.Parse(fin.read(), model)
            self._trees = list(model.trees)

    def batch_predict(self, features, raw_score=False):
        if self._bridge is None:
            return self._batch_predict_local(features, raw_score)
        if self._bridge.role == 'leader':
            return self._batch_predict_leader(features, raw_score)
        return self._batch_predict_follower(features, raw_score)

    def _batch_predict_local(self, features, raw_score):
        N = features.shape[0]
        raw_prediction = []
        for i in range(features.shape[0]):
            score = 0.0
            for tree in self._trees:
                node = tree.nodes[0]
                while node.left_child != 0:
                    assert node.is_owner
                    if features[i, node.feature_id] < node.threshold:
                        node = tree.nodes[node.left_child]
                    else:
                        node = tree.nodes[node.right_child]
                score += node.weight
            raw_prediction.append(score)
        raw_prediction = np.asarray(raw_prediction)
        if raw_score:
            return raw_prediction
        return self._loss.predict(raw_prediction)

    def _batch_predict_leader(self, features, raw_score):
        N = features.shape[0]
        raw_prediction = np.zeros(N, dtype=BST_TYPE)
        for tree in self._trees:
            assignment = np.zeros(N, dtype=np.int32)
            finish_count = 0
            while finish_count != N:
                self._bridge.start(self._bridge.new_iter_id())
                finish_count = 0
                for i in range(N):
                    node = tree.nodes[assignment[i]]
                    if node.left_child == 0:
                        finish_count += 1
                        continue
                    if node.is_owner:
                        if features[i, node.feature_id] < node.threshold:
                            assignment[i] = node.left_child
                        else:
                            assignment[i] = node.right_child
                    else:
                        assignment[i] = -1

                self._bridge.send(
                    self._bridge.current_iter_id, 'leader_assignment',
                    assignment)
                follower_assignment = self._bridge.receive(
                    self._bridge.current_iter_id, 'follower_assignment')
                assignment = np.maximum(assignment, follower_assignment)
                self._bridge.commit()
            for i in range(N):
                raw_prediction[i] += tree.nodes[assignment[i]].weight

        if raw_score:
            return raw_prediction
        return self._loss.predict(raw_prediction)

    def _batch_predict_follower(self, features, raw_score):
        N = features.shape[0]
        for tree in self._trees:
            assignment = np.zeros(N, dtype=np.int32)
            finish_count = 0
            while finish_count != N:
                self._bridge.start(self._bridge.new_iter_id())
                finish_count = 0
                for i in range(N):
                    node = tree.nodes[assignment[i]]
                    if node.left_child == 0:
                        finish_count += 1
                        continue
                    if node.is_owner:
                        if features[i, node.feature_id] < node.threshold:
                            assignment[i] = node.left_child
                        else:
                            assignment[i] = node.right_child
                    else:
                        assignment[i] = -1

                self._bridge.send(
                    self._bridge.current_iter_id, 'follower_assignment',
                    assignment)
                leader_assignment = self._bridge.receive(
                    self._bridge.current_iter_id, 'leader_assignment')
                assignment = np.maximum(assignment, leader_assignment)
                self._bridge.commit()

    def fit(self, features, labels=None, checkpoint_path=None):
        num_examples = features.shape[0]

        # sort feature columns
        binned = BinnedFeatures(features, self._max_bins)

        # initial f(x)
        if len(self._trees) > 0:
            sum_prediction = self.batch_predict(features)
        else:
            sum_prediction = np.zeros(num_examples, dtype=BST_TYPE)

        # start iterations
        while len(self._trees) < self._max_iters:
            if self._bridge is None:
                tree, raw_prediction = self._fit_one_round_local(
                    sum_prediction, binned, labels)
                sum_prediction += raw_prediction
            elif self._bridge.role == 'leader':
                tree, raw_prediction = self._fit_one_round_leader(
                    sum_prediction, binned, labels)
                sum_prediction += raw_prediction
            else:
                tree = self._fit_one_round_follower(binned)

            self._trees.append(tree)

            if checkpoint_path is not None:
                filename = os.path.join(
                    checkpoint_path, 'checkpoint-%04d.proto')
                self.save_model(filename)

    def _fit_one_round_local(self, sum_fx, binned, labels):
        # compute grad and hess
        pred = self._loss.predict(sum_fx)
        grad = self._loss.gradient(sum_fx, pred, labels)
        hess = self._loss.hessian(sum_fx, pred, labels)
        logging.info(
            'Leader starting iteration %d. Metrics are %s',
            len(self._trees), self._loss.metrics(pred, labels))

        grower = BaseGrower(
            binned, grad, hess,
            learning_rate=self._learning_rate,
            max_depth=self._max_depth,
            max_leaves=self._max_leaves,
            l2_regularization=self._l2_regularization,
            grow_policy=self._grow_policy,
            num_parallel=self._num_parallel)
        grower.grow()

        return grower.to_proto(), grower.get_prediction()

    def _fit_one_round_leader(self, sum_fx, binned, labels):
        # compute grad and hess
        self._bridge.start(self._bridge.new_iter_id())
        pred = self._loss.predict(sum_fx)
        grad = self._loss.gradient(sum_fx, pred, labels)
        hess = self._loss.hessian(sum_fx, pred, labels)
        print('metrics: %s'%self._loss.metrics(pred, labels))
        _encrypt_and_send_numbers(
            self._bridge, 'grad', self._public_key, grad)
        _encrypt_and_send_numbers(
            self._bridge, 'hess', self._public_key, hess)
        self._bridge.commit()

        grower = LeaderGrower(
            self._bridge, self._public_key, self._private_key,
            binned, grad, hess,
            learning_rate=self._learning_rate,
            max_depth=self._max_depth,
            max_leaves=self._max_leaves,
            l2_regularization=self._l2_regularization,
            grow_policy=self._grow_policy,
            num_parallel=self._num_parallel)
        grower.grow()

        return grower.to_proto(), grower.get_prediction()

    def _fit_one_round_follower(self, binned):
        # compute grad and hess
        self._bridge.start(self._bridge.new_iter_id())
        grad = np.asarray(_receive_encrypted_numbers(
            self._bridge, 'grad', self._public_key))
        hess = np.asarray(_receive_encrypted_numbers(
            self._bridge, 'hess', self._public_key))
        self._bridge.commit()
        logging.info(
            'Follower starting iteration %d.',
            len(self._trees))

        grower = FollowerGrower(
            self._bridge, self._public_key,
            binned, grad, hess,
            learning_rate=self._learning_rate,
            max_depth=self._max_depth,
            max_leaves=self._max_leaves,
            l2_regularization=self._l2_regularization,
            grow_policy=self._grow_policy,
            num_parallel=self._num_parallel)
        grower.grow()

        return grower.to_proto()


# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module contains functions for adding the Autograd interface
to a PennyLane Device class.
"""
# pylint: disable=too-many-arguments
import autograd
from autograd.numpy.numpy_boxes import ArrayBox

import pennylane as qml
from pennylane import numpy as np

def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=2):
    """Execute a batch of tapes with Autograd parameters on a device.

    Args:
        tapes (Sequence[.QuantumTape]): batch of tapes to execute
        device (.Device): Device to use to execute the batch of tapes.
            If the device does not provide a ``batch_execute`` method,
            by default the tapes will be executed in serial.
        execute_fn (callable): The execution function used to execute the tapes
            during the forward pass. This function must return a tuple ``(results, jacobians)``.
            If ``jacobians`` is an empty list, then ``gradient_fn`` is used to
            compute the gradients during the backwards pass.
        gradient_kwargs (dict): dictionary of keyword arguments to pass when
            determining the gradients of tapes
        gradient_fn (callable): the gradient function to use to compute quantum gradients
        _n (int): a positive integer used to track nesting of derivatives, for example
            if the nth-order derivative is requested.
        max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies
            the maximum order of derivatives to support. Increasing this value allows
            for higher order derivatives to be extracted, at the cost of additional
            (classical) computational overhead during the backwards pass.

    Returns:
        list[list[float]]: A nested list of tape results. Each element in
        the returned list corresponds in order to the provided tapes.
    """
    for tape in tapes:
        # set the trainable parameters
        params = tape.get_parameters(trainable_only=False)
        tape.trainable_params = qml.math.get_trainable_indices(params)

    parameters = autograd.builtins.tuple(
        [autograd.builtins.list(t.get_parameters()) for t in tapes]
    )

    return _execute(
        parameters,
        tapes=tapes,
        device=device,
        execute_fn=execute_fn,
        gradient_fn=gradient_fn,
        gradient_kwargs=gradient_kwargs,
        _n=_n,
        max_diff=max_diff,
    )[0]

@autograd.extend.primitive
def _execute(
    parameters,
    tapes=None,
    device=None,
    execute_fn=None,
    gradient_fn=None,
    gradient_kwargs=None,
    _n=1,
    max_diff=2,
):  # pylint: disable=dangerous-default-value,unused-argument
    """Autodifferentiable wrapper around ``Device.batch_execute``.

    The signature of this function is designed to work around Autograd restrictions.
    Note that the ``parameters`` argument is dependent on the ``tapes`` argument;
    this function should always be called as follows:

    >>> parameters = [autograd.builtins.list(t.get_parameters()) for t in tapes])
    >>> parameters = autograd.builtins.tuple(parameters)
    >>> _execute(parameters, tapes=tapes, device=device)

    In particular:

    - ``parameters`` is dependent on the provided tapes: always extract them as above
    - ``tapes`` is a *required* argument
    - ``device`` is a *required* argument

    The private argument ``_n`` is used to track nesting of derivatives, for example
    if the nth-order derivative is requested. Do not set this argument unless you
    understand the consequences!
    """
    with qml.tape.Unwrap(*tapes):
        res, jacs = execute_fn(tapes, **gradient_kwargs)

    for i, r in enumerate(res):
        res[i] = np.tensor(r)

        if res[i].dtype == np.dtype("object"):
            # For backwards compatibility, we flatten ragged tape outputs
            res[i] = np.hstack(r)

    return res, jacs

def vjp(
    ans,
    parameters,
    tapes=None,
    device=None,
    execute_fn=None,
    gradient_fn=None,
    gradient_kwargs=None,
    _n=1,
    max_diff=2,
):  # pylint: disable=dangerous-default-value,unused-argument
    """Returns the vector-Jacobian product operator for a batch of quantum tapes.

    Args:
        ans (array): the result of the batch tape execution
        parameters (list[list[Any]]): Nested list of the quantum tape parameters.
            This argument should be generated from the provided list of tapes.
        tapes (Sequence[.QuantumTape]): batch of tapes to execute
        device (.Device): Device to use to execute the batch of tapes.
            If the device does not provide a ``batch_execute`` method,
            by default the tapes will be executed in serial.
        execute_fn (callable): The execution function used to execute the tapes
            during the forward pass. This function must return a tuple ``(results, jacobians)``.
            If ``jacobians`` is an empty list, then ``gradient_fn`` is used to
            compute the gradients during the backwards pass.
        gradient_fn (callable): the gradient function to use to compute quantum gradients
        gradient_kwargs (dict): dictionary of keyword arguments to pass when
            determining the gradients of tapes
        _n (int): a positive integer used to track nesting of derivatives, for example
            if the nth-order derivative is requested.
        max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies
            the maximum number of derivatives to support. Increasing this value allows
            for higher order derivatives to be extracted, at the cost of additional
            (classical) computational overhead during the backwards pass.

    Returns:
        function: this function accepts the backpropagation
        gradient output vector, and computes the vector-Jacobian product
    """

    def grad_fn(dy):
        """Returns the vector-Jacobian product with given
        parameter values and output gradient dy"""

        dy = dy[0]
        jacs = ans[1]

        if jacs:
            # Jacobians were computed on the forward pass (mode="forward")
            # No additional quantum evaluations needed; simply compute the VJPs directly.
            vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(dy, jacs)]

        else:
            # Need to compute the Jacobians on the backward pass (accumulation="backward")

            if isinstance(gradient_fn, qml.gradients.gradient_transform):
                # Gradient function is a gradient transform.

                # Generate and execute the required gradient tapes
                if _n == max_diff:
                    with qml.tape.Unwrap(*tapes):
                        vjp_tapes, processing_fn = qml.gradients.batch_vjp(
                            tapes,
                            dy,
                            gradient_fn,
                            reduction="append",
                            gradient_kwargs=gradient_kwargs,
                        )

                        vjps = processing_fn(execute_fn(vjp_tapes)[0])

                else:
                    vjp_tapes, processing_fn = qml.gradients.batch_vjp(
                        tapes, dy, gradient_fn, reduction="append", gradient_kwargs=gradient_kwargs
                    )

                    # This is where the magic happens. Note that we call ``execute``.
                    # This recursion, coupled with the fact that the gradient transforms
                    # are differentiable, allows for arbitrary order differentiation.
                    vjps = processing_fn(
                        execute(
                            vjp_tapes,
                            device,
                            execute_fn,
                            gradient_fn,
                            gradient_kwargs,
                            _n=_n + 1,
                            max_diff=max_diff,
                        )
                    )

            else:
                # Gradient function is not a gradient transform
                # (e.g., it might be a device method).
                # Note that unlike the previous branch:
                #
                # - there is no recursion here
                # - gradient_fn is not differentiable
                #
                # so we cannot support higher-order derivatives.
                with qml.tape.Unwrap(*tapes):
                    jacs = gradient_fn(tapes, **gradient_kwargs)

                vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(dy, jacs)]

        return [qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps]

    return grad_fn

autograd.extend.defvjp(_execute, vjp, argnums=[0])


# -*- coding: utf-8 -*-
"""
Generates circuits for quantum error correction with surface code patches.
"""
import copy
import warnings
from abc import ABC, abstractmethod

import numpy as np
import networkx as nx
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute

try:
    from qiskit import Aer

    HAS_AER = True
except ImportError:
    from qiskit import BasicAer

    HAS_AER = False

class LatticeError(Exception):
    pass

class _Face:
    """
    Abstract class for a single *face* of the surface code, which is described
    by a syndrome qubit and the four data qubits that surround it. If the face
    exists on an edge of the surface code lattice, some of the data qubits may
    be None.
    """

    def __init__(self, syndrome, top_l, top_r, bot_l, bot_r):
        """
        Initialize a face by passing in the single QubitRegisters of the circuit
        in which this face will be embedded.
        """
        self.syndrome = syndrome
        self.top_l = top_l
        self.top_r = top_r
        self.bot_l = bot_l
        self.bot_r = bot_r

    @abstractmethod
    def entangle(self, circ):
        pass

class _FaceX(_Face):
    """
    X-syndrome face of the rotated surface code.
    """

    def entangle(self, circ):
        """
        Traverse in reverse "Z" pattern
        """
        if (self.top_r and not self.top_l) or (self.bot_r and not self.bot_l):
            raise LatticeError("Inconsistent X syndrome connections")

        circ.h(self.syndrome)
        if self.top_r:
            circ.cx(self.syndrome, self.top_r)
            circ.cx(self.syndrome, self.top_l)
        if self.bot_r:
            circ.cx(self.syndrome, self.bot_r)
            circ.cx(self.syndrome, self.bot_l)
        circ.h(self.syndrome)

class _FaceZ(_Face):
    """
    Z-syndrome face of the rotated surface code.
    """

    def entangle(self, circ):
        """
        Traverse in reverse "N" pattern
        """
        if (self.top_r and not self.bot_r) or (self.top_l and not self.bot_l):
            raise LatticeError("Inconsistent Z syndrome connections")

        if self.top_r:
            circ.cx(self.top_r, self.syndrome)
            circ.cx(self.bot_r, self.syndrome)
        if self.top_l:
            circ.cx(self.top_l, self.syndrome)
            circ.cx(self.bot_l, self.syndrome)

class RotatedSurfaceCodeLattice:
    """
    This class is essentially a helper class that translates between the lattice
    abstraction of the surface code and the physical qubits in the circuit. This
    promotes a clean separation of concerns and wraps the implementation details
    of the lattice + code entanglement ordering.
    """

    def __init__(self, d, data_register, mx_register, mz_register):
        """
        Initializes an instance of the rotated surface code lattice with our
        chosen layout and numbering.

        Args:
            d (int): surface code distance
            data_register (QuantumRegister): grouped register of all data qubits
            mx_register (QuantumRegister): grouped register of all measure-x qubits
            mz_register (QuantumRegister): grouped register of all measure-z qubits
        """
        self.d = d
        self.measure_x = []
        self.measure_z = []

        per_row_x = (d ** 2 - 1) // 2 // (d + 1)
        per_row_z = (d ** 2 - 1) // 2 // (d - 1)
        for mx in mx_register:
            idx = mx.index
            row = idx // per_row_x
            offset = idx % per_row_x
            start = (row - 1) * d
            row_parity = row % 2

            if row == 0:  # First row
                top_l, top_r = None, None
                bot_l = data_register[idx * 2]
                bot_r = data_register[idx * 2 + 1]
            elif row == d:  # Last row
                bot_l, bot_r = None, None
                top_l = data_register[idx * 2 + 1]
                top_r = data_register[idx * 2 + 2]
            else:
                top_l = data_register[start + (offset * 2) + row_parity]
                top_r = data_register[start + (offset * 2) + row_parity + 1]
                bot_l = data_register[start + d + (offset * 2) + row_parity]
                bot_r = data_register[start + d + (offset * 2) + row_parity + 1]
            self.measure_x.append(_FaceX(mx, top_l, top_r, bot_l, bot_r))

        for mz in mz_register:
            idx = mz.index
            row = idx // per_row_z
            offset = idx % per_row_z
            start = row * d
            row_parity = row % 2

            top_l = data_register[start + (offset * 2) - row_parity]
            top_r = data_register[start + (offset * 2) - row_parity + 1]
            bot_l = data_register[start + d + (offset * 2) - row_parity]
            bot_r = data_register[start + d + (offset * 2) - row_parity + 1]

            # Overwrite edge column syndromes
            if row_parity == 0 and offset == per_row_z:  # Last column
                top_r, bot_r = None, None
            elif row_parity == 1 and offset == 0:  # First column
                top_l, bot_l = None, None

            self.measure_z.append(_FaceZ(mz, top_l, top_r, bot_l, bot_r))

    def entangle(self, circ):
        """
        Entangles the entire surface code by calling the entangle method of each
        syndrome face. Within a face, order is determined by the delegated
        method. Order between faces should not matter here, since the projection
        will be determined by the measurement order.
        """
        for syndrome in (self.measure_x, self.measure_z):
            for face in syndrome:
                face.entangle(circ)

    def parse_readout(self, readout_string):
        """
        Helper method to turn a result string (e.g. 1 10100000 10010000) into an
        appropriate logical readout value and XOR-ed syndrome locations
        according to our grid coordinate convention.
        """
        syn_len = (self.d ** 2 - 1) // 2
        chunks = readout_string.split(" ")

        int_syndromes = [int(x, base=2) for x in chunks[-1:0:-1]]
        xor_syndromes = [a ^ b for (a, b) in zip(int_syndromes, int_syndromes[1:])]

        mask_Z = "1" * syn_len
        mask_X = mask_Z + "0" * syn_len
        X_syndromes = [(x & int(mask_X, base=2)) >> syn_len for x in xor_syndromes]
        Z_syndromes = [x & int(mask_Z, base=2) for x in xor_syndromes]

        X = []
        for T, syndrome in enumerate(X_syndromes):
            for loc in range(syn_len):
                if syndrome & 1 << loc:
                    X.append((T, -0.5 + loc, 0.5 + loc % 2))

        Z = []
        for T, syndrome in enumerate(Z_syndromes):
            for loc in range(syn_len):
                if syndrome & 1 << loc:
                    Z.append((T, 0.5 + loc // 2, 0.5 + loc % 2 * 2 - loc // 2))

        return (
            int(chunks[0]),
            {"X": X, "Z": Z,},
        )

class SurfaceCodeLogicalQubit(QuantumCircuit):
    """
    A single logical surface code qubit. At the physical level, this wraps a
    circuit, so we chose to subclass and extend QuantumCircuit.
    """

    def __init__(self, d, *args, **kwargs):
        """
        Initializes a new QuantumCircuit for this logical qubit and calculates
        the underlying surface code lattice ordering.

        Args:
            d (int): Number of physical "data" qubits. Only odd d is possible!
        """
        if d % 2 != 1:
            raise ArgumentError("Surface code distance must be odd!")
        self.__d = d
        self.__T = 0
        self.__num_data = d ** 2
        self.__num_syn = (d ** 2 - 1) // 2

        self.__data = QuantumRegister(self.__num_data, "data")
        self.__mz = QuantumRegister(self.__num_syn, "mz")
        self.__mx = QuantumRegister(self.__num_syn, "mx")

        # We implement and assume the rotated lattice only, but this can be
        # imagined to accept other layouts in the future.
        self.__lattice = RotatedSurfaceCodeLattice(d, self.__data, self.__mx, self.__mz)

        # Spare ancilla (e.g. for readout)
        self.__ancilla = QuantumRegister(1, name="ancilla")
        super().__init__(self.__data, self.__mz, self.__mx, self.__ancilla)

    def stabilize(self):
        """
        Run a single round of stabilization (entangle and measure).
        """
        syndrome_readouts = ClassicalRegister(
            self.__num_syn * 2, name="c{}".format(self.__T)
        )
        self.add_register(syndrome_readouts)
        self.__T += 1

        self.__lattice.entangle(self)
        self.barrier()
        self.measure(self.__mz, syndrome_readouts[0 : self.__num_syn])
        self.measure(self.__mx, syndrome_readouts[self.__num_syn : self.__num_syn * 2])
        self.reset(self.__mz)
        self.reset(self.__mx)
        self.barrier()

    def identity(self):
        """
        Inserts an identity on the data and syndrome qubits. This is a hack to
        create an isolated error model.
        """
        [
            self.id(x)
            for register in (self.__data, self.__mz, self.__mx)
            for x in register
        ]
        self.barrier()

    def hadamard_reset(self):
        """
        A hack to initialize a + and - logical qubit for now...
        """
        [self.reset(x) for x in self.__data]
        [self.h(x) for x in self.__data]
        self.barrier()

    def logical_x(self):
        """
        Logical X operator on the qubit.
        """
        for i in range(0, self.__num_data, self.__d):
            self.x(self.__data[i])
        self.barrier()

    def logical_z(self):
        """
        Logical Z operator on the qubit.
        """
        for i in range(self.__d):
            self.z(self.__data[i])
        self.barrier()

    def readout_z(self):
        """
        Convenience method to read-out the logical-Z projection.
        """
        readout = ClassicalRegister(1, name="readout")
        self.add_register(readout)

        self.reset(self.__ancilla)
        for i in range(self.__d):
            self.cx(self.__data[i], self.__ancilla)
        self.measure(self.__ancilla, readout)
        self.barrier()

    def readout_x(self):
        """
        Convenience method to read-out the logical-X projection.
        """
        readout = ClassicalRegister(1, name="readout")
        self.add_register(readout)

        self.reset(self.__ancilla)
        self.h(self.__ancilla)
        for i in range(0, self.__num_data, self.__d):
            self.cx(self.__ancilla, self.__data[i])
        self.h(self.__ancilla)
        self.measure(self.__ancilla, readout)
        self.barrier()

    def parse_readout(self, readout_string):
        return self.__lattice.parse_readout(readout_string)


import os,os.path,shutil
import numpy as np
import pickle

import dmdd
import dmdd_efficiencies as eff

def check_min_mass(element='fluorine', Qmin=1., v_esc=544., v_lag=220., mx_guess=1.):
    experiment = dmdd.Experiment('test',element,Qmin, 40.,100., eff.efficiency_unit)
    res = experiment.find_min_mass(v_esc=v_esc, v_lag=v_lag, mx_guess=mx_guess)
    print res,'GeV'
    if res<0:
        print 'Problem: try another mx_guess...'

def make_UVmodels(return_models=False):
    SI_Higgs = dmdd.UV_Model('SI_Higgs', ['mass', 'sigma_si'], fixed_params={'fnfp_si': 1})
    milicharge = dmdd.UV_Model('Milicharge', ['mass', 'sigma_si_massless'], fixed_params={'fnfp_si_massless': 0})
    SD_flavoruniversal = dmdd.UV_Model('SD_fu', ['mass','sigma_sd'], fixed_params={'fnfp_sd': -1.1})
    anapole = dmdd.UV_Model('Anapole', ['mass','sigma_anapole'])
    magdip_heavy = dmdd.UV_Model('Mag.dip.heavy', ['mass','sigma_magdip'])
    magdip_0 = dmdd.UV_Model('Mag.dip.light', ['mass','sigma_magdip_massless'])
    elecdip_heavy = dmdd.UV_Model('Elec.dip.heavy', ['mass','sigma_elecdip'])
    elecdip_0 = dmdd.UV_Model('Elec.dip.light', ['mass','sigma_elecdip_massless'])
    f1 = dmdd.UV_Model('f1', ['mass','sigma_f1'], fixed_params={'fnfp_f1': 1.})
    f2_Higgs = dmdd.UV_Model('f2_Higgs', ['mass','sigma_f2'], fixed_params={'fnfp_f2': -0.05})
    #f2_flavoruniversal = dmdd.UV_Model('f2_flavor-universal', ['mass','sigma_f2'], fixed_params={'fnfp_f2': 1.})
    f3_Higgs = dmdd.UV_Model('f3_Higgs', ['mass','sigma_f3'], fixed_params={'fnfp_f3': -0.05})
    #f3_flavoruniversal = dmdd.UV_Model('f3_flavor-universal', ['mass','sigma_f3'], fixed_params={'fnfp_f3': 1.})
    LS = dmdd.UV_Model('LS', ['mass','sigma_LS'], fixed_params={'fnfp_LS': 0.})

    models = [SI_Higgs, milicharge, SD_flavoruniversal, anapole,
              magdip_heavy, magdip_0, elecdip_heavy, elecdip_0,
              f1, f2_Higgs, f3_Higgs, LS]

    if return_models:
        return models

def make_experiments(return_experiments=False):
    xe = dmdd.Experiment('Xe','xenon',5., 40.,100., eff.efficiency_unit)
    ge = dmdd.Experiment('Ge','germanium',0.4, 100.,100., eff.efficiency_unit)
    if return_experiments:
        return [xe,ge]

def test_MultinestRun(mass=50,test_fits=False):

    SI_Higgs = dmdd.UV_Model('SI_Higgs', ['mass', 'sigma_si'], fixed_params={'fnfp_si': 1})
    elecdip_heavy = dmdd.UV_Model('Elec.dip.heavy', ['mass','sigma_elecdip'])

    experiment = make_experiments(return_experiments=True)

    simmodel = SI_Higgs
    fitmodel1 = SI_Higgs
    fitmodel2 = elecdip_heavy

    pardic = {'sigma_si': 70.,'mass': mass}
    simname = 'simtest'

    testrun1 = dmdd.MultinestRun(simname, experiment, simmodel, pardic,
                                 fitmodel1, prior_ranges={'mass':(1,1000),
                                                          'sigma_si':(0.001,100000),
                                                          'sigma_elecdip':(0.001,100000)})
    data1 = np.loadtxt(testrun1.simulations[0].datafile)

    pardic = {'sigma_si': 70.0007,'mass': mass}
    testrun2 = dmdd.MultinestRun(simname, experiment, simmodel, pardic,
                                 fitmodel2, empty_run=False,
                                 prior_ranges={'mass':(1,1000),
                                                'sigma_si':(0.001,100000),
                                                'sigma_elecdip':(0.001,100000)})
    data2 = np.loadtxt(testrun1.simulations[0].datafile)

    #simulation datafile should be created only for the first instance of MultinestRun:
    assert np.allclose(data1, data2)

    if test_fits:

        testrun1.fit()
        testrun1.visualize()
        testrun2.fit()
        testrun2.visualize()

        if (not os.path.exists(testrun1.chains_file)) or (not os.path.exists(testrun1.pickle_file)) or (not os.path.exists(testrun1.stats_file)):
            raise AssertionError('Stats or chains or pickle are not created or are erased.')

        plotfile1 = testrun1.chainspath + '2d_posterior_mass_vs_sigma_si.pdf'
        plotfile2 = testrun1.chainspath + '{}_theoryfitdata_Ge.pdf'.format(simname)
        plotfile3 = testrun1.chainspath + '{}_theoryfitdata_Xe.pdf'.format(simname)
        if (not os.path.exists(plotfile1)) or (not os.path.exists(plotfile2)) or (not os.path.exists(plotfile3)):
            raise AssertionError('Plots are not created or are erased.')

        if (not os.path.exists(testrun2.chains_file)) or (not os.path.exists(testrun2.pickle_file)) or (not os.path.exists(testrun2.stats_file)):
            raise AssertionError('Stats or chains or pickle are not created.')

        plotfile1 = testrun2.chainspath + '2d_posterior_mass_vs_sigma_elecdip.pdf'
        plotfile2 = testrun2.chainspath + '{}_theoryfitdata_Ge.pdf'.format(simname)
        plotfile3 = testrun2.chainspath + '{}_theoryfitdata_Xe.pdf'.format(simname)
        if (not os.path.exists(plotfile1)) or (not os.path.exists(plotfile2)) or (not os.path.exists(plotfile3)):
            raise AssertionError('Plots are not created.')

def test_UVrate():

    experiment = dmdd.Experiment('Xe','xenon',5., 40.,10000., eff.efficiency_unit)
    models = make_UVmodels(return_models=True)
    mass = 40.
    qs = np.array([15.])
    v_lag = 200.
    v_rms = 100.
    v_esc = 600.
    rho_x = 0.4

    sigma_names = {}
    fnfp_names = {}
    fnfp_vals = {}
    for m in models:
        sigma_names[m.name] = m.param_names[1]
        if len(m.fixed_params)>0:
            fnfp_names[m.name] = m.fixed_params.keys()[0]
            fnfp_vals[m.name] = m.fixed_params.values()[0]
        else:
            fnfp_names[m.name] = None
            fnfp_vals[m.name] = None

    dRdQs = np.zeros(len(models))
    Rs = np.zeros(len(models))
    for i,m in enumerate(models):
        kwargs = {sigma_names[m.name]:1.}
        if fnfp_names[m.name] is not None:
            kwargs[fnfp_names[m.name]] = fnfp_vals[m.name]

        dRdQs[i] = dmdd.rate_UV.dRdQ(qs, mass=mass, element=experiment.element,
                                        v_lag=v_lag, v_rms=v_rms, v_esc=v_esc, rho_x=rho_x,
                                        **kwargs)
        Rs[i] = dmdd.rate_UV.R(eff.efficiency_unit, mass=mass, element=experiment.element,
                                        Qmin=experiment.Qmin, Qmax=experiment.Qmax,
                                        v_lag=v_lag, v_rms=v_rms, v_esc=v_esc, rho_x=rho_x,
                                        **kwargs)
    #print 'dRdQs = {}\n'.format(dRdQs)
    #print 'Rs = {}\n'.format(Rs)
    dRdQs_correct = [  1.27974652e-12,   1.67031585e-13,   6.28936205e-13,   7.76864477e-13,
                       7.71724584e-13,   5.66164037e-13,   8.40579288e-13,   6.16678247e-13,
                       4.72480605e-13,   2.59857470e-16,   9.59390104e-16,   1.14295679e-13]

    Rs_correct = [  6.15358778e-11,   3.10857259e-11,   3.14982315e-11,   4.14119198e-11,
                    1.82181891e-11,   3.84877268e-11,   2.35638282e-11,   5.50063883e-11,
                    1.34702925e-11,   5.82472177e-15,   1.64213483e-14,   2.26028126e-12]

    assert np.allclose(dRdQs_correct, dRdQs)
    assert np.allclose(Rs_correct, Rs)

    ###
    qs = np.array([8.3,15.7])
    logtest1 = dmdd.rate_UV.loglikelihood(qs, eff.efficiency_unit, mass=mass,
                                              sigma_si=1.,fnfp_si=1.,
                                                element=experiment.element,
                                                Qmin=experiment.Qmin, Qmax=experiment.Qmax,
                                                exposure=experiment.exposure,energy_resolution=True,
                                                v_lag=v_lag, v_rms=v_rms, v_esc=v_esc, rho_x=rho_x)
    logtest2 = dmdd.rate_UV.loglikelihood(qs, eff.efficiency_unit, mass=mass,
                                              sigma_si=1.,fnfp_si=1.,
                                                element=experiment.element,
                                                Qmin=experiment.Qmin, Qmax=experiment.Qmax,
                                                exposure=experiment.exposure,energy_resolution=False,
                                                v_lag=v_lag, v_rms=v_rms, v_esc=v_esc, rho_x=rho_x)

    #print 'logtest1_correct={}'.format(logtest1)
    #print 'logtest2_correct={}'.format(logtest2)
    logtest1_correct=-19.6890210901
    logtest2_correct=-13.4627188661

    print('correct={}  got={}\n'.format(logtest1_correct,logtest1))
    print('correct={}  got={}\n'.format(logtest2_correct,logtest2))

    assert np.isclose(logtest1_correct, logtest1)
    assert np.isclose(logtest2_correct, logtest2)