#!/usr/bin/env python
'''
mcu: Modeling and Crystallographic Utilities
Copyright (C) 2019 Hung Q. Pham. All Rights Reserved.

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

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

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Email: Hung Q. Pham <pqh3.14@gmail.com>
'''

import numpy as np
import re, argparse
from ..utils.misc import check_exist
from ..vasp import const
from ..cell import utils as cell_utils

lattice_MATCH = re.compile(r'''
[\w\W]*
 Lattice [ ]* vectors [ ]* \:
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n[ ]\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

species_MATCH = re.compile(r'''
[ ]*
 Species [ ]* \:
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n[ ]\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

kmesh_MATCH = re.compile(r'''
[\w\W]* k-point [ ]* grid [ ]* \: [ ]* (?P<kx>\S+) [ ]* (?P<ky>\S+) [ ]* (?P<kz>\S+)
''', re.VERBOSE)

koffset_MATCH = re.compile(r'''
[\w\W]* k-point [ ]* offset [ ]* \: [ ]* (?P<kx>\S+) [ ]* (?P<ky>\S+) [ ]* (?P<kz>\S+)
''', re.VERBOSE)

total_charge_MATCH = re.compile(r'''
[\w\W]*
Total [ ]* nuclear [ ]* charge    [ ]* \: [ ]* (?P<nuclear>\S+) [ ]* \n
Total [ ]* core [ ]* charge       [ ]* \: [ ]* (?P<core>\S+) [ ]* \n
Total [ ]* valence [ ]* charge    [ ]* \: [ ]* (?P<valence>\S+) [ ]* \n
Total [ ]* excess [ ]* charge     [ ]* \: [ ]* (?P<excess>\S+) [ ]* \n
Total [ ]* electronic [ ]* charge [ ]* \: [ ]* (?P<electronic>\S+)
''', re.VERBOSE)

states_charge_MATCH = re.compile(r'''
[\w\W]*
Number [ ]* of [ ]* empty [ ]* states  [ ]* \: [ ]* (?P<empty>\S+) [ ]* \n
Total [ ]* number [ ]* of [ ]* valence [ ]* states [ ]* \: [ ]* (?P<valence>\S+) [ ]* \n
Total [ ]* number [ ]* of [ ]* core [ ]* states [ ]* \: [ ]* (?P<core>\S+)
''', re.VERBOSE)

############ For SCF loops ##################
energies_MATCH = re.compile(r'''
[ ]*
 Energies [ ]* \:
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n[ ]\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

charges_MATCH = re.compile(r'''
[ ]*
 Charges [ ]* \:
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n[ ]\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

def read_info(filename):
    '''Read the INFO.OUT file'''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, "r") as data_file:
        data = data_file.read()

        # spin
        spin = int('spin-polarised' in data)
        soc = 'spin-orbit coupling' in data
        if soc: spin = False

        # Cell information
        # row-vectors instead of column-vectors as in elk output
        tmp = lattice_MATCH.match(data)['content']
        lattice = np.float64(tmp.split()).reshape(3,3, order='F')

        # Get atom info
        species_blocks = []
        atom = []
        magnetic = []
        frac_coords = []
        for atoms in species_MATCH.finditer(data):
            atoms_data = atoms.groupdict()['content']
            tmp = atoms_data.split('\n')
            species = {}
            species['atom type'] = re.findall(r"\((?P<val>\S+)\)", tmp[0])[0]
            species['parameters'] = tmp[1].split(':')[1].strip()
            species['name'] = tmp[2].split(':')[1].strip()
            species['nuclear charge'] = float(tmp[3].split(':')[1])
            species['electronic charge'] = float(tmp[4].split(':')[1])
            species['atomic mass'] = float(tmp[5].split(':')[1])
            species['muffin-tin radius'] = float(tmp[6].split(':')[1])
            species['muffin-tin radial'] = int(tmp[7].split(':')[1])
            species['muffin-tin inner'] = int(tmp[8].split(':')[1])
            for atm in tmp[10:]:
                atom.append(species['atom type'])
                atm_split = atm.split()
                frac_coords.append(np.float64(atm_split[2:5]))
                magnetic.append(np.float64(atm_split[5:]))

            species_blocks.append(species)

        # Make the cell tuple
        atom_number = cell_utils.convert_atomtype(atom)
        cell = (lattice * const.AUTOA, frac_coords, atom_number)

        # kmesh
        tmp = kmesh_MATCH.match(data)
        kmesh = np.int64([tmp['kx'], tmp['ky'], tmp['kz']])
        tmp = koffset_MATCH.match(data)
        kmesh_shift = np.float64([tmp['kx'], tmp['ky'], tmp['kz']])

        # total charge
        charge_data = total_charge_MATCH.match(data)
        charge = {}
        charge['nuclear'] = float(charge_data['nuclear'])
        charge['core'] = float(charge_data['core'])
        charge['valence'] = float(charge_data['valence'])
        charge['excess'] = float(charge_data['excess'])
        charge['electronic'] = float(charge_data['electronic'])

        # Total states:
        states_data = states_charge_MATCH.match(data)
        states = {}
        states['empty'] = states_data['empty']
        states['core'] = states_data['core']
        states['valence'] = states_data['valence']

        # Collect the SCF cycles
        energies_blocks = []
        for cycle in energies_MATCH.finditer(data):
            energies_data = cycle.groupdict()['content']
            tmp = energies_data.split('\n')
            energies = {}
            energies['Fermi'] = float(tmp[1].split(':')[1])
            energies['sum of eigenvalues'] = float(tmp[2].split(':')[1])
            energies['electron kinetic'] = float(tmp[3].split(':')[1])
            energies['core electron kinetic'] = float(tmp[4].split(':')[1])
            energies['Coulomb'] = float(tmp[5].split(':')[1])
            energies['Coulomb potential'] = float(tmp[6].split(':')[1])
            energies['nuclear-nuclear'] = float(tmp[7].split(':')[1])
            energies['electron-nuclear'] = float(tmp[8].split(':')[1])
            energies['Hartree'] = float(tmp[9].split(':')[1])
            energies['Madelung'] = float(tmp[10].split(':')[1])
            energies['xc potential'] = float(tmp[11].split(':')[1])
            energies['exchange'] = float(tmp[12].split(':')[1])
            energies['correlation'] = float(tmp[13].split(':')[1])
            energies['electron entropic'] = float(tmp[14].split(':')[1])
            energies['total energy'] = float(tmp[15].split(':')[1])
            energies_blocks.append(energies)

        out = {}
        out['Species'] = species_blocks
        out['Energies'] = energies_blocks
        out['charge'] = charge
        out['states'] = states
        out['atom'] = atom
        out['cell'] = cell
        out['soc'] = soc
        out['spin'] = spin

        return out

def read_kpoints(filename):
    '''Read the KPOINTS.OUT file'''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, "r") as data_file:
        data = data_file.read()
        tmp = np.float64(" ".join(data.split('\n')[1:-1]).split()).reshape(-1, 6)
        return tmp[:,1:4], tmp[:,4]           # kpts and kpts weight

eigval_MATCH = re.compile(r'''
[ ]*
 \([ ]* state, [ ]* eigenvalue [ ]* and [ ]* occupancy [ ]* below\)
(?P<content>
  [\s\S]*?(?=\n.*?[ ] $|\n[ ]\n)  # match everything until next blank line or EOL
)
''', re.VERBOSE)

def read_eigval(filename):
    '''Read the KPOINTS.OUT file
       Return:
        eigvals[spin, kpt, band, eigenvalues]
        occ[spin, kpt, band, occupancy]
    '''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, "r") as data_file:
        data = data_file.read()

        eigvals = []
        occ = []
        for eigval in eigval_MATCH.finditer(data):
            eigval_data = eigval['content']
            tmp = np.float64(eigval_data.split()).reshape(-1, 3)
            eigvals.append(tmp[:,1])
            occ.append(tmp[:,2])

        eigvals = np.asarray(eigvals)
        nkpts = eigvals.shape[0]
        nband = eigvals.shape[1]
        occ = np.asarray(occ)

        # Check if this is a spin-polarized calculation
        if (occ[0] <= 1.0).any() and abs((-occ[0]).argsort() - np.arange(occ[0].shape[0])).sum() > 1e-3:
            assert np.mod(nband,2) == 0, "This is not a spin-polarised calculation"
            nband = nband//2

        eigvals = eigvals.reshape(nkpts, nband, -1).transpose(2, 0, 1)
        occ = occ.reshape(nkpts, nband, -1).transpose(2, 0, 1)
        return eigvals, occ

############ Read BAND***.OUT - begin ##################
def read_bandlines(filename):
    '''Read the BANDLINES.OUT file
    '''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, "r") as data_file:
        data = np.float64(data_file.read().split()).reshape(-1,2)
        nkpoint = data.shape[0] // 2
        return data[2*np.arange(nkpoint), 0] / const.AUTOA

def read_band(filename, spin_polarized=False):
    '''Read BAND.OUT or BAND_Sss_Aaaaa.OUT

       Depends on the tasks, BAND.OUT or BAND_Sss_Aaaaa.OUT contains different projection
       For tasks 20:    kpt  |  e (a.u.)
       For tasks 21:    kpt  |  e (a.u.)  |  Total  |   s   |   p   |   d   |   f   |
       For tasks 22:    kpt  |  e (a.u.)  |   s     |  p(-1,0,+1)  |  d(-2,-1,0,+1,+2)  |  f(-3,-2,-1,0,+1,+2,+3)
       For tasks 23:    kpt  |  e (a.u.)  |   up    |  down |

       Return:
                bands[spin, kpt, band-th, projected contribution]
    '''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, 'r') as data_file:
        data = data_file.read()
        tmp = data.split('\n \n')[:-1]

        # Figure out how many kpt and how many column
        nband = len(tmp)
        tmp2 = tmp[0].split('\n')
        nkpts = len(tmp2)
        ncol = len(tmp2[0].split())
        bands = []
        for i, band in enumerate(tmp):
            formatted_band = np.float64(band.split()).reshape(-1, ncol)
            if i == 0: proj_kpath = formatted_band[:,0] / const.AUTOA           # in Angstrom-1
            bands.append(formatted_band[:,1:])

        if spin_polarized:
            assert np.mod(nband,2) == 0, "This is not a spin-polarised calculation"
            nband = nband//2

        bands = np.asarray(bands).reshape(-1, nband, nkpts, ncol - 1).transpose(0, 2, 1, 3) * const.AUTOEV
        return proj_kpath, bands

def read_plot1d(filename):
    '''Read the plot1d of elk.in'''

    with open(filename, 'r') as data_file:
        data = data_file.read()
        assert 'plot1d' in data, "plot1d block cannot be found in the elk.in"
        tmp_with_comments = data.split('\n')
        tmp = [item for item in tmp_with_comments if not '!' in item]   # remove comments
        for i, item in enumerate(tmp):
            if 'plot1d' in item:
                where_plot1d = i
                break

        nkpoint, npts = np.int64(tmp[where_plot1d + 1].split()[:2])
        sym_kvecs = np.float64([tmp[i].split()[:3] for i in range(where_plot1d + 2, where_plot1d + nkpoint + 2)])

        return sym_kvecs, npts

############ Read DOS - begin ##################
def read_dos(filename, spin_polarized=False):
    '''Read TDOS.OUT, IDOS.OUT, PDOS_Sss_Aaaaa.OUT

       For PDOS_Sss_Aaaaa.OUT:
       e (a.u.)  |   s     |  p(-1,0,+1)  |  d(-2,-1,0,+1,+2)  |  f(-3,-2,-1,0,+1,+2,+3)
       Return:
                dos[spin, e, dos]
    '''
    assert check_exist(filename), 'Cannot find : ' + filename
    with open(filename, 'r') as data_file:
        data = data_file.read()
        tmp = data.split('\n \n')[:-1]

        # Figure out how many kpt and how many column
        ndos = len(tmp)
        nepsilon = len(tmp[0].split('\n'))
        dos = []
        for i, block in enumerate(tmp):
            formatted_dos = np.float64(block.split()).reshape(nepsilon, 2)
            if i == 0: epsilon = formatted_dos[:,0] * const.AUTOEV          # in eV
            dos.append(formatted_dos[:,1])

        if spin_polarized:
            assert np.mod(ndos,2) == 0, "This is not a spin-polarised calculation"
            ndos = ndos//2

        dos = np.asarray(dos).reshape(-1, ndos, nepsilon).transpose(0, 2, 1)
        return epsilon, dos


#!/usr/bin/env python
# encoding: utf-8
"""
Emcee-based sampler for singlepix SED MCMC.
"""

import time
import numpy as np
from collections import OrderedDict

import emcee
import fsps
from astropy.table import Table, hstack, Column

from sedbot.photconv import micro_jy_to_luminosity
from sedbot.chain import SinglePixelChain

class SinglePixelSampler(object):
    """Emcee-based sampler for single pixel SEDs.

    Parameters
    ----------
    model : object
        A model instance (see :mod:`sedbot.models`).
    n_walkers : int
        Number of emcee walkers.
    """
    def __init__(self, model, n_walkers=100):
        super(SinglePixelSampler, self).__init__()
        self.model = model
        self.n_walkers = n_walkers
        self._sampler = None
        self._run_time = 0.
        self._call_time = np.nan

    def generate_initial_point(self, x0, x_sigma):
        """Generate a chain starting point given an array of Gaussian
        dispersions to build balls from for each parameter.

        Parameters
        ----------
        x0 : ndarray
            Initial points in parameter space. Shape ``(n_params,)``.
        x_sigma : ndarray
            Gaussian standard deviations to disperse starting points for each
            walkers. Shape ``(n_params,)``.
        """
        ndim = len(x0)
        assert len(x0) == len(x_sigma), "Lengths do not match"
        assert ndim == self.model.n_params
        p0 = np.random.randn(ndim * self.n_walkers).reshape((self.n_walkers,
                                                             ndim))
        for i, (x, sigma) in enumerate(zip(x0, x_sigma)):
            p0[:, i] = sigma * p0[:, i] + x
            priorf = self.model.priors[self.model.param_names[i]]
            p0[:, i][p0[:, i] < priorf.lower_limit] = priorf.lower_limit
            p0[:, i][p0[:, i] > priorf.upper_limit] = priorf.upper_limit
        return p0

    def sample(self, n_steps, theta0=None, chain=None):
        """Sample for `n_steps` Emcee steps.

        Parameters
        ----------
        n_steps : int
            Number of emcee steps. The total number of samples will be
            ``n_steps x n_walkers``.
        theta0 : ndarray
            Initial locations of the walkers; shape
        """
        self.sampler = emcee.EnsembleSampler(
            self.n_walkers,
            self.model.n_params, self.model)

        # Do burn-in + run
        with EmceeTimer(n_steps, self.n_walkers) as emcee_timer:
            self.sampler.run_mcmc(theta0, n_steps)
        print(emcee_timer)

        self._run_time += emcee_timer.interval
        self._call_time = emcee_timer.seconds_per_call

    @property
    def table(self):
        """An :class:`astropy.table.Table` with the chain."""
        if self.sampler is None:
            return None
        msuns = np.array([fsps.get_filter(n).msun_ab
                          for n in self.model.computed_bands])
        meta = OrderedDict((
            ('observed_bands', self.model.observed_bands),
            ('instruments', self.model.instruments),
            ('computed_bands', self.model.computed_bands),
            ('msun_ab', msuns),
            ('d', self.model.d),  # expected distance in parsecs
            ('band_indices', self.model.band_indices),
            ('theta_params', self.model.param_names),
            ('compute_time', self._run_time),
            ('step_time', self._call_time),
            ('sed', self.model._sed),
            ('sed_err', self.model._err),
            ('pixels', self.model.pixel_metadata),
            ('area', self.model._area),
            ('n_walkers', self.n_walkers),
            ("f_accept", self.sampler.acceptance_fraction),
            ("acor", self.sampler.acor)))

        # Convert flatchain into a structured array
        nwalkers, nsteps, ndim = self.sampler.chain.shape
        flatchain_arr = self.sampler.chain[:, :, :].reshape((-1, ndim))
        dt = [(n, np.float) for n in self.model.param_names]
        flatchain = np.empty(flatchain_arr.shape[0], dtype=np.dtype(dt))
        for i, n in enumerate(self.model.param_names):
            flatchain[n][:] = flatchain_arr[:, i]

        # Flatten the blob list and make a structured array
        blobs = self.sampler.blobs
        blobchain = np.empty(nwalkers * nsteps, self.model.blob_dtype)
        blobchain.fill(np.nan)
        for i in xrange(nsteps):
            for j in xrange(nwalkers):
                for k in self.model.blob_dtype.names:
                    blobchain[k][i * self.n_walkers + j] = blobs[i][j][k]

        chain_table = Table(flatchain, meta=meta)
        blob_table = Table(blobchain)
        tbl = SinglePixelChain(hstack((chain_table, blob_table),
                                      join_type='exact'))

        # Add M/L computations for each computed band.
        for i, (band_name, msun) in enumerate(zip(self.model.computed_bands,
                                                  msuns)):
            # Either use expected distance or the chain distance
            # FIXME fragile code
            if 'd' in self.model.param_names:
                d = np.array(tbl['d'])
            else:
                d = self.model.d
            logLsol = micro_jy_to_luminosity(tbl['model_sed'][:, i], msun, d)
            ml = tbl['logMstar'] - logLsol
            colname = "logML_{0}".format(band_name)
            tbl.add_column(Column(name=colname, data=ml))

        return tbl

class EmceeTimer(object):
    """Timer for emcee runs. Computes total run time and mean time of each
    likelihood call.

    Example::

    >>> with EmceeTimer(n_steps, n_walkers) as emceet:
    >>>     sampler.run_mcmc(p0, n_steps)
    >>> print emceet

    Parameters
    ----------
    nsteps : int
        Number of emcee steps.
    nwalkers : int
        Number of emcee walkers.
    """
    def __init__(self, nsteps, nwalkers):
        super(EmceeTimer, self).__init__()
        self._nsteps = nsteps
        self._nwalkers = nwalkers
        self._start = None
        self._stop = None
        self._endtime = None
        self._interval = None

    def __enter__(self):
        self._start = time.clock()
        return self

    def __exit__(self, *args):
        self._end = time.clock()
        self._endtime = time.localtime()
        self._interval = self._end - self._start

    def __str__(self):
        dt, unit = self._human_time(self._interval)
        ct, cunit = self._human_time(self.seconds_per_call)
        enddate = time.strftime('%Y-%m-%d %H:%M:%S', self._endtime)
        l1 = "Finished emcee run at {enddate}"
        l2 = "\tRun time: {dt} {unit}"
        l3 = "\t{ct} {cunit} per call"
        return "\n".join((l1, l2, l3)).format(dt=dt, unit=unit,
                                              ct=ct, cunit=cunit,
                                              enddate=enddate)

    def _human_time(self, interval):
        """Return a time interval scaled to human-usable units."""
        if interval < 60.:
            dt, unit = interval, "seconds"
        if interval >= 60.:
            dt, unit = interval / 60., "minutes"
        elif interval >= 3600.:
            dt, unit = interval / 3600., "hours"
        elif interval >= 86400:
            dt, unit = interval / 86400., "days"
        return dt, unit

    @property
    def seconds_per_call(self):
        """The mean number of seconds elapsed per likelihood call.

        Returns
        -------
        interval : float
            Description
        """
        return self._interval / float(self._nsteps * self._nwalkers)

    @property
    def interval(self):
        """The timer interval in secods."""
        return self._interval


import numpy
def BREMSCASC(J11,EGAMMA,X0,Y0,Z0,T0,GDCX,GDCY,GDCZ,ILOW):
	# IMPLICIT #real*8(A-H,O-Z)
	# IMPLICIT #integer*8(I-N)
	#CHARACTER*6 SCR(17),SCR1(17)
	# COMMON/COMP/LCMP,LCFLG,LRAY,LRFLG,LPAP,LPFLG,LBRM,LBFLG,LPEFLG
	# COMMON/COMPTOUT/EGAM,EELEC,THETAG,THETAE
	# COMMON/PRIM4/MSUM1,MCOMP1,MRAYL1,MPAIR1,MPHOT1,MVAC1
	# COMMON/GENCAS/ELEV[17,79],NSDEG(17),AA[17],BB[17],SCR,SCR1
	# COMMON/MIXC/PRSH(6,3,17,17),ESH(6,3,17),AUG(6,3,17,17,17),RAD[6,3,17,17],PRSHBT(6,3,17),IZ[6,3],INIOCC(6,3,17),ISHLMX(6,3),AMZ[6,3]
	# COMMON/UPD/NOCC(6,3,17),AUGR(6,3,17,17,17),RADR(6,3,17,17)
	# COMMON/CALCASB/IONSUM(10),IFLSUM(10),ESTORE(10,28),EPHOTON(10,28),DRXE(10,28),DRYE(10,28),DRZE(10,28),DRX(10,28),DRY(10,28),DRZ[10,28]
	# COMMON/CALCAS1B/IONSUM1(10),IFLSUM1(10),ESTOR1(10,28),EPHOTG1(10,28),DRXE1(10,28),DRYE1(10,28),DRZE1(10,28),DRX1(10,28),DRY1(10,28),DRZ1(10,28)
	# COMMON/CALCAS2B/IONSUM2(10),IFLSUM2(10),ESTOR2(10,28),EPHOTG2(10,28),DRXE2(10,28),DRYE2(10,28),DRZE2(10,28),DRX2(10,28),DRY2(10,28),DRZ2(10,28)
	# COMMON/CALCAS3B/IONSUM3(10),IFLSUM3(10),ESTOR3(10,28),EPHOTG3(10,28),DRXE3(10,28),DRYE3(10,28),DRZE3(10,28),DRX3(10,28),DRY3(10,28),DRZ3(10,28)
	# COMMON/CALCAS4B/IONSUM4(10),IFLSUM4(10),ESTOR4(10,28),EPHOTG4(10,28),DRXE4(10,28),DRYE4(10,28),DRZE4(10,28),DRX4(10,28),DRY4(10,28),DRZ4(10,28)
	# COMMON/CALCAS5B/IONSUM5(10),IFLSUM5(10),ESTOR5(10,28),EPHOTG5(10,28),DRXE5(10,28),DRYE5(10,28),DRZE5(10,28),DRX5(10,28),DRY5(10,28),DRZ5(10,28)
	# COMMON/RESB/IONSM(10),IFLSM(10),ESTOR(10,28),EPHOT(10,28),X10(10,28),Y10(10,28),Z10(10,28),DRX01(10,28),DRY01(10,28),DRZ01(10,28)
	# COMMON/GENB1/IONF1(10),ESTF1(10,28),X11(10,28),Y11(10,28),Z11(10,28),DRX11(10,28),DRY11(10,28),DRZ11(10,28)
	# COMMON/GENB2/IONF2(10),ESTF2(10,28),X21(10,28),Y21(10,28),Z21(10,28),DRX21(10,28),DRY21(10,28),DRZ21(10,28)
	# COMMON/GENB3/IONF3(10),ESTF3(10,15),X31(10,15),Y31(10,15),Z31(10,15),DRX31(10,15),DRY31(10,15),DRZ31(10,15)
	# COMMON/GENB4/IONF4(10),ESTF4(10,12),X41(10,12),Y41(10,12),Z41(10,12),DRX41(10,12),DRY41(10,12),DRZ41(10,12)
	# COMMON/GENB5/IONF5(10),ESTF5(10,5),X51(10,5),Y51(10,5),Z51(10,5),DRX51(10,5),DRY51(10,5),DRZ51(10,5)
	# COMMON/COUTE/ECMP(10),ECDRX(10),ECDRY(10),ECDRZ[10],XCPOS(10),YCPOS(10),ZCPOS(10),KCGAS(10),LCGAS(10),ICSHELL(10)
	# COMMON/COUTTB/TT(10),TTP
	# COMMON/PPSTRB/NPTP,EPPST[2],XPP[2],YPP[2],ZPP[2],DRXPP[2],DRYPP[2],DRZPP[2]
	# COMMON/COMP/
	global LCMP,LCFLG,LRAY,LRFLG,LPAP,LPFLG,LBRM,LBFLG,LPEFLG
	#COMMON/COMPTOUT/
	global EGAM,EELEC,THETAG,THETAE
	#COMMON/PRIM4/
	global MSUM1,MCOMP1,MRAYL1,MPAIR1,MPHOT1,MVAC1
	#COMMON/GENCAS/
	global ELEV#(17,79)
	global NSDEG#(17)
	global AA#(17)
	global BB#(17)
	global SCR,SCR1
	#COMMON/MIXC/
	global PRSH#(6,3,17,17)
	global ESH#(6,3,17)
	global AUG#(6,3,17,17,17)
	global RAD#(6,3,17,17)
	global PRSHBT#(6,3,17)
	global IZ#(6,3)
	global INIOCC#(6,3,17)
	global ISHLMX#(6,3)
	global AMZ#(6,3)
	#COMMON/UPD/
	global NOCC#(6,3,17)
	global AUGR#(6,3,17,17,17)
	global RADR#(6,3,17,17)
	#COMMON/CALCASB/
	global IONSUM#(10)
	global IFLSUM#(10)
	global ESTORE#(10,28)
	global EPHOTON#(10,28)
	global DRXE#(10,28)
	global DRYE#(10,28)
	global DRZE#(10,28)
	global DRX#(10,28)
	global DRY#(10,28)
	global DRZ#(10,28)
	#COMMON/CALCAS1B/
	global IONSUM1#(10)
	global IFLSUM1#(10)
	global ESTOR1#(10,28)
	global EPHOTG1#(10,28)
	global DRXE1#(10,28)
	global DRYE1#(10,28)
	global DRZE1#(10,28)
	global DRX1#(10,28)
	global DRY1#(10,28)
	global DRZ1#(10,28)
	#COMMON/CALCAS2B/
	global IONSUM2#(10)
	global IFLSUM2#(10)
	global ESTOR2#(10,28)
	global EPHOTG2#(10,28)
	global DRXE2#(10,28)
	global DRYE2#(10,28)
	global DRZE2#(10,28)
	global DRX2#(10,28)
	global DRY2#(10,28)
	global DRZ2#(10,28)
	#COMMON/CALCAS3B/
	global IONSUM3#(10)
	global IFLSUM3#(10)
	global ESTOR3#(10,28)
	global EPHOTG3#(10,28)
	global DRXE3#(10,28)
	global DRYE3#(10,28)
	global DRZE3#(10,28)
	global DRX3#(10,28)
	global DRY3#(10,28)
	global DRZ3#(10,28)
	#COMMON/CALCAS4B/
	global IONSUM4#(10)
	global IFLSUM4#(10)
	global ESTOR4#(10,28)
	global EPHOTG4#(10,28)
	global DRXE4#(10,28)
	global DRYE4#(10,28)
	global DRZE4#(10,28)
	global DRX4#(10,28)
	global DRY4#(10,28)
	global DRZ4#(10,28)
	#COMMON/CALCAS5B/
	global IONSUM5#(10)
	global IFLSUM5#(10)
	global ESTOR5#(10,28)
	global EPHOTG5#(10,28)
	global DRXE5#(10,28)
	global DRYE5#(10,28)
	global DRZE5#(10,28)
	global DRX5#(10,28)
	global DRY5#(10,28)
	global DRZ5#(10,28)
	#COMMON/RESB/
	global IONSM#(10)
	global IFLSM#(10)
	global ESTOR#(10,28)
	global EPHOT#(10,28)
	global X10#(10,28)
	global Y10#(10,28)
	global Z10#(10,28)
	global DRX01#(10,28)
	global DRY01#(10,28)
	global DRZ01#(10,28)
	#COMMON/GENB1/
	global IONF1#(10)
	global ESTF1#(10,28)
	global X11#(10,28)
	global Y11#(10,28)
	global Z11#(10,28)
	global DRX11#(10,28)
	global DRY11#(10,28)
	global DRZ11#(10,28)
	#COMMON/GENB2/
	global IONF2#(10)
	global ESTF2#(10,28)
	global X21#(10,28)
	global Y21#(10,28)
	global Z21#(10,28)
	global DRX21#(10,28)
	global DRY21#(10,28)
	global DRZ21#(10,28)
	#COMMON/GENB3/
	global IONF3#(10)
	global ESTF3#(10,15)
	global X31#(10,15)
	global Y31#(10,15)
	global Z31#(10,15)
	global DRX31#(10,15)
	global DRY31#(10,15)
	global DRZ31#(10,15)
	#COMMON/GENB4/
	global IONF4#(10)
	global ESTF4#(10,12)
	global X41#(10,12)
	global Y41#(10,12)
	global Z41#(10,12)
	global DRX41#(10,12)
	global DRY41#(10,12)
	global DRZ41#(10,12)
	#COMMON/GENB5/
	global IONF5#(10)
	global ESTF5#(10,5)
	global X51#(10,5)
	global Y51#(10,5)
	global Z51#(10,5)
	global DRX51#(10,5)
	global DRY51#(10,5)
	global DRZ51#(10,5)
	#COMMON/COUTE/
	global ECMP#(10)
	global ECDRX#(10)
	global ECDRY#(10)
	global ECDRZ#(10)
	global XCPOS#(10)
	global YCPOS#(10)
	global ZCPOS#(10)
	global KCGAS#(10)
	global LCGAS#(10)
	global ICSHELL#(10)
	#COMMON/COUTTB/
	global TT#(10)
	global TTP
	#COMMON/PPSTRB/
	global NPTP
	global EPPST#(2)
	global XPP#(2)
	global YPP#(2)
	global ZPP#(2)
	global DRXPP#(2)
	global DRYPP#(2)
	global DRZPP#(2)
	#----------------------------------------------------------------------
	# BREMSSTRAHLUNG CASCADE TREE:
	#   SET OR ZERO SOME VARIABLES
	#   CREATE INTERACTION TREE FOR PE COMPTON RAYLEIGH AND PAIR PRODUCTION
	#   STORE ELECTRON ENERGY DIRECTION COSINES AND POSITION WITH SHELL
	#   LEVEL AND GAS IDENTITY IN COMMON/COUT/ FOR EACH COMPTON AND PE EVENT
	#   STORE PAIR PRODUCTION ELECTRON AND POSITRON DATA IN COMMON/PPSTR/
	#   STORE FINAL CASCADE RESULTS IN COMMON/CASRSB/
	#----------------------------------------------------------------------
	API=numpy.arccos(-1.00)
	TWOPI=2.00*API
	IDBG=0
	#     IF(J11 == 2) IDBG=1
	def GOTO5():
		ILOW=0
		# USE VELOCITY IN METRES/PICOSECONDS
		VV=2.99792458E-4
		#
		LFIX=0
		# ALLOW FLUORESCENCE CALCULATION
		def GOTO123():
			ICONPH=1
			#  PHOTON ORIGIN
			X=X0
			Y=Y0
			Z=Z0
			TSUM=T0
			#  LOAD INITIAL DIRECTION COSINES BEFORE INTERACTION
			DRXS=GDCX
			DRYS=GDCY
			DRZS=GDCZ
			#   INITIAL ENERGY
			ENERGY=EGAMMA
			#
			# ZERO SOME ARRAYS
			EPPST[1]=0.00
			EPPST[2]=0.00
			for K in range(1,10):
				for J in range(1,28):
					EPHOTON[K][J]=0.0
					EPHOTG1[K][J]=0.0
					EPHOTG2[K][J]=0.0
					EPHOTG3[K][J]=0.0
					EPHOTG4[K][J]=0.0
					EPHOTG5[K][J]=0.0
					ESTORE[K][J]=0.0
					ESTOR1[K][J]=0.0
					ESTOR2[K][J]=0.0
					ESTOR3[K][J]=0.0
					ESTOR4[K][J]=0.0
					ESTOR5[K][J]=0.0
					ESTOR[K][J]=0.0
					ESTF1[K][J]=0.0
					ESTF2[K][J]=0.0
				for J in range(1,15):
					ESTF3[K][J]=0.0
				for J in range(1,12):
					ESTF4[K][J]=0.0
				for J in range(1,5):
					ESTF5[K][J]=0.0
				IFLSUM[K]=0
				IFLSUM1[K]=0
				IFLSUM2[K]=0
				IFLSUM3[K]=0
				IFLSUM4[K]=0
				IFLSUM5[K]=0
				IONSUM[K]=0
				IONSUM1[K]=0
				IONSUM2[K]=0
				IONSUM3[K]=0
				IONSUM4[K]=0
				IONSUM5[K]=0
				IONSM[K]=0
				IONF1[K]=0
				IONF2[K]=0
				IONF3[K]=0
				IONF4[K]=0
				IONF5[K]=0
				KCGAS[K]=0
				LCGAS[K]=0
				ICSHELL[K]=0
				ECMP[K]=0.0
			NCOMP=0
			NRAYL=0
			NPAIR=0
			NPHOT=0
			NVAC=0
			NPTP=0
			#
			IFIRST=1
			ISECOND=2
			def GOTO2():
				NTOTI=NCOMP+NRAYL+NPAIR+NPHOT
				# PHOTONS
				ABSO(IFIRST,ENERGY,ISHELL,KGAS,LGAS,DIST)
				#     IF(IDBG == 1) :
				#     WRITE(6,888) ENERGY,ISHELL,J11
				# 888 print(' AFTER ABSO ENERGY=',D12.5,' ISHELL=',I3,' EVENT NO=',I4)
				#     # endIF
				if(ISHELL == -1):
					# BREMSSTRAHLUNG GAMMA TOO LOW IN ENERGY TO IONISE
					ILOW=1
					return
				# endif
				#
				#  CREATE INTERACTION TREE
				flag=0
				if(LPEFLG == 1):
					flag=100
				if(LCFLG == 1):
					flag=10			###############################
				elif(LRFLG == 1):	##							 ##
					flag=20			## Made significant changes  ##
				elif(LPFLG == 1):	##							 ##
					flag=30			###############################
				# COMPTON SCATTERING
				if(flag==10 or flag==0):
					COMPTON(KGAS,LGAS,ENERGY)
					NCOMP=NCOMP+1
					NVAC=NVAC+1
					if(NVAC > 10):
						# MAXIMUM OF 10 PRIMARY INTERACTIONS
						#      NJHIGH=NJHIGH+1
						GOTO123()
					# endif
					# RANDOMISE ANGLE PHI
					R3=DRAND48(RDUM)
					PHI=TWOPI*R3
					# CALCULATE COMPTON ELECTRON DIRECTION COSINES USING THETAE AND PHI
					DRCOS(DRXS,DRYS,DRZS,THETAE,PHI,DRXX,DRYY,DRZZ)
					# FOR COMPTON EFFECT STORE ELECTRON DIRECTION COSINES
					ECDRX[NVAC]=DRXX
					ECDRY[NVAC]=DRYY
					ECDRZ[NVAC]=DRZZ
					# CALCULATE WHICH SHELL HAS VACANCY FROM COMPTON EVENT
					CVAC(KGAS,LGAS,EELEC,KSHELL,KBAD)
					# REJECT EVENT WITHOUT SUFFICIENT ENERGY TO IONISE SHELLS
					if(KBAD == 1):
						GOTO123()
					# STORE ELECTRON ENERGY SHELL VACANCY ISHELL KGAS LGAS POSITION
					ECMP[NVAC]=EELEC
					ICSHELL[NVAC]=KSHELL
					KCGAS[NVAC]=KGAS
					LCGAS[NVAC]=LGAS
					# INTERACTION POSITIONS
					XCPOS[NVAC]=X+DIST*DRXS
					YCPOS[NVAC]=Y+DIST*DRYS
					ZCPOS[NVAC]=Z+DIST*DRZS
					TT[NVAC]=TSUM+DIST/VV
					# UPDATE PHOTON STARTING ENERGY POSITION AND ANGLES
					ENERGY=EGAM
					X=XCPOS[NVAC]
					Y=YCPOS[NVAC]
					Z=ZCPOS[NVAC]
					PHIG=PHI+API
					if(PHIG >= TWOPI):
						PHIG=PHI-API
					DRCOS(DRXS,DRYS,DRZS,THETAG,PHIG,DRXX,DRYY,DRZZ)
					# NEW DIRECTION COSINES
					DRXS=DRXX
					DRYS=DRYY
					DRZS=DRZZ
					# LOOP BACK
					GOTO2()

				if(flag==20 or flag==0):
					# RAYLEIGH SCATTERING
					RAYLEIGH(KGAS,LGAS,ENERGY,THETAR)
					NRAYL=NRAYL+1
					#  CALCULATE ENERGY LOSS IN RAYLEIGH SCATTERING
					RAYLOS(KGAS,LGAS,ENERGY,THETAR,ELRAY)
					#     IF(IDBG == 1) WRITE(6,776) ENERGY,ELRAY
					# 776 print(' AFTER RAYLOS ENERGY=','%.4f' % ,' ELRAY=','%.4f' % )
					# UPDATE X-RAY STARTING ENERGY POSITION AND ANGLES
					ENERGY=ENERGY-ELRAY
					X=X+DIST*DRXS
					Y=Y+DIST*DRYS
					Z=Z+DIST*DRZS
					TSUM=TSUM+DIST/VV
					# RANDOMISE ANGLE PHI
					R3=DRAND48(RDUM)
					PHIR=TWOPI*R3
					DRCOS(DRXS,DRYS,DRZS,THETAR,PHIR,DRXX,DRYY,DRZZ)
					DRXS=DRXX
					DRYS=DRYY
					DRZS=DRZZ
					# LOOP BACK
					GOTO2()

				if(flag==30 or flag==0):
					# PAIR PRODUCTION
					PAIR(KGAS,LGAS,ENERGY,E1,E2,THET1,PHI1,THET2,PHI2)
					NPAIR=NPAIR+1
					if(NPAIR > 2):
						print(' ERROR NPAIR GT 2 =',NPAIR,' IN BREMSCASC EVENT NO =',J11)
						sys.exit()
					# endif
					NPTP=NPAIR
					# STORE ELECTRON AND POSITRON ENERGY POSITION AND ANGLES
					EPPST[1]=E1
					EPPST[2]=E2
					if(NVAC == 0):
						# FIRST INTERACTION IS PAIR PRODUCTION
						XPP[1]=X
						YPP[1]=Y
						ZPP[1]=Z
						XPP[2]=X
						YPP[2]=Y
						ZPP[2]=Z
						TTP=TSUM
						pass
					# endif
					else:
						XPP[1]=X+DIST*DRXS
						YPP[1]=Y+DIST*DRYS
						ZPP[1]=Z+DIST*DRZS
						XPP[2]=XPP[1]
						YPP[2]=YPP[1]
						ZPP[2]=ZPP[1]
						TTP=TSUM+DIST/VV
					DRCOS(DRXS,DRYS,DRZS,THET1,PHI1,DRXX,DRYY,DRZZ)
					DRXPP[1]=DRXX
					DRYPP[1]=DRYY
					DRZPP[1]=DRZZ
					DRCOS(DRXS,DRYS,DRZS,THET2,PHI2,DRXX,DRYY,DRZZ)
					DRXPP[2]=DRXX
					DRYPP[2]=DRYY
					DRZPP[2]=DRZZ
					flag=200
					# PHOTOELECTRIC ABSORPTION
					#  STORE ENERGY ISHELL KGAS
				if(flag==100 or flag==0):
					NVAC=NVAC+1
					if(NVAC > 10):
						# ONLY ALLOW MAXIMUM OF 10 PRIMARY INTERACTIONS
						GOTO123()
						sys.exit()
					# endif

					NPHOT=NPHOT+1
					# ECMP= TOTAL ENERGY= EGAMMA = ELECTRON KINETIC ENERGY+ VACANCY ENERGY
					ECMP[NVAC]=ENERGY
					ICSHELL[NVAC]=ISHELL
					XCPOS[NVAC]=X+DIST*DRXS
					YCPOS[NVAC]=Y+DIST*DRYS
					ZCPOS[NVAC]=Z+DIST*DRZS
					TT[NVAC]=TSUM+DIST/VV
					KCGAS[NVAC]=KGAS
					LCGAS[NVAC]=LGAS
					# FOR PE EFFECT STORE PHOTON INCIDENT ANGLE
					ECDRX[NVAC]=DRXS
					ECDRY[NVAC]=DRYS
					ECDRZ[NVAC]=DRZS
				# LOOP OVER SHELL VACANCIES
				if(flag==200 or flag==0):
					# STORE NUMBER AND TYPE OF PRIMARY INTERACTIONS
					MSUM1=NTOTI
					MCOMP1=NCOMP
					MRAYL1=NRAYL
					MPAIR1=NPAIR
					MPHOT1=NPHOT
					MVAC1=NVAC
					# LOOP OVER SHELL INTERACTIONS
					#     IF(IDBG == 1) :
					#     WRITE(6,54) NVAC
					#  54 print(' NVAC=',I3)
					#    for 6 in range(M1=1,NVAC):
					#     WRITE(6,55) M1, (ESTORE(M1,K1),K1=1,28)
					#  55 print(' M1=',I3,/,' ESTORE(M1,K)=',4(7'%.4f' % ,/))
					#  56 CONTINUE
					#    for 8 in range(M1=1,NVAC):
					#     WRITE(6,57) (EPHOTON(M1,K1),K1=1,28)
					#  57 print(' EPHOTON=',4(7'%.4f' % ,/))
					#  58 CONTINUE
					#     # endIF
					for K in range(1,NVAC):
						CONTROLB(IDBG,K)
					#  COMPRESS AUGER AND FLUORESCENCE DATA INTO BLOCKS
					COMPRESS(IDBG,ENSUM)
					# CATCH DROPPED KSHELL FLUORESCENCE
					if(abs(ENSUM-EGAMMA)> 2200.) :
						DIF=ENSUM-EGAMMA
					if(IDBG == 1):
						print('\n Dif=','%.6f' % DIF,' J11=',J11,'\n')
						GOTO5()
					# endif
					# ADD PAIR DATA AND LOAD INTO COMMON/CASRSB/
					CASRESB()
			GOTO2()
		GOTO123()
	GOTO5()
	return
	# end


# Standard library
import os

# Third-party
import numpy as np
from numpy.lib.recfunctions import append_fields
from scipy.integrate import quad
from scipy.interpolate import interp1d
from astropy.table import Table
from astropy import units as u

# Project
from ._read_mist_models import IsoCmdReader, IsoReader
from .imf import imf_dict, IMFIntegrator
from .. import MIST_PATH
from ..log import logger
from ..filters import phot_system_list, get_filter_names
from ..filters import load_zero_point_converter
from ..util import check_units, fetch_mist_grid_if_needed
phot_str_helper = {p.lower():p for p in phot_system_list}

__all__ = ['fetch_mist_iso_cmd', 'Isochrone', 'MISTIsochrone']

class Isochrone(object):
    """
    Class for storing generic isochrones. It also has methods for calculating
    IMF-weighted properties of a simple stellar population with the given
    age and metallicity.

    Parameters
    ----------
    mini : `~numpy.ndarray`
        Initial stellar masses. The masses must be monotonically increasing.
    mact : `~numpy.ndarray`
        Actual stellar masses after mass loss.
    mag_table : `~astropy.table.Table`, dict, or structured `~numpy.ndarray`
        The stellar magnitudes.
    eep : `~numpy.ndarray`, optional
        The Primary Equivalent Evolutionary Points. Needed to identify phases
        of stellar evolution.
    log_L : `~numpy.ndarray`, optional
        Stellar luminosities.
    log_Teff : `~numpy.ndarray`, optional
        Stellar effective temperatures.
    """

    def __init__(self, mini, mact, mags, eep=None, log_L=None,
                 log_Teff=None):
        self.mini = np.asarray(mini)
        self.mact = np.asarray(mact)
        if (np.diff(self.mini) < 0).sum() > 0:
            raise Exception('Initial masses must be monotonically increasing.')
        self.eep = None if eep is None else np.asarray(eep)
        self.log_L = None if log_L is None else np.asarray(log_L)
        self.log_Teff = None if log_Teff is None else np.asarray(log_Teff)
        if type(mags) == dict or type(mags) == np.ndarray:
            self.mag_table = Table(mags)
        elif type(mags) == Table:
            self.mag_table = mags
        else:
            raise Exception(f'{type(mags)} is not a valid type for mags')

    @property
    def m_min(self):
        """The minimum mass of the isochrone."""
        return self.mini.min()

    @property
    def m_max(self):
        """The maximum mass of the isochrone."""
        return self.mini.max()

    @property
    def filters(self):
        """List of filters in the given photometric system(s)."""
        return self.mag_table.colnames

    @staticmethod
    def from_parsec(file_name, log_age=None, zini=None, num_filters=None):
        """
        Read isochrone generated from the `PARSEC website
        <http://stev.oapd.inaf.it/cgi-bin/cmd>`_. If more than one age and/or
        metallicity is included in the file, you must provide the log_age
        and/or zini parameters.

        Parameters
        ----------
        file_name : str
            Isochrone file name.
        log_age : float, optional
            Log of age in years. You must provided this parameter if there is
            more than one age in the isochrone file. Note this function does
            not interpolate ages, so it must be included in the file.
        zini : float, optional
            Initial metal fraction. You must provided this parameter if there
            is more than one metallicity in the isochrone file. Note this
            function does not interpolate metallicity, so it must be included
            in the file.
        num_filters : int, optional
            Number of filters included in the isochrone file. If None, will
            assume the last non-filter parameter is `mbolmag`.

        Returns
        -------
        iso : `artpop.stars.Isochrone`
            PARSEC Isochrone object.
        """
        if os.path.isfile(file_name):
            file = open(file_name, 'r')
            lines = file.readlines()
            file.close()
            for l in lines:
                if 'Zini' in l.split()[1]:
                    names = l.split()[1:]
                    break
            data = np.loadtxt(file_name)
            parsec = Table(data, names=names)
        else:
            raise Exception(f'{file_name} does not exist.')
        isochrone_full = parsec.copy()
        if log_age is not None:
            age_cut = np.abs(parsec['logAge'] - log_age) < 1e-5
            if age_cut.sum() < 1:
                raise Exception(f'log_age = {log_age} not found.')
            parsec = parsec[age_cut]
        if zini is not None:
            zini_cut = np.abs(parsec['Zini'] - zini) < 1e-8
            if zini_cut.sum() < 1:
                raise Exception(f'Zini = {zini} not found.')
            parsec = parsec[zini_cut]
        if num_filters is None:
            filt_idx = np.argwhere(np.array(names) == 'mbolmag')[0][0] + 1
        else:
            filt_idx = len(names) - num_filters
        iso = Isochrone(mini=parsec['Mini'],
                        mact=parsec['Mass'],
                        mags=parsec[names[filt_idx:]],
                        log_L=parsec['logL'],
                        log_Teff=parsec['logTe'])
        iso.isochrone_full = isochrone_full
        return iso

    def interpolate(self, y_name, x_interp, x_name='mini',
                    slice_interp=np.s_[:], **kwargs):
        """
        Interpolate isochrone.

        Parameters
        ----------
        y_name : str
            Parameter name of interpolated variable.
        x_interp : `~numpy.ndarray`
            New x values to interpolate over.
        x_name : str
            Parameter name of x values. Usually want this to be initial mass.

        Returns
        -------
        y_interp : `~numpy.ndarray`
            Interpolated y values.
        """
        x = self.mag_table[x_name] if x_name in self.filters\
                                   else getattr(self, x_name)
        y = self.mag_table[y_name] if y_name in self.filters\
                                   else getattr(self, y_name)
        x = x[slice_interp]
        y = y[slice_interp]
        y_interp = interp1d(x, y, **kwargs)(x_interp)
        return y_interp

    def mag_to_mass(self, mag, bandpass):
        """
        Interpolate isochrone to calculate initial stellar masses that have
        the given magnitude. Since more than one mass can have the same
        magnitude (from stars being in different phases of stellar evolution),
        we step through each isochrone and check if the given magnitude falls
        between two bins, in which case we interpolate the associated mass.

        Parameters
        ----------
        mag : float
            Magnitude to be converted in to stellar mass(es).
        bandpass : str
            Photometric bandpass of the `mag`.

        Returns
        -------
        mass_interp : `~numpy.ndarray`
            Interpolated mass(es) associated with `mag`.
        """
        y = self.mini
        x = self.mag_table[bandpass]

        if mag < x.min() or mag > x.max():
            raise Exception(f'mag = {mag} is outside the isochrone range.')

        y_interp = []
        # Loop over isochrone and check if desired magnitude is in between two
        # bins. If it is, interpolate the mass. Looping like this is necessary
        # because different masses can have the same magnitude due to
        # being at different phases of stellar evolution.
        for i in range(len(x)):
            if x[i] == mag:
                y_interp.append(y[i])
                continue
            if i < len(x) - 1:
                _x = [x[i], x[i + 1]]
                _y = [y[i], y[i + 1]]
                dx = x[i + 1] - x[i]
                if dx > 0:
                    if (mag <= x[i + 1]) and (mag > x[i]):
                        y_interp.append(interp1d(_x, _y)([mag])[0])
                else:
                    if (mag >= x[i + 1]) and (mag < x[i]):
                        y_interp.append(interp1d(_x, _y)([mag])[0])
        mass_interp = np.array(y_interp)

        return mass_interp

    def calculate_mag_limit(self, imf, bandpass, frac_mass_sampled=None,
                            frac_num_sampled=None, distance=10*u.pc):
        """
        Calculate the limiting faint magnitude to sample a given fraction
        of mass or number of a stellar population. This is used for when you
        only what to sample stars that are brighter than a magnitude limit.
        You must specify either `frac_mass_sampled` or `frac_num_sampled`.

        .. note::
            You must give `frac_mass_sampled` *or* `frac_num_sampled`.

        Parameters
        ----------
        imf : str
            Name for the IMF.
        bandpass : str
            Observation bandpass of the limiting magnitude.
        frac_mass_sampled: float, optional
            Fraction of mass that will be sampled if only stars that are
            brighter than the magnitude limit are sampled.
        frac_num_sampled: float, optional
            Fraction of stars by number that will be sampled if only stars that
            are brighter than the magnitude limit are sampled.
        distance : float or `~astropy.units.Quantity`, optional
            Distance to source. If float is given, the units are assumed
            to be `~astropy.units.Mpc`. Default distance is 10
            `~astropy.units.pc` (i.e., the mags are in absolute units).

        Returns
        -------
        mag_limit : float
            The desired limiting magnitude.
        mass_limit : float
            Mass of stars (in solar masses) that have the limiting magnitude.
        """
        mfint = IMFIntegrator(imf, self.m_min, self.m_max)
        m_vals = np.logspace(np.log10(self.m_min), np.log10(self.m_max), 500)
        err_msg = 'You must give a fraction such that 0 < fraction <= 1.'
        if frac_mass_sampled is not None:
            assert frac_mass_sampled > 0 and frac_mass_sampled <= 1, err_msg
            frac_interp = frac_mass_sampled
            fracs = [mfint.m_integrate(_m, self.m_max, True) for _m in m_vals]
        elif frac_num_sampled is not None:
            assert frac_num_sampled > 0 and frac_num_sampled <= 1, err_msg
            frac_interp = frac_num_sampled
            fracs = [mfint.integrate(_m, self.m_max, True) for _m in m_vals]
        else:
            msg = 'You must give either frac_mass_sampled or frac_num_sampled.'
            raise Exception(msg)
        d = check_units(distance, 'Mpc')
        mass_limit = interp1d(fracs, m_vals)([frac_interp])[0]
        mag_limit = self.interpolate(bandpass, [mass_limit])[0]
        mag_limit = mag_limit + 5 * np.log10(d.to('pc').value) - 5
        return mag_limit, mass_limit

    def nearest_mini(self, m):
        """
        Find the nearest mass to `m` in the initial mass array.

        Parameters
        ----------
        m : float:
            The mass for which we want the nearest initial mass.

        Returns
        -------
        m_nearest : float
            Value of the nearest mass.
        arg_nearest : int
            Argument of the nearest mass.
        """
        arg_nearest = np.abs(self.mini - m).argmin()
        m_nearest = self.mini[arg_nearest]
        return m_nearest, arg_nearest

    def imf_weights(self, imf, m_min_norm=None, m_max_norm=None,
                    norm_type='mass', **kwargs):
        """
        Calculate IMF weights.

        Parameters
        ----------
        imf : str or callable function that takes mass as its argument
            Stellar Initial mass function.
        m_min_norm : None or float, optional
            Minimum mass for the normalization. Must be less than or equal to
            the mini mass of isochrone, which will be used if None is given.
        m_max_norm : float, optional
            Maximum mass for normalization.
        norm_type : str, optional
            Type of IMF normalization (mass or number).

        Returns
        -------
        wght : `~numpy.ndarray`
            IMF weights calculated such that an integral over mass is simply
            given by SUM(m * wght).
        """
        m_min = m_min_norm if m_min_norm else self.mini.min()
        m_max = m_max_norm if m_max_norm else self.mini.max()
        if m_min > self.m_min:
            raise Exception('Minimum mass must be <= isochrone min mass.')
        if callable(imf):
            imf_func = imf
        else:
            imf_func = lambda m: imf_dict[imf](m, norm_type=None)
        m_imf_func = lambda m: m * imf_func(m)
        norm_func = dict(mass=m_imf_func, number=imf_func)[norm_type]
        norm = quad(norm_func, m_min, m_max, **kwargs)[0]

        wght = []
        mini = self.mini
        # Assume mass is constant in each bin and integrate.
        # This means an integral over mass is simply SUM(m * wght)
        # This trick was stolen from Charlie Conroy's FSPS code :)
        for i in range(len(mini)):
            if i == 0:
                m1 = m_min
            else:
                m1 = mini[i] - 0.5 * (mini[i] - mini[i-1])
            if i == len(mini) - 1:
                m2 = mini[i]
            else:
                m2 = mini[i] + 0.5 * (mini[i+1] - mini[i])
            if m2 < m1:
                raise Exception('Masses must be monotonically increasing.')
            wght.append(quad(imf_func, m1, m2, **kwargs)[0])
        wght = np.array(wght) / norm
        return wght

    def ssp_color(self, blue, red, imf='kroupa', **kwargs):
        """
        Calculate IMF-weighted integrated color.

        Parameters
        ----------
        blue : str
            Blue bandpass.
        red : str
            Red bandpass.
        imf : str, optional
            IMF name.

        Returns
        -------
        color : float
            Integrated color of stellar population.
        """
        wght = self.imf_weights(imf, **kwargs)
        lum_blue = np.sum(wght * 10**(-0.4 * self.mag_table[blue]))
        lum_red = np.sum(wght * 10**(-0.4 * self.mag_table[red]))
        color = -2.5 * np.log10(lum_blue / lum_red)
        return color

    def ssp_sbf_mag(self, bandpass, imf='kroupa', **kwargs):
        """
        Calculate IMF-weighted SBF magnitude.

        Parameters
        ----------
        bandpass : str
            Bandpass to of SBF mag.
        imf : str, optional
            IMF name.

        Returns
        -------
        sbf : float
            SBF magnitude of stellar population.
        """
        wght = self.imf_weights(imf, **kwargs)
        lumlum = np.sum(wght * 10**(-0.8 * self.mag_table[bandpass]))
        lum = np.sum(wght * 10**(-0.4 * self.mag_table[bandpass]))
        sbf = -2.5 * np.log10(lumlum / lum)
        return sbf

    def ssp_mag(self, bandpass, imf='kroupa', norm_type='mass', **kwargs):
        """
        Calculate IMF-weighted magnitude.

        Parameters
        ----------
        bandpass : str
            Bandpass to of mean mag.
        imf : str, optional
            IMF name.
        norm_type : str, optional
            Normalization type (mass or number)

        Returns
        -------
        m : float
            Integrated magnitude of stellar population.
        """
        wght = self.imf_weights(imf, norm_type=norm_type, **kwargs)
        lum = np.sum(wght * 10**(-0.4 * self.mag_table[bandpass]))
        m = -2.5 * np.log10(lum)
        return m

    def ssp_mean_star_mass(self, imf):
        """
        Calculate IMF-weighted mean stellar mass, where the IMF is normalized
        by number to one star formed. The mean stellar mass is equivalent to
        dividing a large of samples (distributed according to the IMF) by the
        number of stars.

        Parameters
        ----------
        imf : str:
            Initial mass function name.

        Returns
        -------
        mean_mass : float
            The mean stellar mass.
        """
        w = self.imf_weights(imf, norm_type='number',
                             m_max_norm=self.mini.max())
        mean_mass = np.sum(self.mact * w)
        return mean_mass

    def ssp_surviving_mass(self, imf, m_min=None, m_max=120,
                           add_remnants=True, mlim_bh=40.0, mlim_ns=8.5):
        """
        Calculate IMF-weighted mass normalized such that 1 solar mass is
        formed by the age of the population.

        The initial-mass-dependent remnant formulae are taken
        from `Renzini & Ciotti 1993
        <https://ui.adsabs.harvard.edu/abs/1993ApJ...416L..49R/abstract>`_.

        Parameters
        ----------
        imf : str
            IMF name.
        add_remnants : bool
            If True, add stellar remnants to surviving mass.

        Returns
        -------
        mass : float
            Total surviving mass.
        """
        m_min = m_min if m_min else self.mini.min()
        m_max = m_max if m_max else self.mini.max()
        wght = self.imf_weights(imf, m_min_norm=m_min, m_max_norm=m_max)
        mass = np.sum(self.mact * wght)

        if add_remnants:
            imf_func = lambda m: imf_dict[imf](m, norm_type=None)
            m_imf_func = lambda m: m * imf_func(m)
            norm = quad(m_imf_func, m_min, m_max)[0]

            # BH remnants
            m_low = max(mlim_bh, self.m_max)
            mass = mass + 0.5 * quad(m_imf_func, m_low, m_max)[0] / norm

            # NS remnants
            if self.m_max <= mlim_bh:
                m_low = max(mlim_ns, self.m_max)
                mass = mass + 1.4 * quad(imf_func, m_low, mlim_bh)[0] / norm

            # WD remnants
            if self.m_max < mlim_ns:
                mmax = self.m_max
                mass = mass + 0.48 * quad(imf_func, mmax, mlim_ns)[0] / norm
                mass = mass + 0.077 * quad(m_imf_func, mmax, mlim_ns)[0] / norm

        return mass

def fetch_mist_iso_cmd(log_age, feh, phot_system, mist_path=MIST_PATH,
                       v_over_vcrit=0.4):
    """
    Fetch MIST isochrone grid.

    Parameters
    ----------
    log_age : float
        Logarithm base 10 of the simple stellar population age in years.
    feh : float
        Metallicity [Fe/H] of the simple stellar population.
    phot_system : str
        Name of the photometric system.
    mist_path : str, optional
        Path to MIST isochrone grids. Use this if you want to use a different
        path from the `MIST_PATH` environment variable.
    v_over_vcrit : float, optional
        Rotation rate divided by the critical surface linear velocity. Current
        options are 0.4 (default) and 0.0.

    Returns
    -------
    iso_cmd : `~numpy.ndarray`
        Structured ``numpy`` array with isochrones and stellar magnitudes.
    """

    # fetch the mist grid if necessary
    fetch_mist_grid_if_needed(phot_system, v_over_vcrit, mist_path)

    v = f'{v_over_vcrit:.1f}'
    ver = 'v1.2'
    p = phot_str_helper[phot_system.lower()]
    path = os.path.join(mist_path, 'MIST_' + ver + f'_vvcrit{v}_' + p)
    sign = 'm' if feh < 0 else 'p'
    fn = f'MIST_{ver}_feh_{sign}{abs(feh):.2f}_afe_p0.0_vvcrit{v}_{p}.iso.cmd'
    fn = os.path.join(path, fn)
    iso_cmd = IsoCmdReader(fn, verbose=False)
    iso_cmd = iso_cmd.isocmds[iso_cmd.age_index(log_age)]
    return iso_cmd

class MISTIsochrone(Isochrone):
    """
    Class for fetching and storing MIST isochrones. It also has several methods
    for calculating IMF-weighted photometric properties of a stellar population
    with the given age an metallicity.

    .. note::
        Currently, the models are interpolated in metallicity but not in age.
        Ages are therefore limited to the age grid of the MIST models. The
        [Fe/H] and log(Age/yr) grids are stored as the private class attributes
        ``_feh_grid`` and ``_log_age_grid``.

    Parameters
    ----------
    log_age : float
        Logarithm base 10 of the simple stellar population age in years.
    feh : float
        Metallicity [Fe/H] of the simple stellar population.
    phot_system : str or list-like
        Name of the photometric system(s).
    mist_path : str, optional
        Path to MIST isochrone grids. Use this if you want to use a different
        path from the ``MIST_PATH`` environment variable.
    ab_or_vega : str, optional
        Magnitudes will be in the AB (default) or Vega magnitude system.
    v_over_vcrit : float, optional
        Rotation rate divided by the critical surface linear velocity. Current
        options are 0.4 (default) and 0.0.
    """

    # the age grid
    _log_age_grid = np.arange(5.0, 10.3, 0.05)
    _log_age_min = _log_age_grid.min()
    _log_age_max = _log_age_grid.max()

    # the [Fe/H] metallicity grid
    # mist has feh <= -4, but using <=-3 due to interpolation issues
    _feh_grid = np.concatenate([np.arange(-3.0, -2., 0.5),
                                np.arange(-2.0, 0.75, 0.25)])
    _feh_min = _feh_grid.min()
    _feh_max = _feh_grid.max()

    def __init__(self, log_age, feh, phot_system, mist_path=MIST_PATH,
                 ab_or_vega='ab', v_over_vcrit=0.4):

        # verify age are metallicity are within model grids
        if log_age < self._log_age_min or log_age > self._log_age_max:
            raise Exception(f'log_age = {log_age} not in range of age grid')
        if feh < self._feh_min or feh > self._feh_max:
            raise Exception(f'feh = {feh} not in range of feh grid')

        self.feh = feh
        self.mist_path = mist_path
        self.phot_system = phot_system
        self.ab_or_vega = ab_or_vega
        self.v_over_vcrit = v_over_vcrit

        # use nearest age (currently not interpolating on age)
        age_diff = np.abs(self._log_age_grid - log_age)
        self.log_age = self._log_age_grid[age_diff.argmin()]
        if age_diff.min() > 1e-6:
            logger.debug('Using nearest log_age = {:.2f}'.format(self.log_age))

        # store phot_system as list to allow multiple photometric systems
        if type(phot_system) == str:
            phot_system = [phot_system]

        # fetch first isochrone grid, interpolating on [Fe/H] if necessary
        self._iso_full = self._fetch_iso(phot_system[0])

        # iterate over photometric systems and fetch remaining isochrones
        filter_dict = get_filter_names()
        filters = filter_dict[phot_system[0]].copy()
        for p in phot_system[1:]:
            filt = filter_dict[p].copy()
            filters.extend(filt)
            _iso = self._fetch_iso(p)
            mags = [_iso[f].data for f in filt]
            self._iso_full = append_fields(self._iso_full, filt, mags)

        # covert magnitudes to ab or vega as necessary
        self.zpt_convert = load_zero_point_converter()
        for filt in filters:
            converter = getattr(self.zpt_convert, f'to_{ab_or_vega.lower()}')
            try:
                m_convert = converter(filt)
            except AttributeError:
                m_convert = 0.0
                logger.warning(f'No AB / Vega conversion found for {filt}.')
            self._iso_full[filt] = self._iso_full[filt] + m_convert

        super(MISTIsochrone, self).__init__(
            mini = self._iso_full['initial_mass'],
            mact = self._iso_full['star_mass'],
            mags = Table(self._iso_full[filters]),
            eep = self._iso_full['EEP'],
            log_L = self._iso_full['log_L'],
            log_Teff = self._iso_full['log_Teff'],
        )

    @property
    def isochrone_full(self):
        """MIST entire isochrone in a structured `~numpy.ndarray`."""
        return self._iso_full

    @staticmethod
    def from_parsec(fn, **kwargs):
        msg = 'PARSEC isochrones do not with MISTIsochrone.'
        raise Exception(msg + ' Use artpop.Isochrone instead.')

    def _fetch_iso(self, phot_system):
        """Fetch MIST isochrone grid, interpolating on [Fe/H] if necessary."""
        if self.feh in self._feh_grid:
            args = [self.log_age, self.feh, phot_system, self.mist_path,
                    self.v_over_vcrit]
            iso = fetch_mist_iso_cmd(*args)
        else:
            iso = self._interp_on_feh(phot_system)
        return iso

    def _interp_on_feh(self, phot_system):
        """Interpolate isochrones between two [Fe/H] grid points."""
        i_feh = self._feh_grid.searchsorted(self.feh)
        feh_lo, feh_hi = self._feh_grid[i_feh - 1: i_feh + 1]

        logger.debug('Interpolating to [Fe/H] = {:.2f} '\
                     'using [Fe/H] = {} and {}'.\
                     format(self.feh, feh_lo, feh_hi))

        mist_0 = fetch_mist_iso_cmd(
            self.log_age, feh_lo, phot_system, self.mist_path)
        mist_1 = fetch_mist_iso_cmd(
            self.log_age, feh_hi, phot_system, self.mist_path)

        y0, y1 = np.array(mist_0.tolist()), np.array(mist_1.tolist())

        x = self.feh
        x0, x1 = feh_lo, feh_hi
        weight = (x - x0) / (x1 - x0)

        len_0, len_1 = len(y0), len(y1)

        # if necessary, extrapolate using trend of the longer array
        if (len_0 < len_1):
            delta = y1[len_0:] - y1[len_0 - 1]
            y0 = np.append(y0, y0[-1] + delta, axis=0)
        elif (len_0 > len_1):
            delta = y0[len_1:] - y0[len_1 - 1]
            y1 = np.append(y1, y1[-1] + delta, axis=0)

        y = y0 * (1 - weight) + y1 * weight
        iso = np.core.records.fromarrays(y.transpose(), dtype=mist_0.dtype)

        return iso

import numpy,from numpy,from scipy,from astropy


"""
Module with reading functionalities for isochrones.
"""

import configparser
import os
import warnings

from typing import Optional, Tuple

import h5py
import numpy as np

from scipy.interpolate import griddata
from typeguard import typechecked

from species.core import box
from species.read import read_model
from species.util import read_util

class ReadIsochrone:
    """
    Class for reading isochrone data from the database.
    """

    @typechecked
    def __init__(self, tag: str) -> None:
        """
        Parameters
        ----------
        tag : str
            Database tag of the isochrone data.

        Returns
        -------
        NoneType
            None
        """

        self.tag = tag

        config_file = os.path.join(os.getcwd(), "species_config.ini")

        config = configparser.ConfigParser()
        config.read(config_file)

        self.database = config["species"]["database"]

        with h5py.File(self.database, "r") as h5_file:
            if f"isochrones/{self.tag}" not in h5_file:
                raise ValueError(f"There is no isochrone data stored "
                                 f"with the selected tag \'{tag}\'.")

    @typechecked
    def get_isochrone(
        self,
        age: float,
        masses: np.ndarray,
        filters_color: Optional[Tuple[str, str]] = None,
        filter_mag: Optional[str] = None,
    ) -> box.IsochroneBox:
        """
        Function for selecting an isochrone.

        Parameters
        ----------
        age : float
            Age (Myr) at which the isochrone data is interpolated.
        masses : np.ndarray
            Masses (:math:`M_\\mathrm{J}`) at which the isochrone
            data is interpolated.
        filters_color : tuple(str, str), None
            Filter names for the color as listed in the file with the
            isochrone data. Not selected if set to ``None`` or if only
            evolutionary tracks are available.
        filter_mag : str, None
            Filter name for the absolute magnitude as listed in the
            file with the isochrone data. Not selected if set to
            ``None`` or if only evolutionary tracks are available.

        Returns
        -------
        species.core.box.IsochroneBox
            Box with the isochrone.
        """

        age_points = np.full(masses.shape[0], age)  # (Myr)

        color = None
        mag_abs = None

        index_teff = 2
        index_logg = 4

        # Read isochrone data

        with h5py.File(self.database, "r") as h5_file:
            model = h5_file[f"isochrones/{self.tag}/evolution"].attrs["model"]
            evolution = np.asarray(h5_file[f"isochrones/{self.tag}/evolution"])

            if model == "baraffe":
                filters = list(h5_file[f"isochrones/{self.tag}/filters"])
                magnitudes = np.asarray(h5_file[f"isochrones/{self.tag}/magnitudes"])

                # Convert the h5py list of filters from bytes to strings
                for i, item in enumerate(filters):
                    if isinstance(item, bytes):
                        filters[i] = item.decode("utf-8")

        if model == "baraffe":
            if filters_color is not None:
                index_color_1 = filters.index(filters_color[0])
                index_color_2 = filters.index(filters_color[1])

            if filter_mag is not None:
                index_mag = filters.index(filter_mag)

            if filters_color is not None:
                mag_color_1 = griddata(
                    points=evolution[:, 0:2],
                    values=magnitudes[:, index_color_1],
                    xi=np.stack((age_points, masses), axis=1),
                    method="linear",
                    fill_value="nan",
                    rescale=False,
                )

                mag_color_2 = griddata(
                    points=evolution[:, 0:2],
                    values=magnitudes[:, index_color_2],
                    xi=np.stack((age_points, masses), axis=1),
                    method="linear",
                    fill_value="nan",
                    rescale=False,
                )

                color = mag_color_1 - mag_color_2

            if filter_mag is not None:
                mag_abs = griddata(
                    points=evolution[:, 0:2],
                    values=magnitudes[:, index_mag],
                    xi=np.stack((age_points, masses), axis=1),
                    method="linear",
                    fill_value="nan",
                    rescale=False,
                )

        teff = griddata(
            points=evolution[:, 0:2],
            values=evolution[:, index_teff],
            xi=np.stack((age_points, masses), axis=1),
            method="linear",
            fill_value="nan",
            rescale=False,
        )

        logg = griddata(
            points=evolution[:, 0:2],
            values=evolution[:, index_logg],
            xi=np.stack((age_points, masses), axis=1),
            method="linear",
            fill_value="nan",
            rescale=False,
        )

        return box.create_box(
            boxtype="isochrone",
            model=self.tag,
            filters_color=filters_color,
            filter_mag=filter_mag,
            color=color,
            magnitude=mag_abs,
            teff=teff,
            logg=logg,
            masses=masses,
        )

    @typechecked
    def get_color_magnitude(
        self,
        age: float,
        masses: np.ndarray,
        model: str,
        filters_color: Tuple[str, str],
        filter_mag: str,
        adapt_logg: bool = False,
    ) -> box.ColorMagBox:
        """
        Function for calculating color-magnitude
        combinations from a selected isochrone.

        Parameters
        ----------
        age : float
            Age (Myr) at which the isochrone data is interpolated.
        masses : np.ndarray
            Masses (:math:`M_\\mathrm{J}`) at which the isochrone
            data is interpolated.
        model : str
            Atmospheric model used to compute the synthetic photometry.
        filters_color : tuple(str, str)
            Filter names for the color as listed in the file with the
            isochrone data. The filter names should be provided in the
            format of the SVO Filter Profile Service.
        filter_mag : str
            Filter name for the absolute magnitude as listed in the
            file with the isochrone data. The value should be equal
            to one of the ``filters_color`` values.
        adapt_logg : bool
            Adapt :math:`\\log(g)` to the upper or lower boundary of
            the atmospheric model grid whenever the :math:`\\log(g)`
            that has been calculated from the isochrone mass and
            radius lies outside the available range of the synthetic
            spectra. Typically :math:`\\log(g)` has only a minor
            impact on the broadband magnitudes and colors.

        Returns
        -------
        species.core.box.ColorMagBox
            Box with the color-magnitude data.
        """

        isochrone = self.get_isochrone(
            age=age, masses=masses, filters_color=None, filter_mag=None
        )

        model1 = read_model.ReadModel(model=model, filter_name=filters_color[0])
        model2 = read_model.ReadModel(model=model, filter_name=filters_color[1])

        param_bounds = model1.get_bounds()

        if model1.get_parameters() == ["teff", "logg", "feh"]:
            if model == "sonora-bobcat":
                iso_feh = float(self.tag[-4:])
            else:
                iso_feh = 0.0

        elif model1.get_parameters() != ["teff", "logg"]:
            raise ValueError(
                "Creating synthetic colors and magnitudes from "
                "isochrones is currently only implemented for "
                "models with only Teff and log(g) as free parameters. "
                "Please contact Tomas Stolker if additional "
                "functionalities are required."
            )

        else:
            iso_feh = None

        mag1 = np.zeros(isochrone.masses.shape[0])
        mag2 = np.zeros(isochrone.masses.shape[0])
        radius = np.zeros(isochrone.masses.shape[0])

        for i, mass_item in enumerate(isochrone.masses):
            model_param = {
                "teff": isochrone.teff[i],
                "logg": isochrone.logg[i],
                "mass": mass_item,
                "distance": 10.0,
            }

            if iso_feh is not None:
                model_param['feh'] = iso_feh

            radius[i] = read_util.get_radius(
                model_param["logg"], model_param["mass"]
            )  # (Rjup)

            if np.isnan(isochrone.teff[i]):
                mag1[i] = np.nan
                mag2[i] = np.nan

                warnings.warn(
                    f"The value of Teff is NaN for the following "
                    f"isochrone sample: {model_param}. Setting "
                    f"the magnitudes to NaN."
                )

            else:
                for item_bounds in param_bounds:
                    if model_param[item_bounds] < param_bounds[item_bounds][0]:
                        if adapt_logg and item_bounds == "logg":
                            warnings.warn(
                                f"The log(g) is {model_param[item_bounds]} but the "
                                f"lower boundary of the model grid is "
                                f"{param_bounds[item_bounds][0]}. Adapting "
                                f"log(g) to {param_bounds[item_bounds][0]} since "
                                f"adapt_logg=True."
                            )

                            model_param["logg"] = param_bounds["logg"][0]

                        else:
                            mag1[i] = np.nan
                            mag2[i] = np.nan

                            warnings.warn(
                                f"The value of {item_bounds} is "
                                f"{model_param[item_bounds]}, which is below "
                                f"the lower bound of the model grid "
                                f"({param_bounds[item_bounds][0]}). Setting the "
                                f"magnitudes to NaN for the following isochrone "
                                f"sample: {model_param}."
                            )

                    elif model_param[item_bounds] > param_bounds[item_bounds][1]:
                        if adapt_logg and item_bounds == "logg":
                            warnings.warn(
                                f"The log(g) is {model_param[item_bounds]} but "
                                f"the upper boundary of the model grid is "
                                f"{param_bounds[item_bounds][1]}. Adapting "
                                f"log(g) to {param_bounds[item_bounds][1]} "
                                f"since adapt_logg=True."
                            )

                            model_param["logg"] = param_bounds["logg"][1]

                        else:
                            mag1[i] = np.nan
                            mag2[i] = np.nan

                            warnings.warn(
                                f"The value of {item_bounds} is "
                                f"{model_param[item_bounds]}, which is above "
                                f"the upper bound of the model grid "
                                f"({param_bounds[item_bounds][1]}). Setting the "
                                f"magnitudes to NaN for the following isochrone "
                                f"sample: {model_param}."
                            )

                if not np.isnan(mag1[i]):
                    mag1[i], _ = model1.get_magnitude(model_param)
                    mag2[i], _ = model2.get_magnitude(model_param)

        if filter_mag == filters_color[0]:
            abs_mag = mag1

        elif filter_mag == filters_color[1]:
            abs_mag = mag2

        else:
            raise ValueError(
                "The argument of filter_mag should be equal to "
                "one of the two filter values of filters_color."
            )

        return box.create_box(
            boxtype="colormag",
            library=model,
            object_type="model",
            filters_color=filters_color,
            filter_mag=filter_mag,
            color=mag1 - mag2,
            magnitude=abs_mag,
            names=None,
            sptype=masses,
            mass=masses,
            radius=radius,
            iso_tag=self.tag,
        )

    @typechecked
    def get_color_color(
        self,
        age: float,
        masses: np.ndarray,
        model: str,
        filters_colors: Tuple[Tuple[str, str], Tuple[str, str]],
    ) -> box.ColorColorBox:
        """
        Function for calculating color-magnitude combinations from a
        selected isochrone.

        Parameters
        ----------
        age : float
            Age (Myr) at which the isochrone data is interpolated.
        masses : np.ndarray
            Masses (:math:`M_\\mathrm{J}`) at which the isochrone
            data is interpolated.
        model : str
            Atmospheric model used to compute the synthetic photometry.
        filters_colors : tuple(tuple(str, str), tuple(str, str))
            Filter names for the colors as listed in the file with the
            isochrone data. The filter names should be provided in the
            format of the SVO Filter Profile Service.

        Returns
        -------
        species.core.box.ColorColorBox
            Box with the color-color data.
        """

        isochrone = self.get_isochrone(
            age=age, masses=masses, filters_color=None, filter_mag=None
        )

        model1 = read_model.ReadModel(model=model, filter_name=filters_colors[0][0])
        model2 = read_model.ReadModel(model=model, filter_name=filters_colors[0][1])
        model3 = read_model.ReadModel(model=model, filter_name=filters_colors[1][0])
        model4 = read_model.ReadModel(model=model, filter_name=filters_colors[1][1])

        if model1.get_parameters() == ["teff", "logg", "feh"]:
            if model == "sonora-bobcat":
                iso_feh = float(self.tag[-4:])
            else:
                iso_feh = 0.0

        elif model1.get_parameters() != ["teff", "logg"]:
            raise ValueError(
                "Creating synthetic colors and magnitudes from "
                "isochrones is currently only implemented for "
                "models with only Teff and log(g) as free parameters. "
                "Please contact Tomas Stolker if additional "
                "functionalities are required."
            )

        else:
            iso_feh = None

        mag1 = np.zeros(isochrone.masses.shape[0])
        mag2 = np.zeros(isochrone.masses.shape[0])
        mag3 = np.zeros(isochrone.masses.shape[0])
        mag4 = np.zeros(isochrone.masses.shape[0])
        radius = np.zeros(isochrone.masses.shape[0])

        for i, mass_item in enumerate(isochrone.masses):
            model_param = {
                "teff": isochrone.teff[i],
                "logg": isochrone.logg[i],
                "mass": mass_item,
                "distance": 10.0,
            }

            if iso_feh is not None:
                model_param['feh'] = iso_feh

            radius[i] = read_util.get_radius(
                model_param["logg"], model_param["mass"]
            )  # (Rjup)

            if np.isnan(isochrone.teff[i]):
                mag1[i] = np.nan
                mag2[i] = np.nan
                mag3[i] = np.nan
                mag4[i] = np.nan

                warnings.warn(
                    f"The value of Teff is NaN for the following isochrone "
                    f"sample: {model_param}. Setting the magnitudes to NaN."
                )

            else:
                for item_bounds in model1.get_bounds():
                    if model_param[item_bounds] < model1.get_bounds()[item_bounds][0]:
                        mag1[i] = np.nan
                        mag2[i] = np.nan
                        mag3[i] = np.nan
                        mag4[i] = np.nan

                        warnings.warn(
                            f"The value of {item_bounds} is "
                            f"{model_param[item_bounds]}, which is "
                            f"below the lower bound of the model grid "
                            f" ({model1.get_bounds()[item_bounds][0]}). "
                            f"Setting the magnitudes to NaN for the "
                            f"following isochrone sample: {model_param}."
                        )

                    elif (
                        model_param[item_bounds] > model1.get_bounds()[item_bounds][1]
                    ):
                        mag1[i] = np.nan
                        mag2[i] = np.nan
                        mag3[i] = np.nan
                        mag4[i] = np.nan

                        warnings.warn(
                            f"The value of {item_bounds} is "
                            f"{model_param[item_bounds]}, which is above "
                            f"the upper bound of the model grid "
                            f"({model1.get_bounds()[item_bounds][1]}). "
                            f"Setting the magnitudes to NaN for the "
                            f"following isochrone sample: {model_param}."
                        )

                if (
                    not np.isnan(mag1[i])
                    and not np.isnan(mag2[i])
                    and not np.isnan(mag3[i])
                    and not np.isnan(mag4[i])
                ):
                    mag1[i], _ = model1.get_magnitude(model_param)
                    mag2[i], _ = model2.get_magnitude(model_param)
                    mag3[i], _ = model3.get_magnitude(model_param)
                    mag4[i], _ = model4.get_magnitude(model_param)

        return box.create_box(
            boxtype="colorcolor",
            library=model,
            object_type="model",
            filters=filters_colors,
            color1=mag1 - mag2,
            color2=mag3 - mag4,
            names=None,
            sptype=masses,
            mass=masses,
            radius=radius,
            iso_tag=self.tag,
        )


import os

import numpy as np
from scipy.interpolate import RectBivariateSpline

from adf11 import Adf11

argon_data = {
    'ionisation' : 'scd89_ar.dat',
    'recombination' : 'acd89_ar.dat',
    'continuum_power' : 'prb89_ar.dat',
    'line_power' : 'plt89_ar.dat',
    'cx_power' : 'prc89_ar.dat',
}

carbon_data = {
    'ionisation' : 'scd96_c.dat',
    'recombination' : 'acd96_c.dat',
    'continuum_power' : 'prb96_c.dat',
    'line_power' : 'plt96_c.dat',
    'cx_power' : 'prc96_c.dat',
}

neon_data = {
    'ionisation' : 'scd96_ne.dat',
    'recombination' : 'acd96_ne.dat',
    'continuum_power' : 'prb96_ne.dat',
    'line_power' : 'plt96_ne.dat',
}

def _element_data(element):
    e = element.lower()
    if e in ['ar', 'argon']:
        return argon_data
    elif e in ['c', 'carbon']:
        return carbon_data
    elif e in ['ne', 'neon']:
        return neon_data
    else:
        raise NotImplementedError('unknown element: %s' % element)

def _full_path(file_):
    """ Figure out the location of the atomic datafiles. """
    module_path = os.path.dirname(os.path.realpath( __file__ ))
    return os.path.join(module_path, '..', 'adas_data', file_)

class AtomicData(object):
    def __init__(self, coefficients):
        """
        Parameters
        ----------
        element : string
            Name of the element.
        coefficients : dict
            Map of the different rate coefficients.
        """
        self.coeffs = coefficients
        self._check_consistency()
        self._make_element_initial_uppercase()

    def copy(self):
        new_coeffs = {}
        for key, value in self.coeffs.iteritems():
            new_coeffs[key] = value.copy()

        return self.__class__(new_coeffs)

    def _check_consistency(self):
        nuclear_charge = set()
        element = set()
        for coeff in self.coeffs.values():
            nuclear_charge.add(coeff.nuclear_charge)
            element.add(coeff.element.lower())

        assert len(nuclear_charge) == 1, 'inconsistent nuclear charge.'
        assert len(element) == 1, 'inconsistent element name.'

        self.nuclear_charge = nuclear_charge.pop()
        self.element = element.pop()

    @classmethod
    def from_element(cls, element):
        element_data = _element_data(element)

        coefficients = {}
        for key in element_data.keys():
            name = _full_path(element_data[key])
            coefficients[key] = RateCoefficient.from_adf11(name)

        return cls(coefficients)

    def _make_element_initial_uppercase(self):
        e = self.element
        self.element = e[0].upper() + e[1:]

class RateCoefficient(object):
    def __init__(self, nuclear_charge, element, log_temperature, log_density,
            log_coeff, name=None):
        self.nuclear_charge = nuclear_charge
        self.element = element
        self.adf11_file = name

        self.log_temperature = log_temperature
        self.log_density = log_density
        self.log_coeff = log_coeff

        self._compute_interpolating_splines()

    @classmethod
    def from_adf11(cls, name):
        adf11_data = Adf11(name).read()
        nuclear_charge = adf11_data['charge']
        element = adf11_data['element']
        filename = adf11_data['name']

        log_temperature = adf11_data['log_temperature']
        log_density = adf11_data['log_density']
        log_coeff = adf11_data['log_coeff']

        return cls(nuclear_charge, element, log_temperature, log_density,
                log_coeff, name=filename)

    def copy(self):
        log_temperature = self.log_temperature.copy()
        log_density = self.log_density.copy()
        log_coeff = self.log_coeff.copy()
        cls = self.__class__(self.nuclear_charge, self.element, log_temperature,
                log_density, log_coeff, self.adf11_file)
        return cls

    def _compute_interpolating_splines(self):
        self.splines = []
        for k in xrange(self.nuclear_charge):
            x = self.log_temperature
            y = self.log_density
            z = self.log_coeff[k]
            self.splines.append(RectBivariateSpline(x, y, z))

    def __call__(self, k, Te, ne):
        """
        Evaulate the ionisation/recombination coefficients of k'th atomic state
        at a given temperature and density.

        Parameters
        ----------
        k : int
            Ionisation stage. k=0 is the neutral atom, k=Z is  the fully
            stripped ion, where Z is the atomic number.
        Te : array_like
            Temperature in [eV].
        ne : array_like
            Density in [m-3].

        Returns
        -------
        c : array_like
            Rate coefficent in [m3/s].
        """
        c = self.log10(k, Te, ne)
        return np.power(10, c)

    def log10(self, k, Te, ne):
        """
        Evaulate the logarithm of ionisation/recombination coefficients of
        k'th atomic state at a given temperature and density.

        Parameters
        ----------
        k : int
            Ionisation stage. k=0 is the neutral atom, k=Z is  the fully
            stripped ion, where Z is the atomic number.
        Te : array_like
            Temperature in [eV].
        ne : array_like
            Density in [m-3].

        Returns
        -------
        c : array_like
            log10(rate coefficent in [m3/s])
        """

        Te, ne = np.broadcast_arrays(Te, ne)
        log_temperature = np.log10(Te)
        log_density = np.log10(ne)

        c = self.splines[k](log_temperature, log_density)
        return c.diagonal()

    @property
    def temperature_grid(self):
        return 10**(self.log_temperature)

    @property
    def density_grid(self):
        return 10**(self.log_density)

from sys import float_info
class ZeroCoefficient(RateCoefficient):
    def __init__(self):
        pass

    def __call__(self, k, Te, ne):
        Te, ne = np.broadcast_arrays(Te, ne)
        return float_info.min * np.ones_like(Te)