"""
This code provides obejcts to generate counts and model counts images with
two dimensional projections. The code interfaces with the ROIAnalysis
object for making the calculations, skymaps.SkyImage object for storing
the data, and the image.ZEA object for plotting.  The high level object
roi_plotting.ROIDisplay can use to access these objects form a high
level plotting interface.

$Header: /nfs/slac/g/glast/ground/cvs/pointlike/python/uw/like/roi_image.py,v 1.42 2017/08/23 16:23:42 zimmer Exp $

author: Joshua Lande
"""
from skymaps import SkyImage,SkyDir,PythonUtilities,Band,WeightedSkyDirList
import numpy as np
import astropy.io.fits as pyfits
import scipy
import scipy.ndimage

from . pypsf import PretendBand
from . roi_diffuse import ROIDiffuseModel_OTF
from . roi_extended import ROIExtendedModel,ROIExtendedModelAnalytic
from . SpatialModels import RadiallySymmetricModel
from . pointspec_helpers import get_default_diffuse_mapper
from . roi_tsmap import TSCalc,TSCalcPySkyFunction
from uw.utilities import keyword_options
from uw.utilities.fitstools import get_fields
from uw.utilities.decorators import memoize
from pypsf import PsfOverlap
import collections
from abc import abstractmethod
import numbers

class ROIImage(object):
    """ This object is suitable for creating a SkyImage object
        and filling it with some physically meaningful
        quantity gotten from an ROIAnalysis object.
        The acutal work is done by subclasses. """

    defaults = (
        ('size',     2,     'size of image in degrees'),
        ('pixelsize',0.1,   'size, in degrees, of pixels'),
        ('galactic', False, 'galactic or equatorial coordinates'),
        ('proj',     'ZEA', 'projection name: can change if desired'),
        ('center',    None, 'Center of image. If None, use roi center.'),
        ('conv_type',   -1, 'Conversion type'),
    )

    @keyword_options.decorate(defaults)
    def __init__(self,roi,**kwargs):
        """ Note, unlike ZEA, can support non-square images. To specify a nonsquare
            image, set the size parameter to a lenght two tuple:

                size=(10,5) # dx=10 degrees, dy=5 degrees.
        """
        keyword_options.process(self, kwargs)

        if self.size < self.pixelsize:
            raise Exception("Can only create images with >=1 pixel in them.")

        self.roi = roi

        self.selected_bands = tuple(self.roi.bands if self.conv_type < 0 else \
            [ band for band in self.roi.bands if band.ct == self.conv_type ])

        # by default, use get energy range and image center from roi.
        if self.center is None: self.center=self.roi.roi_dir

        # set up, then create a SkyImage object to perform the projection
        # to a grid and manage an image
        if not isinstance(self.size,collections.Iterable):

            # make sure size and pixelsize are commensurate (helpful for
            # various downsampling code later).
            self.size = int(self.size/self.pixelsize + 0.01)*self.pixelsize
            self.skyimage = SkyImage(self.center, '', self.pixelsize,
                                     self.size, 1, self.proj, self.galactic, False)
        else:
            self.skyimage = SkyImage(self.center, '', self.pixelsize,
                                     float(self.size[0]), 1, self.proj, self.galactic, False, float(self.size[1]))

        self.fill()

        self.nx, self.ny = self.skyimage.naxis1(), self.skyimage.naxis2()
        self.image=ROIImage.skyimage2numpy(self.skyimage)

    @staticmethod
    def skyimage2numpy(skyimage):
        nx, ny = skyimage.naxis1(), skyimage.naxis2()
        image=np.array(skyimage.image()).reshape((ny, nx))
        return image

    @abstractmethod
    def fill(self): pass

    def get_ZEA(self,axes=None,nticks=None):
        """ axes and nticks can be created by this object's constructor, but are
            more logically specified here. If they are not specified, get values from
            initial object creation. """
        # get out of the object all parameters which should be passed to ZEA.

        if hasattr(self.size,'__iter__'):
            raise Exception("Can only create ZEA object for square objects.")

        zea_dict = dict((d[0],self.__dict__[d[0]]) for d in ZEA.defaults if hasattr(d,'__iter__') and \
                hasattr(self,d[0]))
        if axes is not None: zea_dict['axes']=axes
        if nticks is not None: zea_dict['nticks']=nticks

        from uw.utilities.image import ZEA
        zea=ZEA(self.center,**zea_dict)
        zea.skyimage = self.skyimage
        # recalculate, in case the sky image has changed
        zea.image = ROIImage.skyimage2numpy(self.skyimage)

        # The old one gets removed by python's garbage collector (when zea.skyimage is replaced).
        zea.projector = zea.skyimage.projector()

        zea.vmin,zea.vmax = zea.skyimage.minimum(), zea.skyimage.maximum()
        return zea

    def get_pyfits(self):
        """ Create and return a pyfits object that corresponds to the ROIImage object.
            The fits file created is supposed to be consistent with the internal
            representation that SkyImage/SkyProj uses. """

        if self.galactic:
            ctype1="GLON-%s" % self.proj
            ctype2="GLAT-%s" % self.proj
            # for some reason, SkyDir(0,0,SkyDir.GALACTIC).l() = 360
            crval1,crval2=self.center.l() % 360,self.center.b()
        else:
            ctype1="RA-%s" % self.proj
            ctype2="DEC-%s" % self.proj
            crval1,crval2=self.center.ra(),self.center.dec()

        cdelt1,cdelt2=-self.pixelsize,self.pixelsize

        # from SkyImage.cxx like 92:
        #   "center pixel; WCS convention is that center of a pixel is a half-integer"
        crpix1,crpix2=(self.skyimage.naxis1()+1)/2.0,(self.skyimage.naxis2()+1)/2.0

        values = [
            ["TELESCOP", "GLAST"],
            ["INSTRUME", "LAT"],
            ["DATE-OBS", ""],
            ["DATE-END", ""],
            ["EQUINOX", 2000.0, "Equinox of RA & DEC specifications"],
            ["CTYPE1", ctype1, "[RA|GLON]---%%%, %%% represents the projection method such as AIT"],
            ["CRPIX1", crpix1, "Reference pixel"],
            ["CRVAL1", crval1, "RA or GLON at the reference pixel"],
            ["CDELT1", cdelt1, "X-axis incr per pixel of physical coord at position of ref pixel(deg)"],
            ["CTYPE2", ctype2, "[DEC|GLAT]---%%%, %%% represents the projection method such as AIT"],
            ["CRPIX2", crpix2, "Reference pixel"],
            ["CRVAL2", crval2, "DEC or GLAT at the reference pixel"],
            ["CDELT2", cdelt2, "Y-axis incr per pixel of physical coord at position of ref pixel(deg)"],
            ["CROTA2",  0, "Image rotation (deg)"],
        ]
        for i in values:
            if len(i)>2 and len(i[2])>47: i[2]=i[2][0:47]

        cards = [ pyfits.Card(*i) for i in values]

        header=pyfits.Header(cards=cards)

        hdu=pyfits.PrimaryHDU(data=self.image, header=header)
        fits = pyfits.HDUList([hdu])

        return fits

class ROITSMapImage(ROIImage):
    """ Subclass of ROIImage representing a residual TS map. """

    defaults = ROIImage.defaults + TSCalc.defaults

    @keyword_options.decorate(defaults)
    def __init__(self,*args,**kwargs):
        super(ROITSMapImage,self).__init__(*args,**kwargs)

    def fill(self):

        tscalc = TSCalc(self.roi,**keyword_options.defaults_to_kwargs(self,TSCalc))
        temp=TSCalcPySkyFunction(tscalc)
        self.skyimage.fill(temp.get_pyskyfun())

class CountsImage(ROIImage):
    """ This ROIImage subclass fills the sky image with the observed Fermi counts. """

    @staticmethod
    @memoize
    def process_filedata(roi,selected_bands):
        """ The radius parameter will apply a radius cut.
            Assume that all bands are contiguous (hope this is always true!) """

        radius=roi.sa.maxROI

        emin = min(b.emin for b in selected_bands)
        emax = max(b.emax for b in selected_bands)

        ft1files=roi.sa.pixeldata.ft1files

        cuts = ['ENERGY > %s'% emin,
                'ENERGY < %s'% emax,
                'ZENITH_ANGLE < %s' % roi.sa.pixeldata.zenithcut,
                'THETA < %s' % roi.sa.pixeldata.thetacut,
                'EVENT_CLASS >= %s' % roi.sa.pixeldata.event_class]

        data = get_fields(ft1files,['RA','DEC','TIME','ENERGY','CONVERSION_TYPE'],cuts)
        # convert into skydirs
        skydirs = [ SkyDir(float(data['RA'][i]),float(data['DEC'][i])) for i in xrange(len(data['RA']))]

        # apply the same gti cut used to read in the initial WSDL.
        gti=roi.sa.pixeldata.gti
        good_dirs = []

        front_bins = [b for b in selected_bands if b.ct == 0]
        front_emin = min(b.emin for b in front_bins) if len(front_bins)>0 else None
        front_emax = max(b.emax for b in front_bins) if len(front_bins)>0 else None

        back_bins = [b for b in selected_bands if b.ct == 1]
        back_emin = min(b.emin for b in back_bins) if len(back_bins)>0 else None
        back_emax = max(b.emax for b in back_bins) if len(back_bins)>0 else None

        good_photons = []
        for skydir,time,energy,ct in zip(skydirs,data['TIME'],data['ENERGY'],data['CONVERSION_TYPE']):
            if gti.accept(time) and np.degrees(skydir.difference(roi.roi_dir)) < radius:
                if ct == 0 and \
                   front_emin is not None and front_emax is not None and \
                   energy > front_emin and energy < front_emax:
                    good_photons.append(skydir)
                if ct == 1 and \
                   back_emin is not None and back_emax is not None and \
                   energy > back_emin and energy < back_emax:
                    good_photons.append(skydir)

        return good_photons

    def fill(self):
        dirs = CountsImage.process_filedata(self.roi,self.selected_bands)

        for photon_dir in dirs:
            self.skyimage.addPoint(photon_dir)

class ModelImage(ROIImage):
    """ This ROIImage subclass fills the sky image with the model
        predicted counts for a fermi sky model described by an ROIAnalysis
        object.

        This code is forced to deal with the fact that model intensity
        can vary significantly across a spatial pixel. The rest of the
        pointlike code can avoid this whole issue by scaling the healpix
        pixel size with the PSF to ensure that pixels are always small
        compared to the instrument's intrisic resolution. But since
        model predicted counts maps can be generated of arbitary pixel size,
        this issue must directly be dealt with.

        The solution that this code uses to deal with this issue is to
        simply sample from a grid finer by an integer number of pixels
        in each dimensions.  After calculating the model predictions,
        the nearby blocks of model predictions are averaged to downsample
        to create the model predictions.  This formulation assumes that
        each of the subpixels has the same solid angle and so it only
        suitable for relativly small images where pixels have equal area.
        For that reason, it is advised to use the ZEA projection.

        For point and extended sources, the characteristic scale with
        which the convolution must be small compared to is the PSF. So
        the formula for determining the factor is

        factor = ceil(pixelsize/r10)

        Where pxielsize is the plotting pixel size and r10 is the 10%
        containment radius of the PSF.

        For background sources, the characteristic scale is not the PSF
        but the convolution grid pixelsize. So the formula for determining
        the factor is instead

        factor = ceil(pixelsize/(conv_pixelsize/4))

        Where conv_pixelsize is the size of the convolution grid's pixels.

        For background sources, this algorithm is generally efficiency
        since we except the background to vary on this smaller scale all
        across the image.  But for point and (small) extended sources,
        this algorithm is generally very poor because it requires
        calculating the PSF (or PDF) at many points where the value
        is very close to 0. A better algorithm would be an adaptive
        quadrature integration algorithm which evaluate the integral in
        each pixel, then did a more accurate integral and iterated until
        the integral converged. This would avoid having to evaluate the
        model predictions for a source very finely far from the source.
        On the other hand, adding this feature (presumably to C++
        for optimization) would be very costly, and this code runs
        fast enough...

        """

    defaults = ROIImage.defaults + (
            ('override_point_sources', None, """ If either is specified, use override_point_sources these
                                                 and override_diffuse_sources to generate the image instead
                                                 of the sources in the ROI."""),
            ('override_diffuse_sources', None, 'Same as override_point_sources'),
    )

    @keyword_options.decorate(defaults)
    def __init__(self,*args,**kwargs):
        if kwargs.has_key('proj') and kwargs['proj'] != 'ZEA':
            print "Warning, it is strongly advised to use the 'ZEA projection when creating model counts maps."

        super(ModelImage,self).__init__(*args,**kwargs)

    def fill(self):
        self.wsdl = self.skyimage.get_wsdl()

        self.solid_angle = np.radians(self.pixelsize)**2

        model_counts = np.zeros(len(self.wsdl),dtype=float)

        model_counts += self.all_point_source_counts()
        model_counts += self.all_diffuse_sources_counts()
        model_counts *= self.roi.phase_factor # don't forget about the phase factor!
        #NB -- this will need to be fixed if want to account for bracketing IRFs

        PythonUtilities.set_wsdl_weights(model_counts,self.wsdl)

        self.skyimage.set_wsdl(self.wsdl)

    @staticmethod
    def downsample(myarr,factor):
        """
        Code taken from http://code.google.com/p/agpy/source/browse/trunk/agpy/downsample.py

        Downsample a 1D or 2D array by averaging over *factor* pixels in each axis.
        Crops upper edge if the shape is not a multiple of factor.

        This code is pure numpy and should be fast.
        """
        assert isinstance(factor,numbers.Integral)
        assert len(myarr.shape) <= 2

        if len(myarr.shape) == 1:
            xs = myarr.shape[0]
            assert xs % factor == 0
            dsarr = np.concatenate([[myarr[i::factor]]
                                   for i in range(factor)]).mean(axis=0)
            return dsarr

        elif len(myarr.shape) == 2:
            xs,ys = myarr.shape
            assert xs % factor == 0 and ys % factor == 0
            dsarr = np.concatenate([[myarr[i::factor,j::factor]
                for i in range(factor)]
                for j in range(factor)]).mean(axis=0)
            return dsarr

    def bigger_wsdl(self,band,compare=None):
        """ Want to sample on a grid that is comparable in size (or
            smaller than) 10% of the psf to ensure we get a reasonable
            sampling of the grid. """
        if compare is None:
            r10=band.psf.inverse_integral_on_axis(0.10)
            compare=r10

        self.factor = int(np.ceil(self.pixelsize/compare))

        if self.factor == 1:
            return self.wsdl
        else:
            # hold onto this thing since it is needed by downsample_model
            if not hasattr(self.size,'__iter__'):
                self.fine_skyimage = SkyImage(self.center, '', float(self.pixelsize)/self.factor,
                                         self.size, 1, self.proj, self.galactic, False)
            else:
                self.fine_skyimage = SkyImage(self.center, '', float(self.pixelsize)/self.factor,
                                         float(self.size[0]), 1, self.proj, self.galactic, False, float(self.size[1]))

            wsdl = self.fine_skyimage.get_wsdl()
            return wsdl

    def downsample_model(self,rvals):
        if self.factor==1:
            return rvals
        else:
            rvals = rvals.reshape((self.fine_skyimage.naxis2(), self.fine_skyimage.naxis1()))
            rvals = ModelImage.downsample(rvals,self.factor).flatten()
            return rvals

    @staticmethod
    def get_point_sources(roi,override_point_sources,override_diffuse_sources):
        if override_point_sources is None and override_diffuse_sources is None:
            return roi.psm.point_sources
        if override_point_sources is None:
            return []
        elif not isinstance(override_point_sources,collections.Iterable):
            return [override_point_sources]
        else:
            return override_point_sources

    def all_point_source_counts(self):
        """ Calculate the point source contributions. """
        point_sources=ModelImage.get_point_sources(self.roi,self.override_point_sources,self.override_diffuse_sources)
        if len(point_sources)==0: return 0

        point_counts = np.zeros(len(self.wsdl),dtype=float)

        for band in self.selected_bands:
            cpsf = band.psf.cpsf

            # generate a list of skydirs on a finer grid.
            wsdl = self.bigger_wsdl(band)

            rvals  = np.empty(len(wsdl),dtype=float)

            for nps,ps in enumerate(point_sources):
                # evaluate the PSF at the center of each pixel
                cpsf.wsdl_val(rvals,ps.skydir,wsdl)

                # average the finer grid back to original resolution.
                temp = self.downsample_model(rvals)

                temp *= self.solid_angle #multiply by pixel solid angle
                temp *= band.expected(ps.model) # scale by total expected counts
                point_counts += temp

        return point_counts

    def extended_source_counts(self,extended_model):

        rd = self.roi.roi_dir

        es = extended_model.extended_source
        sm = es.smodel

        extended_counts = np.zeros(len(self.wsdl),dtype=float)

        for band in self.selected_bands:
            extended_model.set_state(band)
            exposure=band.exp.value

            er = exposure(es.spatial_model.center,extended_model.current_energy)/exposure(rd,extended_model.current_energy)
            es_counts = band.expected(sm)*er

            wsdl = self.bigger_wsdl(band)

            es_pix_counts = extended_model._pix_value(wsdl)*self.solid_angle

            es_pix_counts= self.downsample_model(es_pix_counts)

            bg_pix_counts = es_pix_counts * es_counts

            extended_counts += bg_pix_counts

        return extended_counts

    def otf_source_counts(self,bg):
        roi=self.roi

        mo=bg.smodel

        background_counts = np.zeros(len(self.wsdl),dtype=float)

        for band in self.selected_bands:

            ns,bg_points,bg_vector = ROIDiffuseModel_OTF.sub_energy_binning(band,bg.nsimps)

            pi_evals  = np.empty([len(self.wsdl),ns + 1])

            wsdl = self.bigger_wsdl(band,compare=bg.pixelsize/4.0)

            for ne,e in enumerate(bg_points):
                bg.set_state(e,band.ct,band)
                temp = self.downsample_model(bg._pix_value(wsdl))
                pi_evals[:,ne] = temp

            pi_evals *= (self.solid_angle * bg_vector)
            mo_evals  = mo(bg_points)
            pi_counts = (pi_evals * mo_evals).sum(axis=1)

            background_counts += pi_counts

        return background_counts

    def diffuse_source_counts(self,bg):
        if isinstance(bg,ROIDiffuseModel_OTF):
            return self.otf_source_counts(bg)
        elif isinstance(bg,ROIExtendedModel):
            return self.extended_source_counts(bg)
        else:
            raise Exception("Unable to calculate model predictions for diffuse source %s", bg.name)

    @staticmethod
    def get_diffuse_sources(roi,override_point_sources,override_diffuse_sources):

        if override_point_sources is None and override_diffuse_sources is None:
            return roi.dsm.bgmodels
        else:
            mapper=get_default_diffuse_mapper(roi.sa,roi.roi_dir,roi.quiet)
            if override_diffuse_sources is None:
                return []
            elif not isinstance(override_diffuse_sources,collections.Iterable):
                return [mapper(override_diffuse_sources)]
            else:
                return [mapper(ds) for ds in override_diffuse_sources]

    def all_diffuse_sources_counts(self):
        """ Calculate the diffuse source contributions. """

        bgmodels=ModelImage.get_diffuse_sources(self.roi,self.override_point_sources,self.override_diffuse_sources)

        return sum(self.diffuse_source_counts(bg)
                   for bg in bgmodels)

class ResidualImage(ROIImage):
    """ Has the same arguments at both ModelImage and Counts image but
        is a ROIImage object representing the difference between the
        source Counts and the ROI Model. """

    # get unique items in the Model + Counts defaults
    defaults = tuple(set(ModelImage.defaults).union(CountsImage.defaults))

    def __init__(self,roi,**kwargs):
        super(ResidualImage,self).__init__(roi,**kwargs)

        self.model=ModelImage(self.roi,**keyword_options.defaults_to_kwargs(self,ModelImage))
        self.counts=CountsImage(self.roi,**keyword_options.defaults_to_kwargs(self,CountsImage))

        self.image = self.counts.image - self.model.image
        SmoothedImage.add_to_skyimage(self.skyimage,self.image)

class RadialImage(object):
    """ This object is similar to ROIImage but performs a radial
        integral around around a given direction in the sky. """

    defaults = (
            ('center',       None,            'Center of image'),
            ('size',            2, 'Size of image (in degrees)'),
            ('pixelsize', 0.00625, """ size of each image pixel. This is a little misleading because the
                                       size of each pixel varies, but regardless this is used to determine
                                       the total number of pixels with npix=size/pixelsize """),
            ('npix',         None, """ If specified, use this value instead of pixelsize. """),
            ('conv_type',    None,            'Conversion type'),
    )

    @keyword_options.decorate(defaults)
    def __init__(self,roi,**kwargs):
        keyword_options.process(self, kwargs)

        if self.npix is None:
            self.npix = float(self.size)/self.pixelsize

        self.roi = roi

        self.selected_bands = tuple(self.roi.bands if self.conv_type < 0 else \
            [ band for band in self.roi.bands if band.ct == self.conv_type ])

        # by default, use get energy range and image center from roi.
        if self.center is None: self.center=self.roi.roi_dir

        # bins in theta^2
        self.bin_edges_deg = np.linspace(0.0,self.size**2,self.npix+1)
        self.bin_centers_deg = (self.bin_edges_deg[1:] + self.bin_edges_deg[:-1])/2.0

        # two factors of radians b/c theta^2
        self.bin_edges_rad = np.radians(np.radians(self.bin_edges_deg))
        self.bin_centers_rad = np.radians(np.radians(self.bin_centers_deg))

        # the lower and upper agle for each bin.
        self.theta_pairs_rad = zip(np.sqrt(self.bin_edges_rad[:-1]),
                                   np.sqrt(self.bin_edges_rad[1:]))

        self.fill()

    @abstractmethod
    def fill(self):
        """ This should fill up self.image appropriatly."""
        pass

class RadialCounts(RadialImage):
    """ This subclass of RadialImage calculates the counts within
        each radial bin. """

    defaults = RadialImage.defaults

    def fill(self):
        dirs = CountsImage.process_filedata(self.roi,self.selected_bands)
        diffs = [self.center.difference(i) for i in dirs]
        self.image=np.histogram(diffs,bins=np.sqrt(self.bin_edges_rad))[0]

class RadialSource(RadialImage):
    """ Subclass of RadialImage where returns the
        model predicted counts integrated radially
        for an extended source. Unlike RadialModel,
        the PSF is not convolved with the RadialSource
        (useful for comparing the extended source shape
        with the PSF). """

    defaults = RadialImage.defaults + (
            ('extended_source', None, 'A spatial moel'),
    )

    def fill(self):
        if self.extended_source is None:
            raise Exception("RadialSource must be given a spatial_model.")

        es=self.extended_source
        sm = es.model

        if not isinstance(es.spatial_model,RadiallySymmetricModel):
            raise Exception("spatial_model must be an instance of RadiallySymmetricModel.")

        total_counts = sum(band.expected(sm) for band in self.roi.bands)
        solid_angle = RadialModel.solid_angle_cone(np.radians(self.size))/self.npix

        fraction = es.spatial_model.at_r(np.sqrt(self.bin_centers_rad))
        fraction*=solid_angle

        self.image = total_counts*fraction

class RadialModel(RadialImage):

    defaults = RadialImage.defaults + (
            ('override_point_sources', None, """ If either is specified, use override_point_sources these
                                                 and override_diffuse_sources to generate the image instead
                                                 of the sources in the ROI."""),
            ('override_diffuse_sources', None, 'Same as override_point_sources'),
    )

    def fill(self):

        # Create fake bands just big enough to enclose the radial model.
        # This will speed up the convolution.
        # The fake band will cause the diffuse sources be be convolved
        # in a smaller area, which will thus make them more accurate,
        # and create better plots.
        self.smaller_bands = []
        for band in self.selected_bands:
            rad=self.center.difference(self.roi.roi_dir) + np.radians(self.size)
            pb = PretendBand(band.e,band.ct, psf=band.psf, radius_in_rad=rad,
                             sd=self.center, emin=band.emin, emax=band.emax)
            self.smaller_bands.append(pb)

        self.image = np.zeros_like(self.bin_centers_rad)
        self.image += self.all_point_source_counts()
        self.image += self.all_diffuse_sources_counts()
        self.image *= self.roi.phase_factor # don't forget about the phase factor!

    def all_point_source_counts(self):
        """ Calculate the point source contributions. """
        point_sources=ModelImage.get_point_sources(self.roi,self.override_point_sources,self.override_diffuse_sources)
        if len(point_sources)==0: return 0

        point_counts = np.zeros_like(self.bin_centers_rad)

        overlap = PsfOverlap()

        for i,(theta_min,theta_max) in enumerate(self.theta_pairs_rad):

            model_counts=0

            for j,ps in enumerate(point_sources):
                for band in self.selected_bands:

                    # this code requires a redundant call to overlap. Improve if time.
                    # Note that ragged_edge is not appropriate here because our counts
                    # are always sumed in the exact range.
                    fraction=overlap(band,self.center,ps.skydir,radius_in_rad=theta_max, ragged_edge=np.inf) - \
                             overlap(band,self.center,ps.skydir,radius_in_rad=theta_min, ragged_edge=np.inf)

                    model_counts += band.expected(ps.model)*fraction

            point_counts[i] = model_counts

        return point_counts

    def extended_source_counts(self,extended_model):
        if type(extended_model) not in [ROIExtendedModel,ROIExtendedModelAnalytic]:
            raise Exception("Unknown extended model.")

        roi=self.roi
        sm = extended_model.extended_source.model

        extended_counts = np.zeros_like(self.bin_centers_rad)

        for band,smaller_band in zip(self.selected_bands,self.smaller_bands):

            extended_model.set_state(smaller_band)

            if type(extended_model) == ROIExtendedModel:

                nside = RadialModel.get_nside(self.size,self.npix)

                temp_band = Band(nside)
                wsdl = WeightedSkyDirList(temp_band,self.center,np.radians(self.size),True)
                vals=extended_model._pix_value(wsdl)

                rvals=np.empty(len(wsdl),dtype=float)
                PythonUtilities.arclength(rvals,wsdl,self.center)

                # get average value in each ring by averaging values.
                fraction = np.histogram(rvals,weights=vals,bins=np.sqrt(self.bin_edges_rad))[0]/\
                           np.histogram(rvals,bins=np.sqrt(self.bin_edges_rad))[0]

                # multiply intensities by solid angle in ring
                fraction *= RadialModel.solid_angle_cone(np.radians(self.size))/self.npix

            elif type(extended_model) == ROIExtendedModelAnalytic:

                fraction = np.empty_like(self.bin_centers_rad)

                for i,(theta_min,theta_max) in enumerate(self.theta_pairs_rad):

                    fraction[i]=extended_model._overlaps(self.center,band,theta_max) - \
                             extended_model._overlaps(self.center,band,theta_min)

            # total counts from source * fraction of PDF in ring = model predictions in each ring.
            extended_counts += band.expected(sm)*fraction

        return extended_counts

    @staticmethod
    def solid_angle_cone(radius_in_radians):
        return 2*np.pi*(1-np.cos(radius_in_radians))

    @staticmethod
    def get_nside(size,npix,num_points_per_ring=200):
        """ Solid angle of each healpix pixel is 4pi/(12*ns^2)
            Solid angel of each ring is pi*(size)^2/npix
            Want size of each ring > num_points_per_ring*size of each healpix (so that
            we get 20 pixels to sample each ring).

            Solving for ns, the size of the healpixels required, we get the required formula
        """
        total_solid_angle=4*np.pi
        image_solid_angle=RadialModel.solid_angle_cone(np.radians(size))

        solid_angle_per_ring=image_solid_angle/npix

        nside = int(np.ceil(np.sqrt(num_points_per_ring*total_solid_angle/(12*solid_angle_per_ring))))
        return nside

    def otf_source_counts(self,bg):

        roi=self.roi

        mo=bg.smodel

        background_counts = np.zeros_like(self.bin_centers_rad)

        for band,smaller_band in zip(self.selected_bands,self.smaller_bands):

            ns,bg_points,bg_vector = ROIDiffuseModel_OTF.sub_energy_binning(band,bg.nsimps)

            nside = RadialModel.get_nside(self.size,self.npix)

            temp_band = Band(nside)
            wsdl = WeightedSkyDirList(temp_band,self.center,np.radians(self.size),True)

            ap_evals = np.empty([len(self.bin_centers_rad),len(bg_points)])

            for ne,e in enumerate(bg_points):

                bg.set_state(e,band.ct,smaller_band)

                rvals=np.empty(len(wsdl),dtype=float)
                PythonUtilities.arclength(rvals,wsdl,self.center)
                vals=bg._pix_value(wsdl)

                # get average value in each ring by averaging values.
                ap_evals[:,ne] = np.histogram(rvals,weights=vals,bins=np.sqrt(self.bin_edges_rad))[0]/\
                                 np.histogram(rvals,bins=np.sqrt(self.bin_edges_rad))[0]

            # multiply intensities by solid angle in ring
            ap_evals *= RadialModel.solid_angle_cone(np.radians(self.size))/self.npix

            ap_evals          *= bg_vector
            mo_evals           = mo(bg_points)
            background_counts += (ap_evals * mo_evals).sum(axis=1)

        return background_counts

    def diffuse_source_counts(self,bg):
        if isinstance(bg,ROIDiffuseModel_OTF):
            return self.otf_source_counts(bg)
        elif isinstance(bg,ROIExtendedModel):
            return self.extended_source_counts(bg)
        else:
            raise Exception("Unable to calculate model predictions for diffuse source %s", bg.name)

    def all_diffuse_sources_counts(self):
        """ Calculate the diffuse source contributions. """
        bgmodels=ModelImage.get_diffuse_sources(self.roi,self.override_point_sources,self.override_diffuse_sources)

        return sum(self.diffuse_source_counts(bg)
                   for bg in bgmodels)

class SmoothedImage(ROIImage):

    """ Represnts a ROIImage objects that is smoothed. This
        is an abstract base class, but subclasses represent
        smoothed versions of particular ROIImage objects.

        SmoothedCounts -> CountsImage
        SmoothedModel -> ModelImage
        SmoothedResidual -> ResidualImage

        Note, after smoothing, the image is normalized to represent

    """

    smoothed_options = (
        ('kerneltype',      'tophat', 'Type of kernel to use'),
        ('kernel_rad',          0.25, 'Sum counts within radius degrees.'),
        ('per_solid_angle',    False, 'If true, after smoothing divide by solid angle (counts/[deg]^2'),
    )

    defaults = ROIImage.defaults + smoothed_options

    @staticmethod
    def convolve_array(image,width):
        """ Summ all this pixels within a given pixel width

            code taken from

            http://code.google.com/p/agpy/source/browse/trunk/agpy/convolve.py """

        return out

    @staticmethod
    def add_to_skyimage(skyimage,image):
        """ Take in a two dimensional numpy array representing the
            fits data and put it into the skyimage. """
        temp=image.reshape(image.shape[0]*image.shape[1])

        wsdl = skyimage.get_wsdl()
        PythonUtilities.set_wsdl_weights(temp,wsdl)
        skyimage.set_wsdl(wsdl)

    @staticmethod
    def convolve_skyimage(skyimage,kernel):
        """ Take in a skyimage object and sum each pixel with all the
            pixels within a given width. """
        image=ROIImage.skyimage2numpy(skyimage)

        convolved=scipy.ndimage.filters.convolve(image, kernel)

        SmoothedImage.add_to_skyimage(skyimage,convolved)

        return skyimage

    @staticmethod
    def get_kernel(kerneltype,kernel_rad,pixelsize):
        """ returns the kernel. Code inspired by
                http://code.google.com/p/agpy/source/browse/trunk/agpy/convolve.py
        """

        width = int(np.ceil(kernel_rad/pixelsize))

        if kerneltype == 'tophat':
            kernelsize=4*width

            kernel = np.zeros([kernelsize,kernelsize],dtype=float)

            xx,yy = np.indices(kernel.shape)
            rr = np.sqrt((xx-kernel.shape[0]/2.)**2+(yy-kernel.shape[1]/2.)**2)
            kernel[rr<=width] = 1

        elif kerneltype == 'gaussian':
            kernelsize = 8*width

            kernel = np.zeros([kernelsize,kernelsize],dtype=float)
            xx,yy = np.indices(kernel.shape)
            rr = np.sqrt((xx-kernel.shape[0]/2.)**2+(yy-kernel.shape[1]/2.)**2)
            kernel = np.exp(-(rr**2)/(2*width**2))

            # Gaussian kernels should be normalized so they don't change overall normalization
            kernel /= kernel.sum()

        else:
            raise Exception("...")

        return kernel

    @keyword_options.decorate(defaults)
    def __init__(self,*args,**kwargs):

        super(SmoothedImage,self).__init__(*args,**kwargs)

        self.kernel=self.get_kernel(self.kerneltype,self.kernel_rad,self.pixelsize)

        self.kernelsize=self.kernel.shape[0]

        # Make an image, bigger then the desired one, by twice
        # the smooth radius in each direction.
        self.pass_dict=keyword_options.defaults_to_kwargs(self,self.object)
        self.pass_dict['size']=self.size+self.kernelsize*self.pixelsize
        self.smoothed=self.object(self.roi,**self.pass_dict)

        self.smoothed.skyimage=SmoothedImage.convolve_skyimage(self.smoothed.skyimage,self.kernel)
        self.smoothed.image=ROIImage.skyimage2numpy(self.smoothed.skyimage)

        # now, shrink down smoothed image and replace the current skyimage with it

        self.image = self.smoothed.image[self.kernelsize/2:-self.kernelsize/2,self.kernelsize/2:-self.kernelsize/2]

        if self.per_solid_angle:
            # convert from counts to counts per square degree
            self.image = self.image/self.pixelsize**2

        SmoothedImage.add_to_skyimage(self.skyimage,self.image)

class SmoothedCounts(SmoothedImage):

    defaults = CountsImage.defaults + SmoothedImage.smoothed_options

    object=CountsImage

class SmoothedModel(SmoothedImage):

    defaults = ModelImage.defaults + SmoothedImage.smoothed_options

    object=ModelImage

class SmoothedResidual(SmoothedImage):

    defaults = ResidualImage.defaults + SmoothedImage.smoothed_options

    object=ResidualImage

import numpy,import scipy,import astropy


"""
This module provies "UTime" support.
"""
from astropy.time.formats import TimeFromEpoch, erfa

__all__ = ['TimeUTime']

class TimeUTime(TimeFromEpoch):
    """
    Seconds from 1979-01-01 00:00:00 UTC.

    Same as Unix time but this starts 9 years later.
    This time format is included for historical reasons.
    Some people in solar physics prefer using this epoch.

    Examples
    --------
    >>> from astropy.time import Time
    >>> t = Time('2000-01-01T13:53:23')
    >>> print(t.utime)
    662738003.0
    >>> t2 = Time('1979-01-01T00:00:00')
    >>> print(t2.utime)
    0.0
    """
    name = 'utime'
    unit = 1.0 / erfa.DAYSEC  # in days (1 day == 86400 seconds)
    epoch_val = '1979-01-01 00:00:00'
    epoch_val2 = None
    epoch_scale = 'utc'  # Scale for epoch_val class attribute
    epoch_format = 'iso'  # Format for epoch_val class attribute


"""Group of functions for physiological processing."""
from re import findall
import logging

import numpy as np
from itertools import product
from scipy.stats import zscore

__all__ = ('rereferencing', 'bipolarization', 'commonaverage',
           'tal2mni', 'mni2tal')

###############################################################################
###############################################################################
#                               RE-REFERENCING
###############################################################################
###############################################################################

def rereferencing(data, chans, reference, to_ignore=None):
    """Re-reference data.

    Parameters
    ----------
    data : array_like
        The array of data of shape (nchan, npts).
    chans : list
        List of channel names of length nchan.
    reference : int
        The index of the channel to consider as a reference.
    to_ignore : list | None
        List of channels to ignore in the re-referencing.

    Returns
    -------
    datar : array_like
        The re-referenced data.
    channelsr : list
        List of re-referenced channel names.
    consider : list
        List of boolean values of channels that have to be considered
        during the ploting processus.
    """
    # Get shapes :
    nchan, npts = data.shape
    # Get data to use as the reference :
    ref = data[[reference], :]
    name = chans[reference]
    # Build ignore vector :
    consider = np.ones((nchan,), dtype=bool)
    consider[reference] = False
    # Find if some channels have to be ignored :
    if to_ignore is None:
        sl = slice(nchan)
    elif isinstance(to_ignore, (tuple, list, np.ndarray)):
        to_ignore = np.asarray(to_ignore)
        sl = np.arange(nchan)[~to_ignore]
        consider[to_ignore] = False
    # Re-reference data :
    data[sl, :] -= ref
    # Build channel names :
    chan = [k + '-' + name if consider[num]
            else k for num, k in enumerate(chans)]

    return data, chan, consider

def bipolarization(data, chans, to_ignore=None, sep='.'):
    """Bipolarize data.

    Parameters
    ----------
    data : array_like
        The array of data of shape (nchan, npts).
    chans : list
        List of channel names of length nchan.
    to_ignore : list | None
        List of channels to ignore in the bipolarization.
    sep : string | '.'
        Separator to simplify electrode names by removing undesired name
        after the sep. For example, if channel = ['h1.025', 'h2.578']
        and sep='.', the final name will be 'h2-h1'.

    Returns
    -------
    datar : array_like
        The re-referenced data.
    channelsr : list
        List of re-referenced channel names.
    consider : list
        List of boolean values of channels that have to be considered
        during the ploting processus.
    """
    # Variables :
    nchan, npts = data.shape
    consider = np.ones((nchan,), dtype=bool)

    # Preprocess channel names by separating channel names / number:
    chnames, chnums = [], []
    for num, k in enumerate(chans):
        # Remove spaces and separation :
        chans[num] = k.strip().replace(' ', '').split(sep)[0]
        # Get only the name / number :
        if findall(r'\d+', k):
            number = findall(r'\d+', k)[0]
            chnums.append(number)
            chnames.append(k.split(number)[0])
        else:
            chnums.append('')
            chnames.append(k)

    # Find if some channels have to be ignored :
    if to_ignore is None:
        sl = range(nchan)
    elif isinstance(to_ignore, (tuple, list, np.ndarray)):
        to_ignore = np.asarray(to_ignore)
        sl = np.arange(nchan)[~to_ignore]
        consider[to_ignore] = False

    # Bipolarize :
    for num in reversed(range(nchan)):
        # If there's a number :
        if chnums[num] and (num in sl):
            # Get the name of the channel to find :
            chan_to_find = chnames[num] + str(int(chnums[num]) - 1)
            # Search if exist in channel list :
            if chan_to_find in chans:
                # Get the index :
                ind = chans.index(chan_to_find)
                # Substract to data :
                data[num, :] -= data[ind, :]
                # Update channel name :
                chans[num] = chans[num] + '-' + chan_to_find
            else:
                consider[num] = False
        else:
            consider[num] = False

    return data, chans, consider

def commonaverage(data, chans, to_ignore=None):
    """Re-referencement using common average.

    Parameters
    ----------
    data : array_like
        The array of data of shape (nchan, npts).
    chans : list
        List of channel names of length nchan.
    to_ignore : list | None
        List of channels to ignore in the re-referencing.

    Returns
    -------
    datar : array_like
        The re-referenced data.
    channelsr : list
        List of re-referenced channel names.
    consider : list
        List of boolean values of channels that have to be considered
        during the ploting processus.
    """
    # Variables :
    nchan, npts = data.shape
    consider = np.ones((nchan,), dtype=bool)
    # Find if some channels have to be ignored :
    if to_ignore is not None:
        consider[to_ignore] = False
    # Get the mean across  EEG channels :
    eegmean = data[consider].mean(0, keepdims=True)
    # Remove the mean on EEG channels :
    data[consider, :] -= eegmean
    # Update channel name :
    for k in range(len(chans)):
        chans[k] = chans[k] + '-m' if consider[k] else chans[k]
    return data, chans, consider

###############################################################################
###############################################################################
#                               XYZ CONVERSION
###############################################################################
###############################################################################

def _spm_matrix(p):
    """Matrix transformation.

    Parameters
    ----------
    p : array_like
        Vector of floats for defining each tranformation. p must be a vector of
        length 9.

    Returns
    -------
    Pr : array_like
        The tranformed array.
    """
    q = [0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]
    p.extend(q[len(p):12])

    # Translation t :
    t = np.array([[1, 0, 0, p[0]],
                  [0, 1, 0, p[1]],
                  [0, 0, 1, p[2]],
                  [0, 0, 0, 1]])
    # Rotation 1 :
    r1 = np.array([[1, 0, 0, 0],
                   [0, np.cos(p[3]), np.sin(p[3]), 0],
                   [0, -np.sin(p[3]), np.cos(p[3]), 0],
                   [0, 0, 0, 1]])
    # Rotation 2 :
    r2 = np.array([[np.cos(p[4]), 0, np.sin(p[4]), 0],
                   [0, 1, 0, 0],
                   [-np.sin(p[4]), 0, np.cos(p[4]), 0],
                   [0, 0, 0, 1]])
    # Rotation 3 :
    r3 = np.array([[np.cos(p[5]), np.sin(p[5]), 0, 0],
                   [-np.sin(p[5]), np.cos(p[5]), 0, 0],
                   [0, 0, 1, 0],
                   [0, 0, 0, 1]])
    # Translation z :
    z = np.array([[p[6], 0, 0, 0],
                  [0, p[7], 0, 0],
                  [0, 0, p[8], 0],
                  [0, 0, 0, 1]])
    # Translation s :
    s = np.array([[1, p[9], p[10], 0],
                  [0, 1, p[11], 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]])
    return np.linalg.multi_dot([t, r1, r2, r3, z, s])

def tal2mni(xyz):
    """Transform Talairach coordinates into MNI.

    Parameters
    ----------
    xyz : array_like
        Array of Talairach coordinates of shape (n_sources, 3)

    Returns
    -------
    xyz_r : array_like
        Array of MNI coordinates of shape (n_sources, 3)
    """
    # Check xyz to be (n_sources, 3) :
    if (xyz.ndim != 2) or (xyz.shape[1] != 3):
        raise ValueError("The shape of xyz must be (N, 3).")
    n_sources = xyz.shape[0]

    # Transformation matrices, different zooms above/below AC :
    rotn = np.linalg.inv(_spm_matrix([0., 0., 0., .05]))
    upz = np.linalg.inv(_spm_matrix([0., 0., 0., 0., 0., 0., .99, .97, .92]))
    downz = np.linalg.inv(_spm_matrix([0., 0., 0., 0., 0., 0., .99, .97, .84]))

    # Apply rotation and translation :
    xyz = np.dot(rotn, np.c_[xyz, np.ones((n_sources, ))].T)
    tmp = np.array(xyz)[2, :] < 0.
    xyz[:, tmp] = np.dot(downz, xyz[:, tmp])
    xyz[:, ~tmp] = np.dot(upz, xyz[:, ~tmp])
    return np.array(xyz[0:3, :].T)

def mni2tal(xyz):
    """Transform MNI coordinates into Talairach.

    Parameters
    ----------
    xyz : array_like
        Array of MNI coordinates of shape (n_sources, 3)

    Returns
    -------
    xyz_r : array_like
        Array of Talairach coordinates of shape (n_sources, 3)
    """
    # Check xyz to be (n_sources, 3) :
    if (xyz.ndim != 2) or (xyz.shape[1] != 3):
        raise ValueError("The shape of xyz must be (N, 3).")
    n_sources = xyz.shape[0]

    # Transformation matrices, different zooms above/below AC :
    up_t = _spm_matrix([0., 0., 0., .05, 0., 0., .99, .97, .92])
    down_t = _spm_matrix([0., 0., 0., .05, 0., 0., .99, .97, .84])
    xyz = np.c_[xyz, np.ones((n_sources, ))].T

    tmp = np.array(xyz)[2, :] < 0.
    xyz[:, tmp] = np.dot(down_t, xyz[:, tmp])
    xyz[:, ~tmp] = np.dot(up_t, xyz[:, ~tmp])
    return np.array(xyz[0:3, :].T)


from .config import on_rtd

import os, re, sys
import warnings
import logging

if not on_rtd:
    import pandas as pd
    import numpy as np
    from scipy.interpolate import LinearNDInterpolator as interpnd
    import numpy.random as rand
    import matplotlib.pyplot as plt

    from astropy import constants as const

    #Define useful constants
    G = const.G.cgs.value
    MSUN = const.M_sun.cgs.value
    RSUN = const.R_sun.cgs.value

    from .extinction import EXTINCTION, LAMBDA_EFF, extcurve, extcurve_0
    from .interp import interp_value, interp_values

else:
    G = 6.67e-11
    MSUN = 1.99e33
    RSUN = 6.96e10

from .config import ISOCHRONES
from .grid import ModelGrid

def get_ichrone(models, bands=None, default=False, **kwargs):
    """Gets Isochrone Object by name, or type, with the right bands

    If `default` is `True`, then will set bands
    to be the union of bands and default_bands
    """
    if isinstance(models, Isochrone):
        return models

    def actual(bands, ictype):
        if bands is None:
            return list(ictype.default_bands)
        elif default:
            return list(set(bands).union(set(ictype.default_bands)))
        else:
            return bands

    if type(models) is type(type):
        ichrone = models(actual(bands, models))
    elif models=='dartmouth':
        from isochrones.dartmouth import Dartmouth_Isochrone
        ichrone = Dartmouth_Isochrone(bands=actual(bands, Dartmouth_Isochrone), **kwargs)
    elif models=='dartmouthfast':
        from isochrones.dartmouth import Dartmouth_FastIsochrone
        ichrone = Dartmouth_FastIsochrone(bands=actual(bands, Dartmouth_FastIsochrone), **kwargs)
    elif models=='mist':
        from isochrones.mist import MIST_Isochrone
        ichrone = MIST_Isochrone(bands=actual(bands, MIST_Isochrone), **kwargs)
    elif models=='padova':
        from isochrones.padova import Padova_Isochrone
        ichrone = Padova_Isochrone(bands=actual(bands, Padova_Isochrone), **kwargs)
    elif models=='basti':
        from isochrones.basti import Basti_Isochrone
        ichrone = Basti_Isochrone(bands=actual(bands, Basti_Isochrone), **kwargs)
    else:
        raise ValueError('Unknown stellar models: {}'.format(models))
    return ichrone

class Isochrone(object):
    """
    Basic isochrone class. Everything is a function of mass, log(age), Fe/H.

    Can be instantiated directly, but will typically be used with a pre-defined
    subclass, such as :class:`dartmouth.Dartmouth_Isochrone`.  All parameters
    must be array-like objects of the same length, with the exception of ``mags``,
    which is a dictionary of such array-like objects.

    :param m_ini:
        Array of initial mass values [msun].
    :type m_ini: array-like

    :param age:
        log10(age) [yr]

    :param feh:
        Metallicity [dex]

    :param m_act:
        Actual mass; same as m_ini if mass loss not implemented [msun]

    :param logL:
        log10(luminosity) [solar units]

    :param Teff:
        Effective temperature [K]

    :param logg:
        log10(surface gravity) [cgs]

    :param mags:
        Dictionary of absolute magnitudes in different bands
    :type mags: ``dict``

    :param tri:
        Triangulation object used
        to initialize the interpolation functions.
        If pre-computed triangulation not provided, then the constructor
        will calculate one.  This might take several minutes, so be patient.
        Much better to use pre-computed ones, as provided in, e.g.,
        :class:`dartmouth.Dartmouth_Isochrone`.
    :type tri: :class:`scipy.spatial.qhull.Delaunay`, optional

    :param minage,maxage:
        If desired, a minimum or maximum age can be manually entered.

    """
    def __init__(self,m_ini,age,feh,m_act,logL,Teff,logg,mags,tri=None,
                 minage=None, maxage=None, ext_table=False):
        """Warning: if tri object not provided, this will be very slow to be created.
        """

        self.minage = age.min()
        self.maxage = age.max()
        self.minmass = m_act.min()
        self.maxmass = m_act.max()
        self.minfeh = feh.min()
        self.maxfeh = feh.max()

        self.ext_table = ext_table

        if minage is not None:
            logging.warning("minage and maxage keywords are deprecated." + \
                          "Use instead the .set_bounds(age=(lo, hi)) attribute of StarModel.")
            self.minage = minage
        if maxage is not None:
            logging.warning("minage and maxage keywords are deprecated." + \
                          "Use instead the .set_bounds(age=(lo, hi)) attribute of StarModel.")
            self.maxage = maxage

        L = 10**logL

        if tri is None:
            points = np.zeros((len(m_ini),3))
            points[:,0] = m_ini
            points[:,1] = age
            points[:,2] = feh
            fn = interpnd(points,m_act)
            self.tri = fn.tri
        else:
            self.tri = tri
            self.mass = interpnd(self.tri,m_act)

        self._data = {'mass':m_act,
                    'logL':logL,
                    'logg':logg,
                    'logTeff':np.log10(Teff),
                    'mags':mags}
        self._props = ['mass', 'logL', 'logg', 'logTeff']

        self.bands = list(mags.keys())

        self._mag = {band:interpnd(self.tri,mags[band]) for band in self.bands}

        self.mag = {b : self._mag_fn(b) for b in self.bands}

    def __getstate__(self):
        odict = self.__dict__.copy()
        del odict['mag'] # This can't be pickled
        return odict

    def __setstate__(self, odict):
        self.__dict__ = odict
        self.__dict__['mag'] = {b : self._mag_fn(b) for b in self.bands}

    def _prop(self, prop, *args):
        if prop not in self._props:
            raise ValueError('Cannot call this function with {}.'.format(prop))
        attr = '_{}'.format(prop)
        if not hasattr(self, attr):
            setattr(self, attr, interpnd(self.tri, self._data[prop]))
        fn = getattr(self, attr)
        return fn(*args)

    def mass(self, *args):
        return self._prop('mass', *args)

    def logL(self, *args):
        return self._prop('logL', *args)

    def logg(self, *args):
        return self._prop('logg', *args)

    def logTeff(self, *args):
        return self._prop('logTeff', *args)

    def radius(self, *args):
        return np.sqrt(G*self.mass(*args)*MSUN/10**self.logg(*args))/RSUN

    def Teff(self, *args):
        return 10**self.logTeff(*args)

    def density(self, *args):
        """ Mean density in g/cc
        """
        M = self.mass(*args) * MSUN
        V = 4./3 * np.pi * (self.radius(*args) * RSUN)**3
        return  M/V

    def delta_nu(self, *args):
        """Returns asteroseismic delta_nu in uHz

        reference: https://arxiv.org/pdf/1312.3853v1.pdf, Eq (2)
        """
        return 134.88 * np.sqrt(self.mass(*args) / self.radius(*args)**3)

    def nu_max(self, *args):
        """Returns asteroseismic nu_max in uHz

        reference: https://arxiv.org/pdf/1312.3853v1.pdf, Eq (3)
        """
        return 3120.* (self.mass(*args) /
                        (self.radius(*args)**2 * np.sqrt(self.Teff(*args)/5777.)))

    def _mag_fn(self, band):
        def fn(mass, age, feh, distance=10, AV=0.0, x_ext=0., ext_table=self.ext_table):
            if x_ext==0.:
                ext = extcurve_0
            else:
                ext = extcurve(x_ext)
            if ext_table:
                A = AV*EXTINCTION[band]
            else:
                A = AV*ext(LAMBDA_EFF[band])
            dm = 5*np.log10(distance) - 5
            return self._mag[band](mass, age, feh) + dm + A
        return fn

    def __call__(self, mass, age, feh,
                 distance=None, AV=0.0,
                 return_df=True, bands=None):
        """
        Returns all properties (or arrays of properties) at given mass, age, feh

        :param mass, age, feh:
            Mass, log(age), metallicity.  Can be float or array_like.

        :param distance:
            Distance in pc.  If passed, then mags will be converted to
            apparent mags based on distance (and ``AV``).

        :param AV:
            V-band extinction (magnitudes).

        :param return_df: (optional)
            If ``True``, return :class:``pandas.DataFrame`` containing all model
            parameters at each input value; if ``False``, return dictionary
            of the same.

        :param bands: (optional)
            List of photometric bands in which to return magnitudes.
            Must be subset of ``self.bands``.  If not set, then will
            default to returning all available bands.

        :return:
            Either a :class:`pandas.DataFrame` or a dictionary containing
            model values evaluated at input points.
        """
        # Broadcast inputs to the same shape
        mass, age, feh = [np.array(a) for a in np.broadcast_arrays(mass, age, feh)]
        args = (mass, age, feh)
        Ms = self.mass(*args)*1
        Rs = self.radius(*args)*1
        logLs = self.logL(*args)*1
        loggs = self.logg(*args)*1
        Teffs = self.Teff(*args)*1
        if bands is None:
            bands = self.bands
        if distance is not None:
            args += (distance, AV)
        mags = {band:1*self.mag[band](*args) for band in bands}
        # if distance is not None:
        #     dm = 5*np.log10(distance) - 5
        #     for band in mags:
        #         A = AV*EXTINCTION[band]
        #         mags[band] = mags[band] + dm + A

        props = {'age':age,'mass':Ms,'radius':Rs,'logL':logLs,
                'logg':loggs,'Teff':Teffs,'mag':mags}

        if not return_df:
            return props
        else:
            d = {}
            for key in props.keys():
                if key=='mag':
                    for m in props['mag'].keys():
                        d['{}_mag'.format(m)] = props['mag'][m]
                else:
                    d[key] = props[key]
            df = pd.DataFrame(d)
            return df

    def agerange(self, m, feh=0.0):
        """
        For a given mass and feh, returns the min and max allowed ages.
        """
        ages = np.arange(self.minage, self.maxage, 0.01)
        rs = self.radius(m, ages, feh)
        w = np.where(np.isfinite(rs))[0]
        return ages[w[0]],ages[w[-1]]

    def evtrack(self,m,feh=0.0,minage=None,maxage=None,dage=0.02,
                return_df=True):
        """
        Returns evolution track for a single initial mass and feh.

        :param m:
            Initial mass of desired evolution track.

        :param feh: (optional)
            Metallicity of desired track.  Default = 0.0 (solar)

        :param minage, maxage: (optional)
            Minimum and maximum log(age) of desired track. Will default
            to min and max age of model isochrones.

        :param dage: (optional)
            Spacing in log(age) at which to evaluate models.  Default = 0.02

        :param return_df: (optional)
            Whether to return a ``DataFrame`` or dicionary.  Default is ``True``.

        :return:
            Either a :class:`pandas.DataFrame` or dictionary
            representing the evolution
            track---fixed mass, sampled at chosen range of ages.

        """
        if minage is None:
            minage = self.minage
        if maxage is None:
            maxage = self.maxage
        ages = np.arange(minage,maxage,dage)
        Ms = self.mass(m,ages,feh)
        Rs = self.radius(m,ages,feh)
        logLs = self.logL(m,ages,feh)
        loggs = self.logg(m,ages,feh)
        Teffs = self.Teff(m,ages,feh)
        mags = {band:self.mag[band](m,ages,feh) for band in self.bands}

        props = {'age':ages,'mass':Ms,'radius':Rs,'logL':logLs,
                'logg':loggs, 'Teff':Teffs, 'mag':mags}

        if not return_df:
            return props
        else:
            d = {}
            for key in props.keys():
                if key=='mag':
                    for m in props['mag'].keys():
                        d['{}_mag'.format(m)] = props['mag'][m]
                else:
                    d[key] = props[key]
            try:
                df = pd.DataFrame(d)
            except ValueError:
                df = pd.DataFrame(d, index=[0])
            return df

    def isochrone(self,age,feh=0.0,minm=None,maxm=None,dm=0.02,
                  return_df=True,distance=None,AV=0.0):
        """
        Returns stellar models at constant age and feh, for a range of masses

        :param age:
            log10(age) of desired isochrone.

        :param feh: (optional)
            Metallicity of desired isochrone (default = 0.0)

        :param minm, maxm: (optional)
            Mass range of desired isochrone (will default to max and min available)

        :param dm: (optional)
            Spacing in mass of desired isochrone.  Default = 0.02 Msun.

        :param return_df: (optional)
            Whether to return a :class:``pandas.DataFrame`` or dictionary.  Default is ``True``.

        :param distance:
            Distance in pc.  If passed, then mags will be converted to
            apparent mags based on distance (and ``AV``).

        :param AV:
            V-band extinction (magnitudes).

        :return:
            :class:`pandas.DataFrame` or dictionary containing results.

        """
        if minm is None:
            minm = self.minmass
        if maxm is None:
            maxm = self.maxmass
        ms = np.arange(minm,maxm,dm)
        ages = np.ones(ms.shape)*age

        Ms = self.mass(ms,ages,feh)
        Rs = self.radius(ms,ages,feh)
        logLs = self.logL(ms,ages,feh)
        loggs = self.logg(ms,ages,feh)
        Teffs = self.Teff(ms,ages,feh)
        mags = {band:self.mag[band](ms,ages,feh) for band in self.bands}
        #for band in self.bands:
        #    mags[band] = self.mag[band](ms,ages)
        if distance is not None:
            dm = 5*np.log10(distance) - 5
            for band in mags:
                A = AV*EXTINCTION[band]
                mags[band] = mags[band] + dm + A

        props = {'M':Ms,'R':Rs,'logL':logLs,'logg':loggs,
                'Teff':Teffs,'mag':mags}

        if not return_df:
            return props
        else:
            d = {}
            for key in props.keys():
                if key=='mag':
                    for m in props['mag'].keys():
                        d['{}_mag'.format(m)] = props['mag'][m]
                else:
                    d[key] = props[key]
            try:
                df = pd.DataFrame(d)
            except ValueError:
                df = pd.DataFrame(d, index=[0])
            return df

    def random_points(self,n,minmass=None,maxmass=None,
                      minage=None,maxage=None,
                      minfeh=None,maxfeh=None):
        """
        Returns n random mass, age, feh points, none of which are out of range.

        :param n:
            Number of desired points.

        :param minmass, maxmass: (optional)
            Desired allowed range.  Default is mass range of ``self``.

        :param minage, maxage: (optional)
            Desired allowed range.  Default is log10(age) range of
            ``self``.

        :param minfehs, maxfeh: (optional)
            Desired allowed range.  Default is feh range of ``self``.

        :return:
            :class:`np.ndarray` arrays of randomly selected mass, log10(age),
            and feh values
            within allowed ranges.  Used, e.g., to initialize random walkers for
            :class:`StarModel` fits.

        .. todo::

            Should change this to drawing from priors!  Current implementation
            is a bit outdated.
        """
        if minmass is None:
            minmass = self.minmass
        if maxmass is None:
            maxmass = self.maxmass
        if minage is None:
            minage = self.minage
        if maxage is None:
            maxage = self.maxage
        if minfeh is None:
            minfeh = self.minfeh
        if maxfeh is None:
            maxfeh = self.maxfeh

        ms = rand.uniform(minmass,maxmass,size=n)
        ages = rand.uniform(minage,maxage,size=n)
        fehs = rand.uniform(minage,maxage,size=n)

        Rs = self.radius(ms,ages,fehs)
        bad = np.isnan(Rs)
        nbad = bad.sum()
        while nbad > 0:
            ms[bad] = rand.uniform(minmass,maxmass,size=nbad)
            ages[bad] = rand.uniform(minage,maxage,size=nbad)
            fehs[bad] = rand.uniform(minfeh,maxfeh,size=nbad)
            Rs = self.radius(ms,ages,fehs)
            bad = np.isnan(Rs)
            nbad = bad.sum()
        return ms,ages,fehs

class MagFunction(object):
    def __init__(self, ic, band, icol):
        self.ic = ic
        self.band = band
        self.icol = icol
        self.x_ext = ic.x_ext
        self.ext_table = ic.ext_table

        if self.x_ext==0.:
            ext = extcurve_0
        else:
            ext = extcurve(x_ext)
        if self.ext_table:
            self.AAV = EXTINCTION[self.band]
        else:
            self.AAV = ext(LAMBDA_EFF[self.band])

    def __call__(self, mass, age, feh, distance=10, AV=0.0, x_ext=None, ext_table=False):
        if x_ext is not None:
            if x_ext==0.:
                ext = extcurve_0
            else:
                ext = extcurve(x_ext)

            if ext_table:
                AAV = EXTINCTION[self.band]
            else:
                AAV = ext(LAMBDA_EFF[self.band])
        else:
            AAV = self.AAV

        A = AV*AAV
        dm = 5*np.log10(distance) - 5
        mag = self.ic.interp_value(mass, age, feh, self.icol)
        return mag + dm + A

class FastIsochrone(Isochrone):
    """Alternative isochrone implementation for large grids, faster likelihoods

    This implementation allows faster point interpolations than
    the triangulation method in the base :class:`Isochrone` class;
    however, for simulating large populations (passing arrays of parameters
    at a time), the base :class:`Isochrone` is still faster.  This is
    also the go-to implementation for grids that are too large to make the
    Delaunay method feasible (e.g., significantly over 100,000 points in the grid.)

    However, edge effects might be a bit more problematic here than in the base class,
    meaning fitting very evolved stars might be an issue. Also, don't trust this subclass
    for the :function:`Isochrone.agerange` function, for the same reason.

    Subclasses must set the appropriate attributes, and then things should work.
    """
    name = 'default'
    modelgrid = ModelGrid
    age_col = None
    feh_col = None
    mass_col = None
    loggTeff_col = None
    logg_col = None
    logL_col = None
    default_bands = ('g')

    def __init__(self, bands=None, x_ext=0., ext_table=False, debug=False, **kwargs):
        # df should be indexed by [feh, age]

        if bands is None:
            bands = list(self.default_bands)
        bands = sorted(bands)
        self.bands = bands

        self._df = None
        self.x_ext = 0.
        self.ext_table = ext_table
        self.debug = debug

        self._fehs = None
        self._ages = None
        self._Nfeh = None
        self._Nage = None

        self._minage = None
        self._maxage = None
        self._minmass = None
        self._maxmass = None
        self._minfeh = None
        self._maxfeh = None

        n_common_cols = len(self.modelgrid.get_common_columns(**kwargs))
        self._mag_cols = {b:n_common_cols+i for i,b in enumerate(self.bands)}
        self.mag = {b: MagFunction(self, b, i)
                            for b,i in self._mag_cols.items()}

        #organized array
        self._grid = None
        self._grid_Ns = None

        # kwargs to pass to self.modelgrid
        self.modelgrid_kwargs = kwargs

    def _initialize(self):
        for attr in ['df','Ncols','fehs','ages','Nfeh','Nage',
                     'minage','maxage','minfeh','maxfeh','minmass','maxmass']:
            _ = getattr(self, attr)

    @property
    def df(self):
        if self._df is None:
            self._df = self.modelgrid(self.bands, **self.modelgrid_kwargs).df
        return self._df

    @property
    def Ncols(self):
        return self.df.shape[1]

    @property
    def fehs(self):
        if self._fehs is None:
            self._fehs = self.df.iloc[:, self.feh_col].unique()
        return self._fehs

    @property
    def ages(self):
        if self._ages is None:
            self._ages = self.df.iloc[:, self.age_col].unique()
        return self._ages

    @property
    def Nfeh(self):
        if self._Nfeh is None:
            self._Nfeh = len(self.fehs)
        return self._Nfeh

    @property
    def Nage(self):
        if self._Nage is None:
            self._Nage = len(self.ages)
        return self._Nage

    @property
    def minage(self):
        if self._minage is None:
            self._minage = self.ages.min()
        return self._minage

    @property
    def maxage(self):
        if self._maxage is None:
            self._maxage = self.ages.max()
        return self._maxage

    @property
    def minfeh(self):
        if self._minfeh is None:
            self._minfeh = self.fehs.min()
        return self._minfeh

    @property
    def maxfeh(self):
        if self._maxfeh is None:
            self._maxfeh = self.fehs.max()
        return self._maxfeh

    @property
    def minmass(self):
        if self._minmass is None:
            self._minmass = self.df.iloc[:, self.mass_col].min()
        return self._minmass

    @property
    def maxmass(self):
        if self._maxmass is None:
            self._maxmass = self.df.iloc[:, self.mass_col].max()
        return self._maxmass

    def logTeff(self, mass, age, feh):
        return self.interp_value(mass, age, feh, self.loggTeff_col)

    def logg(self, mass, age, feh):
        return self.interp_value(mass, age, feh, self.logg_col)

    def logL(self, mass, age, feh):
        return self.interp_value(mass, age, feh, self.logL_col)

    def mass(self, *args):
        if np.size(args[0]) > 1:
            return np.array(args[0])
        else:
            return args[0]

    @property
    def grid(self):
        if self._grid is None:
            self._make_grid()
        return self._grid

    @property
    def grid_Ns(self):
        if self._grid_Ns is None:
            self._make_grid()
        return self._grid_Ns

    @property
    def _npz_filename(self):
        keys = list(self.modelgrid_kwargs.keys())
        keys.sort()

        filename = os.path.join(ISOCHRONES, self.name, '{}'.format('-'.join(self.bands)))

        for k in keys:
            filename += '_{}{}'.format(k, self.modelgrid_kwargs[k])

        filename += '.npz'
        return filename

    def _make_grid(self, recalc=False):
        # Read from file if available.
        if os.path.exists(self._npz_filename) and not recalc:
            d = np.load(self._npz_filename)
            self._grid = d['grid']
            self._grid_Ns = d['grid_Ns']
        else:
            df_list = [[self.df.ix[f,a] for f in self.fehs] for a in self.ages]
            lens = np.array([[len(df_list[i][j]) for j in range(self.Nfeh)]
                             for i in range(self.Nage)]).T #just because
            data = np.zeros((self.Nfeh, self.Nage, lens.max(), self.Ncols))

            for i in range(self.Nage):
                for j in range(self.Nfeh):
                    N = lens[j,i]
                    data[j, i, :N, :] = df_list[i][j].values
                    data[j, i, N:, :] = np.nan

            np.savez(self._npz_filename, grid=data, grid_Ns=lens)
            self._grid = data
            self._grid_Ns = lens

    def interp_value(self, mass, age, feh, icol): # 4 is log_g
        if self._ages is None:
            self._initialize()

        try:
            return interp_value(float(mass), float(age), float(feh), icol,
                                self.grid, self.mass_col,
                                self.ages, self.fehs, self.grid_Ns, self.debug)

        except:
            # First, broadcast to common shape.
            b = np.broadcast(mass, age, feh)
            mass = np.resize(mass, b.shape).astype(float)
            age = np.resize(age, b.shape).astype(float)
            feh = np.resize(feh, b.shape).astype(float)

            # Then pass to helper function
            return interp_values(mass, age, feh, icol,
                                self.grid, self.mass_col,
                                self.ages, self.fehs, self.grid_Ns)


"""
Python interface for ceq chemical equilibrium calculator

References:
    "Computer Program for Calculation of Complex Equilibrium Compositions and Applications"
    Nasa Reference Publication 1311, October 1995
    Sanford Gordon and Bonnie J. McBride

    "NASA Glenn Coefficients for Calculating Thermodynamic Properties of Individual Species"
    NASA/TP - 2002-211556, September 2002
    Bonnie J. McBride, Michael J. Zehe, and Sanford Gordon

@author: Nick Gibbons (n.gibbons(at)uq.edu.au)
"""

from string import ascii_letters
from numpy import array, zeros, log
from ctypes import cdll,c_double,POINTER,c_int,byref
from clib import *

letters = set(ascii_letters)
DBPATH='../thermo.inp'
LIBPATH='./libceq.so'
HEADERFILE='./ceq.h'

class EqCalculator(object):
    """ Python interface to low level ceq routines """
    def __init__(self, spnames):
        self.spnames = spnames
        self.nsp = len(spnames)
        self.lib = self.load_ceq_library()

        atoms = []
        M = []
        lewis = []
        for sp in spnames:
            asp, Msp, lsp = self.get_thermo_data(sp)
            atoms.append(asp)
            M.append(Msp)
            lewis.append(lsp)

        self.atoms = atoms
        self.M = array(M)
        self.lewis = array(lewis)

        elements = set()
        for a in atoms:
            for k in a.keys(): elements.add(k)
        elements = list(elements)
        elements.sort()

        self.elements = elements
        self.nel = len(elements)

        a = zeros((len(elements),len(spnames)))
        for i,s in enumerate(atoms):
            for k,v in s.items():
                j = elements.index(k)
                a[j,i] = v

        self.a = a

        return

    def get_thermo_data(self, name):
        """ Stripped down version of get_species from lewis_thermo.py """
        data = self.readdb(name)
        header = data[0].strip().split()
        if header[0]!=name:
            raise IOError("Database read failed! {}!={}".format(name, header[0]))

        info = data[1].strip().split()
        intervals = int(info[0])
        M = float(info[-2])/1e3 # convert to kg/mol from kg/kmol

        atomstrings = [data[1][10+8*i:10+8*(i+1)] for i in range(5)]
        atoms = {}
        for s in atomstrings:
            if s[0]==' ': continue
            elementname,amount = s.split()
            atoms[elementname] = float(amount)

        lewis = []
        for i in range(intervals):
            li = []
            line1 = data[3+3*i].replace('D','E') # Fix Fortran Double nonsense
            li.extend([float(line1[16*j:16*(j+1)]) for j in range(5)])

            line2 = data[4+3*i].replace('D','E') # Fix Fortran Double nonsense
            li.extend([float(line2[16*j:16*(j+1)]) for j in range(2)])
            li.extend([float(line2[16*j:16*(j+1)]) for j in range(3,5)])
            lewis.append(li)

        if len(lewis)!=3:
            lsp3 = self.fix_missing_thermo_segment(lewis, M)
            lewis.append(lsp3)
        assert len(lewis)==3 and len(lewis[0])==9 and len(lewis[1])==9 and len(lewis[2])==9
        return atoms, M, lewis

    def readdb(self,name):
        """ Retrieve species 'name' from the lewis_thermo.db file """
        with open(DBPATH) as fp:

            iterline=iter(fp) # Make an iterator for inner loop behaviour

            # Start looking for the beginning of a species entry
            for line in iterline:
                if not line[0] in letters: continue # Skip if not a letter in first position

                if line.startswith(name):        # This could be the one, check the rest
                    header = line.strip().split()
                    if header[0]!=name: continue # Nope false alarm

                    lines = [line]               # We've found it!
                    for nextline in iterline:
                        if nextline[0] in letters: break
                        lines.append(nextline)
                    break                        # Break the outer for loop and jump to return lines
            else:
                raise Exception("Name: {} not found!".format(name))
        return lines

    def fix_missing_thermo_segment(self, lsp, Mi):
        """ Create a NASA9 polynominal assuming constant Cp above 6000 K"""
        assert len(lsp)==2
        lspa = zeros((3,9))
        lspa[0] = lsp[0] # This should copy data automatically
        lspa[1] = lsp[1] # This should copy data automatically

        T = 5999.99999
        Ru = 8.3144621
        X = array([1.0])
        M = array([Mi])
        Td = c_double(T)

        cp = self.lib.get_cp(T, X, 1, lspa, M)
        h  = self.lib.get_h(T, X, 1, lspa, M)
        s0 = self.lib.get_s0(T, X, 1, lspa, M)

        a2 = cp*Mi/Ru
        b1 = h*Mi/Ru - a2*T
        b2 = s0*Mi/Ru - a2*log(T)
        lsp2 = [0.0, 0.0, a2, 0.0, 0.0, 0.0, 0.0, b1, b2]
        return lsp2

    @staticmethod
    def load_ceq_library(LIBPATH=LIBPATH):
        """ Load the c library and set return types """
        lib = CLib(LIBPATH,[HEADERFILE])
        return lib

    def pt(self, p, T, Xs0, verbose=0):
        """ Call c library to compute equilibrium concentrations at fixed p, T """
        if Xs0.size!=self.nsp: raise Exception('Mismatched array size {}!={}'.format(Xs0.size, self.nsp))
        Xs1 = zeros(Xs0.shape)

        recode = self.lib.pt(p, T, Xs0, self.nsp, self.nel, self.lewis, self.M, self.a, Xs1, verbose)
        if recode!=0: raise Exception("Equilibrium Calc Failed.")
        return Xs1

    def rhou(self, rho, u, Xs0, verbose=0):
        """ Call c library to compute equilibrium concentrations at fixed rho, u """
        if Xs0.size!=self.nsp: raise Exception('Mismatched array size {}!={}'.format(Xs0.size, self.nsp))
        Xs1 = zeros(Xs0.shape)
        Tref = zeros(1)

        recode = self.lib.rhou(rho, u, Xs0, self.nsp, self.nel, self.lewis, self.M, self.a, Xs1, Tref, verbose)
        if recode!=0: raise Exception("Equilibrium Calc Failed.")
        T = Tref[0]
        return Xs1, T

    def ps(self, pt, st, Xs0, verbose=0):
        """ Call c library to compute equilibrium concentrations at fixed p, s """
        if Xs0.size!=self.nsp: raise Exception('Mismatched array size {}!={}'.format(Xs0.size, self.nsp))
        Xs1 = zeros(Xs0.shape)
        Tref = zeros(1)

        recode = self.lib.ps(pt, st, Xs0, self.nsp, self.nel, self.lewis, self.M, self.a, Xs1, Tref, verbose)
        if recode!=0: raise Exception("Equilibrium Calc Failed.")
        T = Tref[0]
        return Xs1, T

    def rhot(self, rho, T, Xs0, verbose=0):
        """ Call c library to compute equilibrium concentrations at fixed rho, T """
        if Xs0.size!=self.nsp: raise Exception('Mismatched array size {}!={}'.format(Xs0.size, self.nsp))
        Xs1 = zeros(Xs0.shape)

        recode = self.lib.rhot(rho, T, Xs0, self.nsp, self.nel, self.lewis, self.M, self.a, Xs1, verbose)
        if recode!=0: raise Exception("Equilibrium Calc Failed.")
        return Xs1

    def get_u(self, X, T):
        """ Call c library to compute internal energy at fixed composition and temperature """
        u = self.lib.get_u(T, X, self.nsp, self.lewis, self.M)
        return u

    def get_h(self, X, T):
        """ Call c library to compute internal energy at fixed composition and temperature """
        h = self.lib.get_h(T, X, self.nsp, self.lewis, self.M)
        return h

    def get_cp(self, X, T):
        """ Call c library to compute internal energy at fixed composition and temperature """
        cp = self.lib.get_cp(T, X, self.nsp, self.lewis, self.M)
        return cp

    def get_s0(self, X, T):
        """ Call c library to compute specific entropy at standard state and arbitrary temperature """
        s0 = self.lib.get_s0(T, X, self.nsp, self.lewis, self.M)
        return s0

    def get_s(self, X, T, p):
        """ Call c library to compute internal entropy at an arbitrary pressure and temperature """
        s = self.lib.get_s(T, p, X, self.nsp, self.lewis, self.M)
        return s

    def batch_pt(self, p, T, Xs0, verbose=0):
        """ Call c library to compute equilibrium concentrations at fixed p, T """
        N, nspcheck = Xs0.shape
        if not Xs0.flags['OWNDATA']: raise Exception("Xs0 Memory Error: Array must own its data")
        if nspcheck!=self.nsp: raise Exception("nsp ({}) != Xs0.shape[1] ({})".format(self.nsp, nspcheck))
        if N!=p.size: raise Exception("p.size ({}) != Xs0.shape[0] ({})".format(p.size, N))
        if N!=T.size: raise Exception("T.size ({}) != Xs0.shape[0] ({})".format(T.size, N))

        Xs1 = zeros(Xs0.shape)

        recode = self.lib.batch_pt(N, p, T, Xs0, self.nsp, self.nel, self.lewis, self.M, self.a, Xs1, verbose)
        if recode!=0: raise Exception("Equilibrium Calc Failed.")
        return Xs1

    def batch_rhou(self, rho, u, Xs0, verbose=0):
        """ Call c library to compute equilibrium concentrations at fixed rho, u """
        N, nspcheck = Xs0.shape
        if not Xs0.flags['OWNDATA']: raise Exception("Xs0 Memory Error: Array must own its data")
        if nspcheck!=self.nsp: raise Exception("nsp ({}) != Xs0.shape[1] ({})".format(self.nsp, nspcheck))
        if N!=rho.size: raise Exception("rho.size ({}) != Xs0.shape[0] ({})".format(rho.size, N))
        if N!=u.size: raise Exception("u.size ({}) != Xs0.shape[0] ({})".format(u.size, N))

        Xs1 = zeros(Xs0.shape)
        T = zeros(rho.shape)

        recode = self.lib.batch_rhou(N,rho,u,Xs0,self.nsp,self.nel,self.lewis,self.M,self.a,Xs1,T,verbose)
        if recode!=0: raise Exception("Equilibrium Calc Failed.")
        return Xs1, T

    def batch_u(self, X, T):
        """ Call c library to compute internal energy at fixed composition and temperature """
        N, nspcheck = X.shape
        if not X.flags['OWNDATA']: raise Exception("X Memory Error: Array must own its data")
        if nspcheck!=self.nsp: raise Exception("nsp ({}) != X.shape[1] ({})".format(self.nsp, nspcheck))
        if N!=T.size: raise Exception("T.size ({}) != X.shape[0] ({})".format(T.size, N))

        u = zeros(T.shape)

        recode = self.lib.batch_u(N, T, X, self.nsp, self.lewis, self.M, u)
        if recode!=0: raise Exception("u calc failed.")
        return u

    def YtoX(self, Y):
        Mmix = 1.0/((Y/self.M).sum())
        X = Y*Mmix/self.M
        return X

    def XtoY(self, X):
        Mmix = (X*self.M).sum()
        Y = X*self.M/Mmix
        return Y

if __name__=='__main__':
    print("Called pyeq main!")