# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function
import glob
import re
import numpy as np
from scipy.interpolate import RegularGridInterpolator
from scipy.interpolate import UnivariateSpline
import healpy as hp
from astropy.io import fits

import pyIrfLoader

pyIrfLoader.Loader_go()

from fermipy import utils
from fermipy import spectrum
from fermipy.utils import edge_to_center
from fermipy.utils import edge_to_width
from fermipy.utils import sum_bins
from fermipy.skymap import HpxMap
from fermipy.hpx_utils import HPX
from fermipy.ltcube import LTCube

evtype_string = {
    1: 'FRONT',
    2: 'BACK',
    4: 'PSF0',
    8: 'PSF1',
    16: 'PSF2',
    32: 'PSF3',
    64: 'EDISP0',
    128: 'EDISP1',
    256: 'EDISP2',
    512: 'EDISP3',
}

def loglog_quad(x, y, dim):

    ys0 = [slice(None)] * y.ndim
    ys1 = [slice(None)] * y.ndim

    xs0 = [None] * y.ndim
    xs1 = [None] * y.ndim

    ys0[dim] = slice(None, -1)
    ys1[dim] = slice(1, None)

    xs0[dim] = slice(None, -1)
    xs1[dim] = slice(1, None)
    log_ratio = np.log(x[xs1] / x[xs0])
    return 0.5 * (y[ys0] * x[xs0] + y[ys1] * x[xs1]) * log_ratio

def bins_per_dec(edges):
    return (len(edges) - 1) / np.log10(edges[-1] / edges[0])

def bitmask_to_bits(mask):

    bits = []
    for i in range(32):
        if mask & (2**i):
            bits += [2**i]

    return bits

def poisson_log_like(c, m):
    return c * np.log(m) - m

def poisson_ts(sig, bkg, bkg_fit=None):

    if bkg_fit is None:
        return 2 * (poisson_log_like(sig + bkg, sig + bkg) -
                    poisson_log_like(sig + bkg, bkg))
    else:
        return 2 * (poisson_log_like(sig + bkg, sig + bkg_fit) -
                    poisson_log_like(sig + bkg, bkg_fit))

def poisson_ts_fast(sig, bkg, bkg_fit=None):

    if bkg_fit is None:
        return 2 * ((sig + bkg) * np.log((sig + bkg) / bkg) - sig)
    else:
        return 2 * ((sig + bkg) * np.log((sig + bkg_fit) / bkg_fit) - sig)

def compute_ext_flux(egy, flux):
    pass

def compute_ps_loc(egy, flux):
    """Solve for the localization precision of a point source with a given flux."""
    pass

def compute_ps_counts(ebins, exp, psf, bkg, fn, egy_dim=0, spatial_model='PointSource',
                      spatial_size=1E-3):
    """Calculate the observed signal and background counts given models
    for the exposure, background intensity, PSF, and source flux.

    Parameters
    ----------
    ebins : `~numpy.ndarray`
        Array of energy bin edges.

    exp : `~numpy.ndarray`
        Model for exposure.

    psf : `~fermipy.irfs.PSFModel`
        Model for average PSF.

    bkg : `~numpy.ndarray`
        Array of background intensities.

    fn : `~fermipy.spectrum.SpectralFunction`

    egy_dim : int
        Index of energy dimension in ``bkg`` and ``exp`` arrays.

    """
    ewidth = utils.edge_to_width(ebins)
    ectr = np.exp(utils.edge_to_center(np.log(ebins)))

    r68 = psf.containment_angle(ectr, fraction=0.68)
    if spatial_model != 'PointSource':
        r68[r68 < spatial_size] = spatial_size

    # * np.ones((len(ectr), 31))
    theta_edges = np.linspace(0.0, 3.0, 31)[np.newaxis, :]
    theta_edges = theta_edges * r68[:, np.newaxis]
    theta = 0.5 * (theta_edges[:, :-1] + theta_edges[:, 1:])
    domega = np.pi * (theta_edges[:, 1:]**2 - theta_edges[:, :-1]**2)

    if spatial_model == 'PointSource':
        sig_pdf = domega * psf.interp(ectr[:, np.newaxis], theta)
    elif spatial_model == 'RadialGaussian':
        sig_pdf = domega * utils.convolve2d_gauss(lambda t: psf.interp(ectr[:, np.newaxis, np.newaxis], t),
                                                  theta, spatial_size / 1.5095921854516636, nstep=2000)
    elif spatial_model == 'RadialDisk':
        sig_pdf = domega * utils.convolve2d_disk(lambda t: psf.interp(ectr[:, np.newaxis, np.newaxis], t),
                                                 theta, spatial_size / 0.8246211251235321)
    else:
        raise ValueError('Invalid spatial model: {}'.format(spatial_model))

    sig_pdf *= (np.pi / 180.)**2
    sig_flux = fn.flux(ebins[:-1], ebins[1:])

    # Background and signal counts
    bkgc = bkg[..., np.newaxis] * domega * exp[..., np.newaxis] * \
        ewidth[..., np.newaxis] * (np.pi / 180.)**2
    sigc = sig_pdf * sig_flux[..., np.newaxis] * exp[..., np.newaxis]

    return sigc, bkgc

def compute_norm(sig, bkg, ts_thresh, min_counts, sum_axes=None, bkg_fit=None,
                 rebin_axes=None):
    """Solve for the normalization of the signal distribution at which the
    detection test statistic (twice delta-loglikelihood ratio) is >=
    ``ts_thresh`` AND the number of signal counts >= ``min_counts``.
    This function uses the Asimov method to calculate the median
    expected TS when the model for the background is fixed (no
    uncertainty on the background amplitude).

    Parameters
    ----------
    sig : `~numpy.ndarray`
        Array of signal amplitudes in counts.

    bkg : `~numpy.ndarray`
        Array of background amplitudes in counts.

    ts_thresh : float
        Test statistic threshold.

    min_counts : float
        Counts threshold.

    sum_axes : list
        Axes over which the source test statistic should be summed.
        By default the summation will be performed over all
        dimensions.

    bkg_fit : `~numpy.ndarray`
        Array of background amplitudes in counts for the fitting
        model.  If None then the fit model will be equal to the data
        model.

    """

    if sum_axes is None:
        sum_axes = np.arange(sig.ndim)

    sig = np.expand_dims(sig, -1)
    bkg = np.expand_dims(bkg, -1)
    sig_sum = np.apply_over_axes(np.sum, sig, sum_axes)
    bkg_sum = np.apply_over_axes(np.sum, bkg, sum_axes)
    bkg_fit_sum = None

    if bkg_fit is not None:
        bkg_fit = np.expand_dims(bkg_fit, -1)
        bkg_fit_sum = np.apply_over_axes(np.sum, bkg_fit, sum_axes)

    sig_rebin = sig
    bkg_rebin = bkg
    bkg_fit_rebin = bkg_fit

    if rebin_axes:
        sig_rebin = sig.copy()
        bkg_rebin = bkg.copy()
        if bkg_fit is not None:
            bkg_fit_rebin = bkg_fit.copy()

        for dim, rebin in zip(sum_axes, rebin_axes):
            sig_rebin = sum_bins(sig_rebin, dim, rebin)
            bkg_rebin = sum_bins(bkg_rebin, dim, rebin)
            if bkg_fit is not None:
                bkg_fit_rebin = sum_bins(bkg_fit_rebin, dim, rebin)

    # Find approx solution using coarse binning and summed arrays
    sig_scale = 10**np.linspace(0.0, 10.0, 51) * (min_counts / sig_sum)
    vals_approx = _solve_norm(sig_rebin, bkg_rebin, ts_thresh, min_counts,
                              sig_scale, sum_axes, bkg_fit_rebin)

    # Refine solution using an interval (0.1,10) around approx
    # solution
    sig_scale = (10**np.linspace(0.0, 1.0, 21) *
                 np.fmax(0.333 * vals_approx[..., None],
                         min_counts / sig_sum))

    vals = _solve_norm(sig, bkg, ts_thresh, min_counts, sig_scale,
                       sum_axes, bkg_fit)

    #sig_scale = 10**np.linspace(0.0, 10.0, 101)*(min_counts / sig_sum)
    # vals = _solve_norm(sig, bkg, ts_thresh, min_counts, sig_scale2,
    #                   sum_axes, bkg_fit)

    return vals

def _solve_norm(sig, bkg, ts_thresh, min_counts, sig_scale, sum_axes,
                bkg_fit=None):

    ts = np.apply_over_axes(np.sum, poisson_ts_fast(sig * sig_scale,
                                                    bkg, bkg_fit),
                            sum_axes)

    vals = np.ones(ts.shape[:-1])
    for idx, v in np.ndenumerate(ts[..., 0]):

        if ts[idx][0] >= ts_thresh:
            vals[idx] = sig_scale[idx][0]
            continue

        m = slice(np.argmin(ts[idx]), None)
        vals[idx] = np.interp(ts_thresh, ts[idx][m], sig_scale[idx][m])
        #fn = UnivariateSpline(ts[idx][m], sig_scale[idx][m], k=2, s=0)
        #vals[idx] = fn(ts_thresh)

    return vals

class ExposureMap(HpxMap):

    def __init__(self, data, hpx):
        HpxMap.__init__(self, data, hpx)

    @classmethod
    def create(cls, ltc, event_class, event_types, ebins):
        """Create an exposure map from a livetime cube.  This method will
        generate an exposure map with the same geometry as the
        livetime cube (nside, etc.).

        Parameters
        ----------
        ltc : `~fermipy.irfs.LTCube`
            Livetime cube object.

        event_class : str
            Event class string.

        event_types : list
            List of event type strings, e.g. ['FRONT','BACK'].

        ebins :  `~numpy.ndarray`
            Energy bin edges in MeV.

        """

        evals = np.sqrt(ebins[1:] * ebins[:-1])
        exp = np.zeros((len(evals), ltc.hpx.npix))
        for et in event_types:
            aeff = create_aeff(event_class, et, evals, ltc.costh_center)
            exp += np.sum(aeff.T[:, :, np.newaxis] *
                          ltc.data[:, np.newaxis, :], axis=0)

        hpx = HPX(ltc.hpx.nside, ltc.hpx.nest,
                  ltc.hpx.coordsys, ebins=ebins)
        return cls(exp, hpx)

class PSFModel(object):
    """Class that stores a pre-computed model of the PSF versus energy.

    """

    def __init__(self, dtheta, energies, cth_bins, exp, psf, wts):
        """Create a PSFModel.

        Parameters
        ----------
        dtheta : `~numpy.ndarray`
            Array of angular offsets in degrees at which the PSF is
            evaluated.

        energies : `~numpy.ndarray`
            Array of energies in MeV at which the PSF is evaluated.

        cth_bins : `~numpy.ndarray`
            Interval in cosine of the incidence angle for which this
            model of the PSF was generated.

        exp : `~numpy.ndarray`
            Array of exposure vs. energy in cm^2 s.

        psf : `~numpy.ndarray`
            2D array of PSF values evaluated on an NxM grid of N
            offset angles and M energies (defined by ``dtheta`` and
            ``energies``).

        wts : `~numpy.ndarray`
            Array of weights vs. energy.  These are used to evaluate
            the bin-averaged PSF model.

        """

        self._dtheta = dtheta
        self._log_energies = np.log10(energies)
        self._energies = energies
        self._cth_bins = cth_bins
        self._cth = utils.edge_to_center(cth_bins)
        self._scale_fn = None
        self._exp = exp
        self._psf = psf
        self._wts = wts
        self._psf_fn = RegularGridInterpolator((self._dtheta, self._log_energies),
                                               np.log(self._psf),
                                               bounds_error=False,
                                               fill_value=None)
        self._wts_fn = RegularGridInterpolator((self._log_energies,),
                                               np.log(self._wts),
                                               bounds_error=False,
                                               fill_value=None)

    def eval(self, ebin, dtheta, scale_fn=None):
        """Evaluate the PSF at the given energy bin index.

        Parameters
        ----------
        ebin : int
            Index of energy bin.

        dtheta : array_like
            Array of angular separations in degrees.

        scale_fn : callable
            Function that evaluates the PSF scaling function.
            Argument is energy in MeV.
        """

        if scale_fn is None and self.scale_fn is not None:
            scale_fn = self.scale_fn

        if scale_fn is None:
            scale_factor = 1.0
        else:
            dtheta = dtheta / scale_fn(self.energies[ebin])
            scale_factor = 1. / scale_fn(self.energies[ebin])**2

        vals = 10**np.interp(dtheta, self.dtheta, np.log10(self.val[:, ebin]))
        return vals * scale_factor

    def interp(self, energies, dtheta, scale_fn=None):
        """Evaluate the PSF model at an array of energies and angular
        separations.

        Parameters
        ----------
        energies : array_like
            Array of energies in MeV.

        dtheta : array_like
            Array of angular separations in degrees.

        scale_fn : callable
            Function that evaluates the PSF scaling function.
            Argument is energy in MeV.
        """

        if scale_fn is None and self.scale_fn:
            scale_fn = self.scale_fn

        log_energies = np.log10(energies)

        shape = (energies * dtheta).shape
        scale_factor = np.ones(shape)

        if scale_fn is not None:
            dtheta = dtheta / scale_fn(energies)
            scale_factor = 1. / scale_fn(energies)**2

        vals = np.exp(self._psf_fn((dtheta, log_energies)))
        return vals * scale_factor

    def interp_bin(self, egy_bins, dtheta, scale_fn=None):
        """Evaluate the bin-averaged PSF model over the energy bins ``egy_bins``.

        Parameters
        ----------
        egy_bins : array_like
            Energy bin edges in MeV.

        dtheta : array_like
            Array of angular separations in degrees.

        scale_fn : callable
            Function that evaluates the PSF scaling function.
            Argument is energy in MeV.
        """

        npts = 4
        egy_bins = np.exp(utils.split_bin_edges(np.log(egy_bins), npts))
        egy = np.exp(utils.edge_to_center(np.log(egy_bins)))
        log_energies = np.log10(egy)

        vals = self.interp(egy[None, :], dtheta[:, None],
                           scale_fn=scale_fn)
        wts = np.exp(self._wts_fn((log_energies,)))
        wts = wts.reshape((1,) + wts.shape)
        vals = np.sum(
            (vals * wts).reshape((vals.shape[0], int(vals.shape[1] / npts), npts)), axis=2)
        vals /= np.sum(wts.reshape(wts.shape[0],
                                   int(wts.shape[1] / npts), npts), axis=2)
        return vals

    def containment_angle(self, energies=None, fraction=0.68, scale_fn=None):
        """Evaluate the PSF containment angle at a sequence of energies."""

        if energies is None:
            energies = self.energies

        vals = self.interp(energies[np.newaxis, :], self.dtheta[:, np.newaxis],
                           scale_fn=scale_fn)
        dtheta = np.radians(self.dtheta[:, np.newaxis] * np.ones(vals.shape))
        return self._calc_containment(dtheta, vals, fraction)

    def containment_angle_bin(self, egy_bins, fraction=0.68, scale_fn=None):
        """Evaluate the PSF containment angle averaged over energy bins."""

        vals = self.interp_bin(egy_bins, self.dtheta, scale_fn=scale_fn)
        dtheta = np.radians(self.dtheta[:, np.newaxis] * np.ones(vals.shape))
        return self._calc_containment(dtheta, vals, fraction)

    def _calc_containment(self, dtheta, vals, fraction=0.68):

        delta = dtheta[1:] - dtheta[:-1]
        ctr = 0.5 * (dtheta[1:] + dtheta[:-1])

        avg_val = 0.5 * (vals[1:, :] * np.sin(dtheta[1:, :]) +
                         vals[:-1, :] * np.sin(dtheta[:-1, :]))

        csum = delta * avg_val * 2 * np.pi
        csum = np.cumsum(csum, axis=0)
        theta = np.zeros(csum.shape[1])

        for i in range(csum.shape[1]):
            theta[i] = np.degrees(np.interp(fraction, csum[:, i],
                                            dtheta[1:, i]))

        return theta

    def set_scale_fn(self, scale_fn):
        self._scale_fn = scale_fn

    @property
    def scale_fn(self):
        return self._scale_fn

    @property
    def dtheta(self):
        return self._dtheta

    @property
    def log_energies(self):
        return self._log_energies

    @property
    def energies(self):
        return self._energies

    @property
    def val(self):
        return self._psf

    @property
    def exp(self):
        return self._exp

    @classmethod
    def create(cls, skydir, ltc, event_class, event_types, energies, cth_bins=None,
               ndtheta=500, use_edisp=False, fn=None, nbin=64):
        """Create a PSFModel object.  This class can be used to evaluate the
        exposure-weighted PSF for a source with a given observing
        profile and energy distribution.

        Parameters
        ----------
        skydir : `~astropy.coordinates.SkyCoord`

        ltc : `~fermipy.irfs.LTCube`

        energies : `~numpy.ndarray`
            Grid of energies at which the PSF will be pre-computed.

        cth_bins : `~numpy.ndarray`
            Bin edges in cosine of the inclination angle.

        use_edisp : bool
            Generate the PSF model accounting for the influence of
            energy dispersion.

        fn : `~fermipy.spectrum.SpectralFunction`
            Model for the spectral energy distribution of the source.

        """

        if isinstance(event_types, int):
            event_types = bitmask_to_bits(event_types)

        if fn is None:
            fn = spectrum.PowerLaw([1E-13, -2.0])

        dtheta = np.logspace(-4, 1.75, ndtheta)
        dtheta = np.insert(dtheta, 0, [0])
        log_energies = np.log10(energies)
        egy_bins = 10**utils.center_to_edge(log_energies)

        if cth_bins is None:
            cth_bins = np.array([0.2, 1.0])

        if use_edisp:
            psf = create_wtd_psf(skydir, ltc, event_class, event_types,
                                 dtheta, egy_bins, cth_bins, fn, nbin=nbin)
            wts = calc_counts_edisp(skydir, ltc, event_class, event_types,
                                    egy_bins, cth_bins, fn, nbin=nbin)
        else:
            psf = create_avg_psf(skydir, ltc, event_class, event_types,
                                 dtheta, energies, cth_bins)
            wts = calc_counts(skydir, ltc, event_class, event_types,
                              egy_bins, cth_bins, fn)

        exp = calc_exp(skydir, ltc, event_class, event_types,
                       energies, cth_bins)

        return cls(dtheta, energies, cth_bins, np.squeeze(exp), np.squeeze(psf),
                   np.squeeze(wts))

def create_irf(event_class, event_type):
    if isinstance(event_type, int):
        event_type = evtype_string[event_type]

    irf_factory = pyIrfLoader.IrfsFactory.instance()
    irfname = '%s::%s' % (event_class, event_type)
    irf = irf_factory.create(irfname)
    return irf

def create_psf(event_class, event_type, dtheta, egy, cth):
    """Create an array of PSF response values versus energy and
    inclination angle.

    Parameters
    ----------
    egy : `~numpy.ndarray`
        Energy in MeV.

    cth : `~numpy.ndarray`
        Cosine of the incidence angle.

    """
    irf = create_irf(event_class, event_type)
    theta = np.degrees(np.arccos(cth))
    m = np.zeros((len(dtheta), len(egy), len(cth)))

    for i, x in enumerate(egy):
        for j, y in enumerate(theta):
            m[:, i, j] = irf.psf().value(dtheta, x, y, 0.0)

    return m

def create_edisp(event_class, event_type, erec, egy, cth):
    """Create an array of energy response values versus energy and
    inclination angle.

    Parameters
    ----------
    egy : `~numpy.ndarray`
        Energy in MeV.

    cth : `~numpy.ndarray`
        Cosine of the incidence angle.

    """
    irf = create_irf(event_class, event_type)
    theta = np.degrees(np.arccos(cth))
    v = np.zeros((len(erec), len(egy), len(cth)))
    m = (erec[:,None] / egy[None,:] < 3.0) & (erec[:,None] / egy[None,:] > 0.33333)
    #    m |= ((erec[:,None] / egy[None,:] < 3.0) &
    #          (erec[:,None] / egy[None,:] > 0.5) & (egy[None,:] < 10**2.5))
    m = np.broadcast_to(m[:,:,None], v.shape)

    try:
        x = np.ones(v.shape)*erec[:,None,None]
        y = np.ones(v.shape)*egy[None,:,None]
        z = np.ones(v.shape)*theta[None,None,:]
        v[m] = irf.edisp().value(np.ravel(x[m]), np.ravel(y[m]), np.ravel(z[m]), 0.0)
    except:
        for i, x in enumerate(egy):
            for j, y in enumerate(theta):
                m = (erec / x < 3.0) & (erec / x > 0.333)
                v[m, i, j] = irf.edisp().value(erec[m], x, y, 0.0)

    return v

def create_aeff(event_class, event_type, egy, cth):
    """Create an array of effective areas versus energy and incidence
    angle.  Binning in energy and incidence angle is controlled with
    the egy and cth input parameters.

    Parameters
    ----------
    event_class : str
        Event class string (e.g. P8R2_SOURCE_V6).

    event_type : list

    egy : array_like
        Evaluation points in energy (MeV).

    cth : array_like
        Evaluation points in cosine of the incidence angle.

    """
    irf = create_irf(event_class, event_type)
    irf.aeff().setPhiDependence(False)
    theta = np.degrees(np.arccos(cth))

    # Exposure Matrix
    # Dimensions are Etrue and incidence angle
    m = np.zeros((len(egy), len(cth)))

    for i, x in enumerate(egy):
        for j, y in enumerate(theta):
            m[i, j] = irf.aeff().value(x, y, 0.0)

    return m

def calc_exp(skydir, ltc, event_class, event_types,
             egy, cth_bins, npts=None):
    """Calculate the exposure on a 2D grid of energy and incidence angle.

    Parameters
    ----------
    npts : int
        Number of points by which to sample the response in each
        incidence angle bin.  If None then npts will be automatically
        set such that incidence angle is sampled on intervals of <
        0.05 in Cos(Theta).

    Returns
    -------
    exp : `~numpy.ndarray`
        2D Array of exposures vs. energy and incidence angle.

    """

    if npts is None:
        npts = int(np.ceil(np.max(cth_bins[1:] - cth_bins[:-1]) / 0.025))

    exp = np.zeros((len(egy), len(cth_bins) - 1))
    cth_bins = utils.split_bin_edges(cth_bins, npts)
    cth = edge_to_center(cth_bins)
    ltw = ltc.get_skydir_lthist(skydir, cth_bins).reshape(-1, npts)
    for et in event_types:
        aeff = create_aeff(event_class, et, egy, cth)
        aeff = aeff.reshape(exp.shape + (npts,))
        exp += np.sum(aeff * ltw[np.newaxis, :, :], axis=-1)

    return exp

def create_avg_rsp(rsp_fn, skydir, ltc, event_class, event_types, x,
                   egy, cth_bins, npts=None):
    """Calculate the weighted response function.
    """
    if npts is None:
        npts = int(np.ceil(np.max(cth_bins[1:] - cth_bins[:-1]) / 0.05))

    wrsp = np.zeros((len(x), len(egy), len(cth_bins) - 1))
    exps = np.zeros((len(egy), len(cth_bins) - 1))

    cth_bins = utils.split_bin_edges(cth_bins, npts)
    cth = edge_to_center(cth_bins)
    ltw = ltc.get_skydir_lthist(skydir, cth_bins)
    ltw = ltw.reshape(-1, npts)

    for et in event_types:
        rsp = rsp_fn(event_class, et, x, egy, cth)
        aeff = create_aeff(event_class, et, egy, cth)
        rsp = rsp.reshape(wrsp.shape + (npts,))
        aeff = aeff.reshape(exps.shape + (npts,))
        wrsp += np.sum(rsp * aeff[np.newaxis, :, :, :] *
                       ltw[np.newaxis, np.newaxis, :, :], axis=-1)
        exps += np.sum(aeff * ltw[np.newaxis, :, :], axis=-1)

    exps_inv = np.zeros_like(exps)
    exps_inv[exps > 0] = 1./exps[exps>0]
    wrsp *= exps_inv[np.newaxis, :, :]
    return wrsp

def create_avg_psf(skydir, ltc, event_class, event_types, dtheta,
                   egy, cth_bins, npts=None):
    """Generate model for exposure-weighted PSF averaged over incidence
    angle.

    Parameters
    ----------
    egy : `~numpy.ndarray`
        Energies in MeV.

    cth_bins : `~numpy.ndarray`
        Bin edges in cosine of the incidence angle.
    """

    return create_avg_rsp(create_psf, skydir, ltc,
                          event_class, event_types,
                          dtheta, egy,  cth_bins, npts)

def create_avg_edisp(skydir, ltc, event_class, event_types, erec,
                     egy, cth_bins, npts=None):
    """Generate model for exposure-weighted DRM averaged over incidence
    angle.

    Parameters
    ----------
    egy : `~numpy.ndarray`
        True energies in MeV.

    cth_bins : `~numpy.ndarray`
        Bin edges in cosine of the incidence angle.
    """
    return create_avg_rsp(create_edisp, skydir, ltc,
                          event_class, event_types,
                          erec, egy,  cth_bins, npts)

def create_wtd_psf(skydir, ltc, event_class, event_types, dtheta,
                   egy_bins, cth_bins, fn, nbin=64, npts=1):
    """Create an exposure- and dispersion-weighted PSF model for a source
    with spectral parameterization ``fn``.  The calculation performed
    by this method accounts for the influence of energy dispersion on
    the PSF.

    Parameters
    ----------
    dtheta : `~numpy.ndarray`

    egy_bins : `~numpy.ndarray`
        Bin edges in observed energy.

    cth_bins : `~numpy.ndarray`
        Bin edges in cosine of the true incidence angle.

    nbin : int
        Number of bins per decade in true energy.

    npts : int
        Number of points by which to oversample each energy bin.

    """
    #npts = int(np.ceil(32. / bins_per_dec(egy_bins)))
    egy_bins = np.exp(utils.split_bin_edges(np.log(egy_bins), npts))
    etrue_bins = 10**np.linspace(1.0, 6.5, nbin * 5.5 + 1)
    etrue = 10**utils.edge_to_center(np.log10(etrue_bins))

    psf = create_avg_psf(skydir, ltc, event_class, event_types, dtheta,
                         etrue, cth_bins)
    drm = calc_drm(skydir, ltc, event_class, event_types,
                   egy_bins, cth_bins, nbin=nbin)
    cnts = calc_counts(skydir, ltc, event_class, event_types,
                       etrue_bins, cth_bins, fn)

    wts = drm * cnts[None, :, :]
    wts_norm = np.sum(wts, axis=1)
    wts_norm[wts_norm == 0] = 1.0
    wts = wts / wts_norm[:, None, :]
    wpsf = np.sum(wts[None, :, :, :] * psf[:, None, :, :], axis=2)
    wts = np.sum(wts[None, :, :, :], axis=2)

    if npts > 1:
        shape = (wpsf.shape[0], int(wpsf.shape[1] / npts), npts, wpsf.shape[2])
        wpsf = np.sum((wpsf * wts).reshape(shape), axis=2)
        shape = (wts.shape[0], int(wts.shape[1] / npts), npts, wts.shape[2])
        wpsf = wpsf / np.sum(wts.reshape(shape), axis=2)

    return wpsf

def calc_drm(skydir, ltc, event_class, event_types,
             egy_bins, cth_bins, nbin=64):
    """Calculate the detector response matrix."""
    npts = int(np.ceil(128. / bins_per_dec(egy_bins)))
    egy_bins = np.exp(utils.split_bin_edges(np.log(egy_bins), npts))

    etrue_bins = 10**np.linspace(1.0, 6.5, nbin * 5.5 + 1)
    egy = 10**utils.edge_to_center(np.log10(egy_bins))
    egy_width = utils.edge_to_width(egy_bins)
    etrue = 10**utils.edge_to_center(np.log10(etrue_bins))
    edisp = create_avg_edisp(skydir, ltc, event_class, event_types,
                             egy, etrue, cth_bins)
    edisp = edisp * egy_width[:, None, None]
    edisp = sum_bins(edisp, 0, npts)
    return edisp

def calc_counts(skydir, ltc, event_class, event_types,
                egy_bins, cth_bins, fn, npts=1):
    """Calculate the expected counts vs. true energy and incidence angle
    for a source with spectral parameterization ``fn``.

    Parameters
    ----------
    skydir : `~astropy.coordinate.SkyCoord`

    ltc : `~fermipy.irfs.LTCube`

    egy_bins : `~numpy.ndarray`
        Bin edges in observed energy in MeV.

    cth_bins : `~numpy.ndarray`
        Bin edges in cosine of the true incidence angle.

    npts : int
        Number of points by which to oversample each energy bin.
    """
    #npts = int(np.ceil(32. / bins_per_dec(egy_bins)))
    egy_bins = np.exp(utils.split_bin_edges(np.log(egy_bins), npts))
    exp = calc_exp(skydir, ltc, event_class, event_types,
                   egy_bins, cth_bins)
    dnde = fn.dnde(egy_bins)
    cnts = loglog_quad(egy_bins, exp * dnde[:, None], 0)
    cnts = sum_bins(cnts, 0, npts)
    return cnts

def calc_counts_edisp(skydir, ltc, event_class, event_types,
                      egy_bins, cth_bins, fn, nbin=16, npts=1):
    """Calculate the expected counts vs. observed energy and true
    incidence angle for a source with spectral parameterization ``fn``.

    Parameters
    ----------
    skydir : `~astropy.coordinate.SkyCoord`

    ltc : `~fermipy.irfs.LTCube`

    egy_bins : `~numpy.ndarray`
        Bin edges in observed energy in MeV.

    cth_bins : `~numpy.ndarray`
        Bin edges in cosine of the true incidence angle.

    nbin : int
        Number of points per decade with which to sample true energy.

    npts : int
        Number of points by which to oversample each reconstructed energy bin.

    """
    #npts = int(np.ceil(32. / bins_per_dec(egy_bins)))

    # Split energy bins
    egy_bins = np.exp(utils.split_bin_edges(np.log(egy_bins), npts))
    etrue_bins = 10**np.linspace(1.0, 6.5, nbin * 5.5 + 1)
    drm = calc_drm(skydir, ltc, event_class, event_types,
                   egy_bins, cth_bins, nbin=nbin)
    cnts_etrue = calc_counts(skydir, ltc, event_class, event_types,
                             etrue_bins, cth_bins, fn)

    cnts = np.sum(cnts_etrue[None, :, :] * drm[:, :, :], axis=1)
    cnts = sum_bins(cnts, 0, npts)
    return cnts

def calc_wtd_exp(skydir, ltc, event_class, event_types,
                 egy_bins, cth_bins, fn, nbin=16):
    """Calculate the effective exposure.

    Parameters
    ----------
    skydir : `~astropy.coordinates.SkyCoord`

    ltc : `~fermipy.irfs.LTCube`

    nbin : int
        Number of points per decade with which to sample true energy.

    """
    cnts = calc_counts_edisp(skydir, ltc, event_class, event_types,
                             egy_bins, cth_bins, fn, nbin=nbin)
    flux = fn.flux(egy_bins[:-1], egy_bins[1:])
    return cnts / flux[:, None]

def plot_hpxmap(hpxmap, **kwargs):

    import matplotlib as mpl
    import matplotlib.pyplot as plt
    from matplotlib.colors import PowerNorm, Normalize, LogNorm

    zidx = kwargs.pop('zidx', None)

    kwargs_imshow = {'norm': None,
                     'vmin': None, 'vmax': None}

    gamma = kwargs.get('gamma', 2.0)
    zscale = kwargs.get('zscale', None)
    cbar = kwargs.get('cbar', None)
    cbar_label = kwargs.get('cbar_label', '')
    title = kwargs.get('title', '')
    levels = kwargs.get('levels', None)
    rot = kwargs.get('rot', None)

    kwargs_imshow['vmin'] = kwargs.get('vmin', None)
    kwargs_imshow['vmax'] = kwargs.get('vmax', None)

    cmap = mpl.cm.get_cmap(kwargs.get('cmap', 'jet'))
    cmap.set_under('white')
    kwargs_imshow['cmap'] = cmap

    if zscale == 'pow':
        vmed = np.median(hpxmap.counts)
        vmax = max(hpxmap.counts)
        vmin = min(1.1 * hpxmap.counts[hpxmap.counts > 0])
        kwargs_imshow['norm'] = PowerNorm(gamma=gamma, clip=True)
    elif zscale == 'log':
        kwargs_imshow['norm'] = LogNorm()
    else:
        kwargs_imshow['norm'] = Normalize(clip=True)

    from healpy import projaxes as PA

    fig = plt.gcf()
    if 'sub' in kwargs:
        sub = kwargs['sub']
        nrows, ncols, idx = sub / 100, (sub % 100) / 10, (sub % 10)
        c, r = (idx - 1) % ncols, (idx - 1) / ncols
        margins = (0.01, 0.0, 0.0, 0.02)
        extent = (c * 1. / ncols + margins[0],
                  1. - (r + 1) * 1. / nrows + margins[1],
                  1. / ncols - margins[2] - margins[0],
                  1. / nrows - margins[3] - margins[1])
        extent = (extent[0] + margins[0],
                  extent[1] + margins[1],
                  extent[2] - margins[2] - margins[0],
                  extent[3] - margins[3] - margins[1])
    else:
        extent = (0.02, 0.05, 0.96, 0.9)

    ax = hp.projaxes.HpxMollweideAxes(fig, extent, coord=None, rot=rot,
                                      format='%g', flipconv='astro')

    ax.set_title(title)
    fig.add_axes(ax)

    if zidx is not None:
        data = hpxmap.data[zidx]
    elif hpxmap.data.ndim == 2:
        data = np.sum(hpxmap.data, axis=0)
    else:
        data = hpxmap.data

    img0 = ax.projmap(data, nest=hpxmap.hpx.nest, xsize=1600, coord='C',
                      **kwargs_imshow)

    if levels:
        cs = ax.contour(img0, extent=ax.proj.get_extent(),
                        levels=levels, colors=['k'],
                        interpolation='nearest')

    hp.visufunc.graticule(verbose=False, lw=0.5, color='k')

    if cbar is not None:

        im = ax.get_images()[0]
        cb_kw = dict(orientation='vertical',
                     shrink=.6, pad=0.05)
        cb_kw.update(cbar)
        cb = fig.colorbar(im, **cb_kw)  # ,format='%.3g')
        cb.set_label(cbar_label)
        #, ticks=[min, max])
#            cb.ax.xaxis.set_label_text(cbar_label)


'''
this is not the original prox fitting, but a modified version just for post-processing
our generated results. The input is the smplx body parameters, and the optimization
is based on the scene sdf and the contact loss
'''

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import pickle
import sys, os, glob
import pdb
import json
import argparse
import numpy as np
import open3d as o3d

sys.path.append('/home/yzhang/workspaces/smpl-env-gen-3d-internal')

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import init
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable

import smplx
from human_body_prior.tools.model_loader import load_vposer
import chamfer_pytorch.dist_chamfer as ext

from cvae import BodyParamParser, GeometryTransformer

class FittingOP:
    def __init__(self, fittingconfig, lossconfig):

        for key, val in fittingconfig.items():
            setattr(self, key, val)

        for key, val in lossconfig.items():
            setattr(self, key, val)

        self.vposer, _ = load_vposer(self.vposer_ckpt_path, vp_model='snapshot')
        self.body_mesh_model = smplx.create(self.human_model_path, model_type='smplx',
                                       gender='neutral', ext='npz',
                                       num_pca_comps=12,
                                       create_global_orient=True,
                                       create_body_pose=True,
                                       create_betas=True,
                                       create_left_hand_pose=True,
                                       create_right_hand_pose=True,
                                       create_expression=True,
                                       create_jaw_pose=True,
                                       create_leye_pose=True,
                                       create_reye_pose=True,
                                       create_transl=True,
                                       batch_size=self.batch_size
                                       )
        self.vposer.to(self.device)
        self.body_mesh_model.to(self.device)

        self.xhr_rec = Variable(torch.randn(1,75).to(self.device), requires_grad=True)
        self.optimizer = optim.Adam([self.xhr_rec], lr=self.init_lr_h)

        ## read scene sdf
        with open(self.scene_sdf_path+'.json') as f:
                sdf_data = json.load(f)
                grid_min = np.array(sdf_data['min'])
                grid_max = np.array(sdf_data['max'])
                grid_dim = sdf_data['dim']
        sdf = np.load(self.scene_sdf_path + '_sdf.npy').reshape(grid_dim, grid_dim, grid_dim)

        self.s_grid_min_batch = torch.tensor(grid_min, dtype=torch.float32, device=self.device).unsqueeze(0)
        self.s_grid_max_batch = torch.tensor(grid_max, dtype=torch.float32, device=self.device).unsqueeze(0)
        self.s_sdf_batch = torch.tensor(sdf, dtype=torch.float32, device=self.device).unsqueeze(0)

        ## read scene vertices
        scene_o3d = o3d.io.read_triangle_mesh(self.scene_verts_path)
        scene_verts = np.asarray(scene_o3d.vertices)
        self.s_verts_batch = torch.tensor(scene_verts, dtype=torch.float32, device=self.device).unsqueeze(0)

    def cal_loss(self, xhr, cam_ext):

        ### reconstruction loss
        loss_rec = self.weight_loss_rec*F.l1_loss(xhr, self.xhr_rec)
        xh_rec = GeometryTransformer.convert_to_3D_rot(self.xhr_rec)

        ### vposer loss
        vposer_pose = xh_rec[:,16:48]
        loss_vposer = self.weight_loss_vposer * torch.mean(vposer_pose**2)

        ### contact loss
        body_param_rec = BodyParamParser.body_params_encapsulate_batch(xh_rec)
        joint_rot_batch = self.vposer.decode(body_param_rec['body_pose_vp'],
                                           output_type='aa').view(self.batch_size, -1)

        body_param_ = {}
        for key in body_param_rec.keys():
            if key in ['body_pose_vp']:
                continue
            else:
                body_param_[key] = body_param_rec[key]

        smplx_output = self.body_mesh_model(return_verts=True,
                                              body_pose=joint_rot_batch,
                                              **body_param_)
        body_verts_batch = smplx_output.vertices #[b, 10475,3]
        body_verts_batch = GeometryTransformer.verts_transform(body_verts_batch, cam_ext)

        vid, fid = GeometryTransformer.get_contact_id(
                                body_segments_folder=self.contact_id_folder,
                                contact_body_parts=self.contact_part)
        body_verts_contact_batch = body_verts_batch[:, vid, :]

        dist_chamfer_contact = ext.chamferDist()
        contact_dist, _ = dist_chamfer_contact(body_verts_contact_batch.contiguous(),
                                                self.s_verts_batch.contiguous())

        loss_contact = self.weight_contact * torch.mean(torch.sqrt(contact_dist+1e-4)/(torch.sqrt(contact_dist+1e-4)+1.0))

        ### sdf collision loss
        s_grid_min_batch = self.s_grid_min_batch.unsqueeze(1)
        s_grid_max_batch = self.s_grid_max_batch.unsqueeze(1)

        norm_verts_batch = (body_verts_batch - s_grid_min_batch) / (s_grid_max_batch - s_grid_min_batch) *2 -1
        n_verts = norm_verts_batch.shape[1]
        body_sdf_batch = F.grid_sample(self.s_sdf_batch.unsqueeze(1),
                                        norm_verts_batch[:,:,[2,1,0]].view(-1, n_verts,1,1,3),
                                        padding_mode='border')

        # if there are no penetrating vertices then set sdf_penetration_loss = 0
        if body_sdf_batch.lt(0).sum().item() < 1:
            loss_sdf_pene = torch.tensor(0.0, dtype=torch.float32, device=self.device)
        else:
            loss_sdf_pene = body_sdf_batch[body_sdf_batch < 0].abs().mean()

        loss_collision = self.weight_collision*loss_sdf_pene

        return loss_rec, loss_vposer, loss_contact, loss_collision

    def fitting(self, input_data_file):

        with open(input_data_file, 'rb') as f:
            body_param_input = pickle.load(f)

        xh, self.cam_ext, self.cam_int= BodyParamParser.body_params_parse_fitting(body_param_input)
        xhr = GeometryTransformer.convert_to_6D_rot(xh)
        self.xhr_rec.data = xhr.clone()

        T_mat = np.eye(4)
        T_mat[1,:] = np.array([0,-1,0,0])
        T_mat[2,:] = np.array([0,0,-1,0])
        T_mat = torch.tensor(T_mat, dtype=torch.float32, device=self.device)
        T_mat = T_mat.unsqueeze(0)
        trans = torch.matmul(self.cam_ext[:1],T_mat)

        for ii in range(self.num_iter):

            self.optimizer.zero_grad()

            loss_rec, loss_vposer, loss_contact, loss_collision = self.cal_loss(xhr, trans)
            loss = loss_rec + loss_vposer + loss_contact + loss_collision
            if self.verbose:
                print('[INFO][fitting] iter={:d}, l_rec={:f}, l_vposer={:f}, l_contact={:f}, l_collision={:f}'.format(
                                        ii, loss_rec.item(), loss_vposer.item(),
                                        loss_contact.item(), loss_collision.item()) )

            loss.backward(retain_graph=True)
            self.optimizer.step()

        print('[INFO][fitting] fitting finish, returning optimal value')

        xh_rec =  GeometryTransformer.convert_to_3D_rot(self.xhr_rec)

        return xh_rec

    def save_result(self, xh_rec, output_data_file):

        dirname = os.path.dirname(output_data_file)

        if not os.path.exists(dirname):
            os.makedirs(dirname)

        body_param_list = BodyParamParser.body_params_encapsulate(xh_rec)
        print('[INFO] save results to: '+output_data_file)
        for ii, body_param in enumerate(body_param_list):
            # print(body_param['transl'])
            # print(body_param['global_orient'])
            # print()
            body_param['cam_ext'] = self.cam_ext.detach().cpu().numpy()
            body_param['cam_int'] = self.cam_int.detach().cpu().numpy()
            outfile = open(output_data_file, 'wb')
            pickle.dump(body_param, outfile)
            outfile.close()

if __name__=='__main__':

    gen_path = sys.argv[1]
    fit_path = sys.argv[2]

    scene_test_list = ['17DRP5sb8fy-bedroom', '17DRP5sb8fy-familyroomlounge',
                        '17DRP5sb8fy-livingroom', 'sKLMLpTHeUy-familyname_0_1',
                        'X7HyMhZNoso-livingroom_0_16', 'zsNo4HB9uLZ-bedroom0_0',
                        'zsNo4HB9uLZ-livingroom0_13']

    mp3dr_path = '/is/cluster/yzhang/mp3d-rooms'

    for scenename in scene_test_list:

        fittingconfig={
            'scene_verts_path': os.path.join(mp3dr_path, scenename+'.ply'),
            'scene_sdf_path': os.path.join(mp3dr_path, 'sdf/'+scenename),
            'human_model_path': '/is/ps2/yzhang/body_models/VPoser',
            'vposer_ckpt_path': '/is/ps2/yzhang/body_models/VPoser/vposer_v1_0',
            'init_lr_h': 0.1,
            'num_iter': 50,
            'batch_size': 1,
            'device': torch.device("cuda" if torch.cuda.is_available() else "cpu"),
            'contact_id_folder': '/is/cluster/work/yzhang/PROX/body_segments',
            'contact_part': ['back','butt','L_Hand','R_Hand','L_Leg','R_Leg','thighs'],
            'verbose': False
        }

        lossconfig={
            'weight_loss_rec': 1,
            'weight_loss_vposer':0.01,
            'weight_contact': 0.1,
            'weight_collision' : 0.5
        }

        fop = FittingOP(fittingconfig, lossconfig)

        for ii in range(10000):
            input_data_file = os.path.join(gen_path,scenename+'/body_gen_{:06d}.pkl'.format(ii))

            if not os.path.exists(input_data_file):
                continue

            output_data_file = os.path.join(fit_path,scenename+'/body_gen_{:06d}.pkl'.format(ii))
            if os.path.exists(output_data_file):
                continue

            xh_rec = fop.fitting(input_data_file)

            if xh_rec is None:
                continue
            fop.save_result(xh_rec, output_data_file)


class fpoutClass:
  pass

def readout(outfile,atmfile,fpout):
  import h5py
  import numpy as np
  with h5py.File(outfile, 'r') as h5f:
     fpout.inputparams = dict(h5f['inputparams'].attrs)
     fpout.atm = dict(h5f['atm'].attrs)
     fpout.atm['zin'] = np.array(h5f['atm/zin'])
     fpout.atm['tg'] = np.array(h5f['atm/tg'])
     fpout.atm['bfield'] = np.array(h5f['atm/bfield'])
     fpout.atm['dni'] = np.array(h5f['atm/dni'])
     fpout.atm['dnn'] = np.array(h5f['atm/dnn'])
     fpout.atm['mion'] = np.array(h5f['atm/mion'])
     fpout.atm['Zion'] = np.array(h5f['atm/Zion'])
     fpout.atm['Zn'] = np.array(h5f['atm/Zn'])
     fpout.atm['Enion'] = np.array(h5f['atm/Enion'])
     fpout.E = np.array(h5f['E'])
     fpout.mu = np.array(h5f['mu'])
     fpout.z = np.array(h5f['z'])
     fpout.esvol = np.array(h5f['esvol'])
     fpout.heatrate = np.array(h5f['heatrate'])
     fpout.momrate =  np.array(h5f['momrate'])
     fpout.f = np.array(h5f['f'])
  makefpoutstr(fpout,atmfile)

def makefpoutstr(fpout,atmfile):
  import numpy as np
  import const as c
  fpout.atmfile = atmfile
  fpout.nE = fpout.inputparams['nE']
  fpout.nmu = fpout.inputparams['nmu']
  fpout.nz = fpout.inputparams['nz']
  nE = fpout.nE
  nmu = fpout.nmu
  nz = fpout.nz
  fpout.Em = .5*(fpout.E[1:]+fpout.E[0:nE])
  mbeam = fpout.inputparams['mbeam']
  if (fpout.inputparams['inc_relativity']):
    fpout.gma = fpout.E/(mbeam/1e3) + 1e0
    fpout.bta = np.sqrt(1e0 - pow(fpout.gma,-2))
    fpout.gmam = fpout.Em/(mbeam/1e3) + 1e0
    fpout.btam = np.sqrt(1e0 - pow(fpout.gmam,-2))
  else:
    fpout.gma = np.ones(nE+1)
    fpout.bta = np.sqrt(2*fpout.E*1e3/mbeam)
    fpout.gmam = np.ones(nE)
    fpout.btam = np.sqrt(2*fpout.Em*1e3/mbeam)

  fpout.theta = np.arccos(fpout.mu)
  fpout.thetam = fpout.theta[0] if fpout.inputparams['oneD'] else .5*(fpout.theta[1:]+fpout.theta[0:nmu])
  fpout.mum = np.cos(fpout.thetam)

  fpout.zm = .5*(fpout.z[1:]+fpout.z[0:nz])
  #Calculate useful fluxes
  vz = c.clight *np.outer(fpout.mum, fpout.btam)
  vx = c.clight * np.outer(np.sin(fpout.thetam), fpout.btam)
  domega = fpout.esvol/(fpout.E[1:] - fpout.E[0:nE]) # solid angle part of esvol

  fpout.eflux = np.einsum('ijk,jk,k',fpout.f,vz*fpout.esvol,fpout.Em)*1e3*c.ergperev # Energy flux in z direction (erg cm^-2 s^-1)
  fpout.efluxx = np.einsum('ijk,jk,k',fpout.f,vx*fpout.esvol,fpout.Em)*1e3*c.ergperev # Energy flux in x direction (erg cm^-2 s^-1)
  fpout.nflux = np.einsum('ijk,jk',fpout.f,vz*fpout.esvol)# number flux in z direction (particles cm^-2 s^-1)
  fpout.flux = np.einsum('ijk,jk->ik',fpout.f,vz*domega) # number flux distribution in z direction (particles cm^-2 s^-1 keV^-1)
  fpout.flx = np.einsum('ijk,jk,k->ik',fpout.f,domega,fpout.btam)*c.clight # number flux speed distribution. (particles cm^-2 s^-1 keV^-1)

  if (abs(mbeam-c.me) < 1e3): #electron beam
    import fpbrem
    fpout.Eph = fpout.Em
    fpout.brem = fpbrem.fpbrem(fpout.Em,fpout.flx,fpout.Eph,fpout.atm)
    fpout.totbrem = np.einsum('ij,j',fpout.brem,np.diff(fpout.z))
  else:
  # Not electron beam so do not calculate bremsstrahlung
    fpout.brem = 0e0
    fpout.totbrem = 0e0
    fpout.Eph = 0e0

def writefpout(fpout,outfile = 'out.h5'):
  import h5py
  with h5py.File(outfile,'w') as fle:
    ipg = fle.create_group('inputparams')
    ipg.attrs.update(fpout.inputparams)
    atmattrs = {k:fpout.atm[k] for k in ('nIon','nNeutral') if k in fpout.atm}
    atmg = fle.create_group('atm')
    atmg.attrs.update(atmattrs)
    atmg.create_dataset('zin',data=fpout.atm['zin']); atmg.create_dataset('tg',data=fpout.atm['tg']); atmg.create_dataset('bfield',data=fpout.atm['bfield'])
    atmg.create_dataset('dni',data=fpout.atm['dni']); atmg.create_dataset('dnn',data=fpout.atm['dnn']); atmg.create_dataset('mion',data=fpout.atm['mion'])
    atmg.create_dataset('Zion',data=fpout.atm['Zion']); atmg.create_dataset('Zn',data=fpout.atm['Zn']); atmg.create_dataset('Enion',data=fpout.atm['Enion'])
    fle.create_dataset('E',data = fpout.E); fle.create_dataset('mu',data = fpout.mu); fle.create_dataset('z',data = fpout.z)
    fle.create_dataset('esvol',data = fpout.esvol); fle.create_dataset('heatrate',data = fpout.heatrate); fle.create_dataset('momrate',data = fpout.momrate)
    fle.create_dataset('f',data = fpout.f)

def writeparam(fle,outfile,nE,nmu,atmfile,inc_relativity,inc_cc, inc_synchro, inc_magmirror, inc_rc, oneD, reflecttop, reflectbottom, maxiter, tolres, toldiff,
               implicit_theta, mbeam, Zbeam, Ecut, dlt, Eflux, patype, pasigma, resist_fact, Emin, Emax, restart):

  truestr = ".true."
  falsestr = ".false."

  with open(fle,'w') as pfle:
    pfle.write("&control\n")
    pfle.write("nE = "+str(nE) + ",\n")
    pfle.write("nmu = "+str(nmu) + ",\n")
    pfle.write("Emin = "+  "%.17e" % Emin + ",\n")
    pfle.write("Emax = "+ "%.17e" % Emax + ",\n")
    pfle.write("inc_relativity = " + (truestr if inc_relativity else falsestr)  + ",\n")
    pfle.write("inc_CC = " + (truestr if inc_cc else falsestr) + ",\n")
    pfle.write("inc_synchro = " + (truestr if inc_synchro else falsestr) + ",\n")
    pfle.write("inc_magmirror = " + (truestr if inc_magmirror else falsestr) + ",\n")
    pfle.write("inc_RC = "+ (truestr if inc_rc else falsestr) + ",\n")
    pfle.write("oneD = "+ (truestr if oneD else falsestr) + ",\n")
    pfle.write("reflecttop = " +  (truestr if reflecttop else falsestr) + ",\n")
    pfle.write("reflectbottom = "+ (truestr if reflectbottom else falsestr) + ",\n")
    pfle.write("maxiter = "+str(maxiter)+",\n")
    pfle.write("writeoutput = .true.,\n")
    pfle.write("tolres = "+ "%.4e" % tolres +",\n")
    pfle.write("toldiff = " +"%.4e" % toldiff +",\n")
    pfle.write("implicit_theta = "+"%.17e" % implicit_theta+",\n")
    pfle.write("atmfile = '"+atmfile+"',\n")
    pfle.write("outfile = '"+outfile+"',\n")
    pfle.write("mbeam =  "+ "%.17e" % mbeam+",\n")
    pfle.write("Zbeam = "+"%.17e" % Zbeam+",\n")
    pfle.write("Ecut = "+"%.17e" % Ecut +",\n")
    pfle.write("dlt = "+"%.17e" % dlt+",\n")
    pfle.write("eflux= "+"%.17e" % Eflux+",\n")
    pfle.write("patype = "+str(patype)+",\n")
    pfle.write("pasigma = "+"%.17e" % pasigma+",\n")
    pfle.write("resist_fact = "+"%.17e" % resist_fact +",\n")
    pfle.write("restart = "+ (truestr if restart else falsestr) +"\n")
    pfle.write("/\n")

def solver(nE = 100, nmu=60, atmfile='atm.dat', inc_relativity = True, inc_cc = True, inc_synchro=True, inc_magmirror = False, inc_rc = True,
           oneD = False, reflecttop = False, reflectbottom = False, maxiter=100, tolres=1e-3, toldiff = 1e-4, implicit_theta = 1.0, mbeam = 0e0, Zbeam=-1.0,
           Ecut=20.0, dlt=5.0, Eflux=1e11, patype=2, pasigma=.05, resist_fact=1e0, Emin = 1e0, Emax = 0, restart= False, writeout = True, outfile = '',
           nthreads = 0, mpiexec = 'mpiexec'):
   import inspect
   import os
   import numpy as np
   import const as c
   import subprocess
   from random import random

   if (nthreads == 0):
     nthreads = os.cpu_count()
   if (mbeam == 0e0):
     mbeam = c.me
   if (Emax == 0):
     Emax = 3e3*Ecut

   iid = int(random()*1e6) # pick some random label to keep track of which param.cnt and out.dat go together
   paramfile = 'param.'+str(iid)+'.cnt'
   if (outfile == ''):
     outfile = 'out.'+str(iid)+'.h5'
     keepout = False
   else:
     outfile = str(outfile)
     keepout = True

   if (restart != False):
     if (isinstance(restart,fpoutClass)):
#       writefpout(restart,outfile=outfile)
       dorestart = True
       #Setup input parameters using values from restart
       nE = restart.inputparams['nE']; nmu = restart.inputparams['nmu']; atmfile = restart.atmfile; inc_relativity = restart.inputparams['inc_relativity']
       inc_CC = restart.inputparams['inc_CC']; inc_synchro = restart.inputparams['inc_synchro']; inc_magmirror = restart.inputparams['inc_magmirror']
       oneD = restart.inputparams['oneD']; reflecttop = restart.inputparams['reflecttop']; reflectbottom = restart.inputparams['reflectbottom']
       mbeam = restart.inputparams['mbeam']; Zbeam = restart.inputparams['zbeam']; Ecut = restart.inputparams['Ecut']; dlt = restart.inputparams['dlt']
       Eflux = restart.inputparams['eflux']; patype = restart.inputparams['patype']; pasigma = restart.inputparams['pasigma']
       resist_fact = restart.inputparams['resist_fact']; Emin = restart.inputparams['Emin']; Emax = restart.inputparams['Emax']
       if (maxiter is solver.__defaults__[0]): maxiter = restart.inputparams['maxiter']
       if (tolres == 1e-3): tolres = restart.inputparams['tolres']
       if (toldiff == 1e-4): todiff = restart.inputparams['toldiff']
       if (implicit_theta == 1.0): implicit_theta = restart.inputparams['implicit_theta']
       writefpout(restart,outfile)
     else:
       print('Cannot restart from ' + str(type(restart)))
       return
   else:
     dorestart = False

   import readatm
   atm = readatm.readatm(fle = atmfile)
   if (atm == None): return

   writeparam(paramfile,outfile,nE,nmu,atmfile,inc_relativity,inc_cc, inc_synchro, inc_magmirror, inc_rc, oneD, reflecttop, reflectbottom, maxiter, tolres, toldiff,
              implicit_theta, mbeam, Zbeam, Ecut, dlt, Eflux, patype, pasigma, resist_fact, Emin, Emax, dorestart)

   nz = atm['zin'].size
   nthreads = min(nthreads, int(nz/5))

   thisdir = os.path.dirname(inspect.getfile(writeparam))
   mpiexec = str(mpiexec)
   try:
     out = subprocess.run([mpiexec,'-n',str(nthreads),thisdir+'/fp',paramfile],stderr=subprocess.STDOUT)
   except OSError as err:
     print("OS error: {0}".format(err))
     if (os.path.exists(paramfile)): os.remove(paramfile)
     if (os.path.exists(outfile) and not keepout): os.remove(outfile)
     return None
   fpout = fpoutClass()
   readout(outfile,atmfile,fpout)
   os.remove(paramfile)
   if (not keepout): os.remove(outfile)
   return fpout


import sys
import sysconfig
import warnings
import numpy as nu
import ctypes
import ctypes.util
from numpy.ctypeslib import ndpointer
import os
from galpy import potential
from galpy.util import galpyWarning
from galpy.orbit_src.integratePlanarOrbit import _parse_integrator, _parse_tol
#Find and load the library
_lib= None
outerr= None
PY3= sys.version > '3'
if PY3: #pragma: no cover
    _ext_suffix= sysconfig.get_config_var('EXT_SUFFIX')
else:
    _ext_suffix= '.so'
for path in sys.path:
    try:
        _lib = ctypes.CDLL(os.path.join(path,'galpy_integrate_c%s' % _ext_suffix))
    except OSError as e:
        if os.path.exists(os.path.join(path,'galpy_integrate_c%s' % _ext_suffix)): #pragma: no cover
            outerr= e
        _lib = None
    else:
        break
if _lib is None: #pragma: no cover
    if not outerr is None:
        warnings.warn("integrateFullOrbit_c extension module not loaded, because of error '%s' " % outerr,
                      galpyWarning)
    else:
        warnings.warn("integrateFullOrbit_c extension module not loaded, because galpy_integrate_c%s image was not found" % _ext_suffix,
                      galpyWarning)
    _ext_loaded= False
else:
    _ext_loaded= True

def _parse_pot(pot,potforactions=False):
    """Parse the potential so it can be fed to C"""
    #Figure out what's in pot
    if not isinstance(pot,list):
        pot= [pot]
    #Initialize everything
    pot_type= []
    pot_args= []
    npot= len(pot)
    for p in pot:
        if isinstance(p,potential.LogarithmicHaloPotential):
            pot_type.append(0)
            pot_args.extend([p._amp,p._q,p._core2])
        elif isinstance(p,potential.MiyamotoNagaiPotential):
            pot_type.append(5)
            pot_args.extend([p._amp,p._a,p._b])
        elif isinstance(p,potential.PowerSphericalPotential):
            pot_type.append(7)
            pot_args.extend([p._amp,p.alpha])
        elif isinstance(p,potential.HernquistPotential):
            pot_type.append(8)
            pot_args.extend([p._amp,p.a])
        elif isinstance(p,potential.FlattenedNFWPotential):
            pot_type.append(91)
            pot_args.extend([p._amp,p.a,p.q])
        elif isinstance(p,potential.NFWPotential):
            pot_type.append(9)
            pot_args.extend([p._amp,p.a])
        elif isinstance(p,potential.JaffePotential):
            pot_type.append(10)
            pot_args.extend([p._amp,p.a])
        elif isinstance(p,potential.DoubleExponentialDiskPotential):
            pot_type.append(11)
            pot_args.extend([p._amp,p._alpha,p._beta,p._kmaxFac,
                             p._nzeros,p._glorder])
            pot_args.extend([p._glx[ii] for ii in range(p._glorder)])
            pot_args.extend([p._glw[ii] for ii in range(p._glorder)])
            pot_args.extend([p._j0zeros[ii] for ii in range(p._nzeros+1)])
            pot_args.extend([p._dj0zeros[ii] for ii in range(p._nzeros+1)])
            pot_args.extend([p._j1zeros[ii] for ii in range(p._nzeros+1)])
            pot_args.extend([p._dj1zeros[ii] for ii in range(p._nzeros+1)])
            pot_args.extend([p._kp._amp,p._kp.alpha])
        elif isinstance(p,potential.FlattenedPowerPotential):
            pot_type.append(12)
            pot_args.extend([p._amp,p.alpha,p.q2,p.core2])
        elif isinstance(p,potential.interpRZPotential):
            pot_type.append(13)
            pot_args.extend([len(p._rgrid),len(p._zgrid)])
            if p._logR:
                pot_args.extend([p._logrgrid[ii] for ii in range(len(p._rgrid))])
            else:
                pot_args.extend([p._rgrid[ii] for ii in range(len(p._rgrid))])
            pot_args.extend([p._zgrid[ii] for ii in range(len(p._zgrid))])
            if potforactions:
                pot_args.extend([x for x in p._potGrid_splinecoeffs.flatten(order='C')])
            else:
                pot_args.extend([x for x in p._rforceGrid_splinecoeffs.flatten(order='C')])
                pot_args.extend([x for x in p._zforceGrid_splinecoeffs.flatten(order='C')])
            pot_args.extend([p._amp,int(p._logR)])
        elif isinstance(p,potential.IsochronePotential):
            pot_type.append(14)
            pot_args.extend([p._amp,p.b])
        elif isinstance(p,potential.PowerSphericalPotentialwCutoff):
            pot_type.append(15)
            pot_args.extend([p._amp,p.alpha,p.rc])
        elif isinstance(p,potential.MN3ExponentialDiskPotential):
            # Three Miyamoto-Nagai disks
            npot+= 2
            pot_type.extend([5,5,5])
            pot_args.extend([p._amp*p._mn3[0]._amp,
                             p._mn3[0]._a,p._mn3[0]._b,
                             p._amp*p._mn3[1]._amp,
                             p._mn3[1]._a,p._mn3[1]._b,
                             p._amp*p._mn3[2]._amp,
                             p._mn3[2]._a,p._mn3[2]._b])
        elif isinstance(p,potential.KuzminKutuzovStaeckelPotential):
            pot_type.append(16)
            pot_args.extend([p._amp,p._ac,p._Delta])
        elif isinstance(p,potential.PlummerPotential):
            pot_type.append(17)
            pot_args.extend([p._amp,p._b])
        elif isinstance(p,potential.PseudoIsothermalPotential):
            pot_type.append(18)
            pot_args.extend([p._amp,p._a])
    pot_type= nu.array(pot_type,dtype=nu.int32,order='C')
    pot_args= nu.array(pot_args,dtype=nu.float64,order='C')
    return (npot,pot_type,pot_args)

def integrateFullOrbit_c(pot,yo,t,int_method,rtol=None,atol=None,dt=None):
    """
    NAME:
       integrateFullOrbit_c
    PURPOSE:
       C integrate an ode for a FullOrbit
    INPUT:
       pot - Potential or list of such instances
       yo - initial condition [q,p]
       t - set of times at which one wants the result
       int_method= 'leapfrog_c', 'rk4_c', 'rk6_c', 'symplec4_c'
       rtol, atol
       dt= (None) force integrator to use this stepsize (default is to automatically determine one))
    OUTPUT:
       (y,err)
       y : array, shape (len(y0), len(t))
       Array containing the value of y for each desired time in t, \
       with the initial value y0 in the first row.
       err: error message, if not zero: 1 means maximum step reduction happened for adaptive integrators
    HISTORY:
       2011-11-13 - Written - Bovy (IAS)
    """
    rtol, atol= _parse_tol(rtol,atol)
    npot, pot_type, pot_args= _parse_pot(pot)
    int_method_c= _parse_integrator(int_method)
    if dt is None:
        dt= -9999.99

    #Set up result array
    result= nu.empty((len(t),6))
    err= ctypes.c_int(0)

    #Set up the C code
    ndarrayFlags= ('C_CONTIGUOUS','WRITEABLE')
    integrationFunc= _lib.integrateFullOrbit
    integrationFunc.argtypes= [ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=nu.int32,flags=ndarrayFlags),
                               ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.c_double,
                               ctypes.c_double,
                               ctypes.c_double,
                               ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.POINTER(ctypes.c_int),
                               ctypes.c_int]

    #Array requirements, first store old order
    f_cont= [yo.flags['F_CONTIGUOUS'],
             t.flags['F_CONTIGUOUS']]
    yo= nu.require(yo,dtype=nu.float64,requirements=['C','W'])
    t= nu.require(t,dtype=nu.float64,requirements=['C','W'])
    result= nu.require(result,dtype=nu.float64,requirements=['C','W'])

    #Run the C code
    integrationFunc(yo,
                    ctypes.c_int(len(t)),
                    t,
                    ctypes.c_int(npot),
                    pot_type,
                    pot_args,
                    ctypes.c_double(dt),
                    ctypes.c_double(rtol),ctypes.c_double(atol),
                    result,
                    ctypes.byref(err),
                    ctypes.c_int(int_method_c))

    #Reset input arrays
    if f_cont[0]: yo= nu.asfortranarray(yo)
    if f_cont[1]: t= nu.asfortranarray(t)

    return (result,err.value)

def integrateFullOrbit_dxdv_c(pot,yo,dyo,t,int_method,rtol=None,atol=None): #pragma: no cover because not included in v1, uncover when included
    """
    NAME:
       integrateFullOrbit_dxdv_c
    PURPOSE:
       C integrate an ode for a planarOrbit+phase space volume dxdv
    INPUT:
       pot - Potential or list of such instances
       yo - initial condition [q,p]
       dyo - initial condition [dq,dp]
       t - set of times at which one wants the result
       int_method= 'leapfrog_c', 'rk4_c', 'rk6_c', 'symplec4_c'
       rtol, atol
    OUTPUT:
       (y,err)
       y : array, shape (len(y0), len(t))
       Array containing the value of y for each desired time in t, \
       with the initial value y0 in the first row.
       err: error message if not zero, 1: maximum step reduction happened for adaptive integrators
    HISTORY:
       2011-11-13 - Written - Bovy (IAS)
    """
    rtol, atol= _parse_tol(rtol,atol)
    npot, pot_type, pot_args= _parse_pot(pot)
    int_method_c= _parse_integrator(int_method)
    yo= nu.concatenate((yo,dyo))

    #Set up result array
    result= nu.empty((len(t),12))
    err= ctypes.c_int(0)

    #Set up the C code
    ndarrayFlags= ('C_CONTIGUOUS','WRITEABLE')
    integrationFunc= _lib.integrateFullOrbit_dxdv
    integrationFunc.argtypes= [ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.c_int,
                               ndpointer(dtype=nu.int32,flags=ndarrayFlags),
                               ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.c_double,
                               ctypes.c_double,
                               ndpointer(dtype=nu.float64,flags=ndarrayFlags),
                               ctypes.POINTER(ctypes.c_int),
                               ctypes.c_int]

    #Array requirements, first store old order
    f_cont= [yo.flags['F_CONTIGUOUS'],
             t.flags['F_CONTIGUOUS']]
    yo= nu.require(yo,dtype=nu.float64,requirements=['C','W'])
    t= nu.require(t,dtype=nu.float64,requirements=['C','W'])
    result= nu.require(result,dtype=nu.float64,requirements=['C','W'])

    #Run the C code
    integrationFunc(yo,
                    ctypes.c_int(len(t)),
                    t,
                    ctypes.c_int(npot),
                    pot_type,
                    pot_args,
                    ctypes.c_double(rtol),ctypes.c_double(atol),
                    result,
                    ctypes.byref(err),
                    ctypes.c_int(int_method_c))

    #Reset input arrays
    if f_cont[0]: yo= nu.asfortranarray(yo)
    if f_cont[1]: t= nu.asfortranarray(t)

    return (result,err.value)