#!/usr/bin/env python
"""
SQLAlchemy wrapping of x-ray database for data from
     Elam et al, Chantler et al, Waasmaier and Kirfel

Main Class for full Database:  xrayDB
"""

import os
import json
from collections import namedtuple
import numpy as np
from scipy.interpolate import UnivariateSpline

from sqlalchemy import MetaData, create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import SingletonThreadPool

from .utils import elam_spline, as_ndarray

XrayEdge = namedtuple('XrayEdge', ('energy', 'fyield', 'jump_ratio'))
XrayLine = namedtuple('XrayLine', ('energy', 'intensity', 'initial_level',
                                   'final_level'))
ElementData = namedtuple('ElementData', ('Z', 'symbol', 'mass', 'density'))

__version__ = '1.4'

def make_engine(dbname):
    "create engine for sqlite connection"
    return create_engine('sqlite:///%s' % (dbname),
                         poolclass=SingletonThreadPool,
                         connect_args={'check_same_thread': False})

def isxrayDB(dbname):
    """whether a file is a valid XrayDB database

    Args:
        dbname (string): name of XrayDB file

    Returns:
        bool: is file a valid XrayDB

    Notes:
      1. must be a sqlite db file, with tables named 'elements',
        'photoabsorption', 'scattering', 'xray_levels', 'Coster_Kronig',
        'Chantler', 'Waasmaier', and 'KeskiRahkonen_Krause'
    """
    _tables = ('Chantler', 'Waasmaier', 'Coster_Kronig',
               'KeskiRahkonen_Krause', 'xray_levels',
               'elements', 'photoabsorption', 'scattering')
    result = False
    try:
        engine = make_engine(dbname)
        meta = MetaData(engine)
        meta.reflect()
        result = all([t in meta.tables for t in _tables])
    except:
        pass
    return result

class XrayDB():
    """
    Database of Atomic and X-ray Data

    This XrayDB object gives methods to access the Atomic and
    X-ray data in th SQLite3 database xraydb.sqlite.

    Much of the data in this database comes from the compilation
    of Elam, Ravel, and Sieber, with additional data from Chantler,
    and other sources. See the documention and bibliography for
    a complete listing.
    """

    def __init__(self, dbname='xraydb.sqlite', read_only=True):
        "connect to an existing database"
        if not os.path.exists(dbname):
            parent, _ = os.path.split(__file__)
            dbname = os.path.join(parent, dbname)
            if not os.path.exists(dbname):
                raise IOError("Database '%s' not found!" % dbname)

        if not isxrayDB(dbname):
            raise ValueError("'%s' is not a valid X-ray Database file!" % dbname)

        self.dbname = dbname
        self.engine = make_engine(dbname)
        self.conn = self.engine.connect()
        kwargs = {}
        if read_only:
            kwargs = {'autoflush': True, 'autocommit': False}

            def readonly_flush(*args, **kwargs):
                return

            self.session = sessionmaker(bind=self.engine, **kwargs)()
            self.session.flush = readonly_flush
        else:
            self.session = sessionmaker(bind=self.engine, **kwargs)()

        self.metadata = MetaData(self.engine)
        self.metadata.reflect()
        self.tables = self.metadata.tables
        elems = self.tables['elements'].select().execute()
        self.atomic_symbols = [e.element for e in elems.fetchall()]

    def close(self):
        "close session"
        self.session.flush()
        self.session.close()

    def query(self, *args, **kws):
        "generic query"
        return self.session.query(*args, **kws)

    def get_version(self, long=False, with_history=False):
        """
        return sqlite3 database and python library version numbers

        Parameters:
            long (bool): show timestamp and notes of latest version [False]
            with_history (bool): show complete version history [False]

        Returns:
            string: version information
        """
        out = []
        rows = self.tables['Version'].select().execute().fetchall()
        if not with_history:
            rows = rows[-1:]
        if long or with_history:
            for row in rows:
                out.append("XrayDB Version: %s [%s] '%s'" % (row.tag,
                                                             row.date,
                                                             row.notes))
            out.append("Python Version: %s" % __version__)
            out = "\n".join(out)
        else:
            out = "XrayDB Version: %s, Python Version: %s" % (rows[0].tag,
                                                              __version__)
        return out

    def f0_ions(self, element=None):
        """
        return list of ion names supported for the .f0() function.

        Parameters:
            element (string, int, pr None):  atomic number, symbol, or ionic symbol
                    of scattering element.

        Returns:
            list:  if element is None, all 211 ions are returned.
                   if element is not None, the ions for that element are returned

        Example:
            >>> xdb = XrayDB()
            >>> xdb.f0_ions('Fe')
            ['Fe', 'Fe2+', 'Fe3+']

        Notes:
            Z values from 1 to 98 (and symbols 'H' to 'Cf') are supported.

        References:
            Waasmaier and Kirfel
        """
        wtab = self.tables['Waasmaier']
        rows = self.query(wtab)
        if element is not None:
            elem = self.symbol(element)
            rows = rows.filter(wtab.c.element == elem)
        return [str(r.ion) for r in rows.all()]

    def f0(self, ion, q):
        """
        return f0(q) -- elastic X-ray scattering factor from Waasmaier and Kirfel

        Parameters:
            ion (string, int, or None):  atomic number, symbol or ionic symbol
                  of scattering element.
            q (float, list, ndarray): value(s) of q for scattering factors

        Returns:
            ndarray: elastic scattering factors

        Example:
            >>> xdb = XrayDB()
            >>> xdb.f0('Fe', range(10))
            array([ 25.994603  ,   6.55945765,   3.21048827,   1.65112769,
                     1.21133507,   1.0035555 ,   0.81012185,   0.61900285,
                     0.43883403,   0.27673021])

        Notes:
            q = sin(theta) / lambda, where theta = incident angle,
            and lambda = X-ray wavelength

            Z values from 1 to 98 (and symbols 'H' to 'Cf') are supported.
            The list of ionic symbols can be read with the function .f0_ions()

        References:
            Waasmaier and Kirfel
        """
        wtab = self.tables['Waasmaier']
        if isinstance(ion, int):
            row = self.query(wtab).filter(wtab.c.atomic_number == ion).all()[0]
        elif ion not in self.f0_ions():
            raise ValueError('No ion {:s} from Waasmaier table'.format(repr(ion)))
        else:
            row = self.query(wtab).filter(wtab.c.ion == ion.title()).all()[0]
        q = as_ndarray(q)
        f0 = row.offset
        for s, e in zip(json.loads(row.scale), json.loads(row.exponents)):
            f0 += s * np.exp(-e*q*q)
        return f0

    def _from_chantler(self, element, energy, column='f1', smoothing=0):
        """
        return energy-dependent data from Chantler table

        Parameters:
            element (string or int): atomic number or symbol.
            eneregy (float or ndarray):
        columns: f1, f2, mu_photo, mu_incoh, mu_total

        Notes:
           this function is meant for internal use.
        """
        ctab = self.tables['Chantler']
        elem = self.symbol(element)
        row = self.query(ctab).filter(ctab.c.element == elem).one()

        energy = as_ndarray(energy)
        emin, emax = min(energy), max(energy)

        te = np.array(json.loads(row.energy))
        nemin = max(0, -3 + max(np.where(te <= emin)[0]))
        nemax = min(len(te), 3 + max(np.where(te <= emax)[0]))

        te = te[nemin:nemax+1]
        if column == 'mu':
            column = 'mu_total'
        ty = np.array(json.loads(getattr(row, column)))[nemin:nemax+1]
        ty[np.where(abs(ty) < 1.e-99)] =  1.e-99
        if column == 'f1':
            out = UnivariateSpline(te, ty, s=smoothing)(energy)
        else:
            out = np.exp(np.interp(np.log(energy),
                                   np.log(te),
                                   np.log(ty)))
        if isinstance(out, np.ndarray) and len(out) == 1:
            out = out[0]
        return out

    def chantler_energies(self, element, emin=0, emax=1.e9):
        """
        return array of energies (in eV) at which data is
        tabulated in the Chantler tables for a particular element.

        Parameters:
            element (string or int): atomic number or symbol
            emin (float): minimum energy (in eV) [0]
            emax (float): maximum energy (in eV) [1.e9]

        Returns:
            ndarray: energies

        References:
            Chantler

        Notes:
            returns 2 energies below emin and above emax to better
            enable interpolation
        """
        ctab = self.tables['Chantler']
        elem = self.symbol(element)
        row = self.query(ctab).filter(ctab.c.element == elem).one()
        te = np.array(json.loads(row.energy))

        if emin <= min(te):
            nemin = 0
        else:
            nemin = max(0, -1 + max(np.where(te <= emin)[0]))
        if emax > max(te):
            nemax = len(te)
        else:
            nemax = min(len(te), 2 + max(np.where(te <= emax)[0]))
        return te[nemin:nemax+1]

    def f1_chantler(self, element, energy, **kws):
        """
        returns f1 -- real part of anomalous X-ray scattering factor
        for selected input energy (or energies) in eV.

        Parameters:
            element (string or int): atomic number or symbol
            energy (float or ndarray): energies (in eV).

        Returns:
            ndarray: real part of anomalous scattering factor

        References:
            Chantler
        """
        return self._from_chantler(element, energy, column='f1', **kws)

    def f2_chantler(self, element, energy, **kws):
        """
        returns f2 -- imaginary part of anomalous X-ray scattering factor
        for selected input energy (or energies) in eV.

        Parameters:
            element (string or int): atomic number or symbol
            energy (float or ndarray): energies (in eV).

        Returns:
            ndarray: imaginary part of anomalous scattering factor

        References:
            Chantler
        """
        return self._from_chantler(element, energy, column='f2', **kws)

    def mu_chantler(self, element, energy, incoh=False, photo=False):
        """
        returns X-ray mass attenuation coefficient, mu/rho in cm^2/gr
        for selected input energy (or energies) in eV.
        default is to return total attenuation coefficient.

        Parameters:
            element (string or int): atomic number or symbol
            energy (float or ndarray): energies (in eV).
            photo (bool): return only the photo-electric contribution [False]
            incoh (bool): return only the incoherent contribution [False]

        Returns:
            ndarray: mass attenuation coefficient in cm^2/gr

        References:
            Chantler
        """
        col = 'mu_total'
        if photo:
            col = 'mu_photo'
        elif incoh:
            col = 'mu_incoh'
        return self._from_chantler(element, energy, column=col)

    def _elem_data(self, element):
        "return data from elements table: internal use"
        etab = self.tables['elements']
        row = self.query(etab)
        if isinstance(element, int):
            row = row.filter(etab.c.atomic_number == element).one()
        else:
            elem = element.title()
            if not elem in self.atomic_symbols:
                raise ValueError("unknown element '%s'" % repr(elem))
            row = row.filter(etab.c.element == elem).one()
        return ElementData(int(row.atomic_number),
                           row.element.title(),
                           row.molar_mass, row.density)

    def atomic_number(self, element):
        """
        return element's atomic number

        Parameters:
            element (string or int): atomic number or symbol

        Returns:
            integer: atomic number
        """
        return self._elem_data(element).Z

    def symbol(self, element):
        """
        return element symbol

        Parameters:
            element (string or int): atomic number or symbol

        Returns:
            string: element symbol
        """
        return self._elem_data(element).symbol

    def molar_mass(self, element):
        """
        return molar mass of element

        Parameters:
            element (string or int): atomic number or symbol

        Returns:
            float: molar mass of element in amu
        """
        return self._elem_data(element).mass

    def density(self, element):
        """
        return density of pure element

        Parameters:
            element (string or int): atomic number or symbol

        Returns:
            float: density of element in gr/cm^3
        """
        return self._elem_data(element).density

    def xray_edges(self, element):
        """
        returns dictionary of X-ray absorption edge energy (in eV),
        fluorescence yield, and jump ratio for an element.

        Parameters:
            element (string or int): atomic number or symbol

        Returns:
            dictionary:  keys of edge (iupac symbol), and values of
                         XrayEdge namedtuple of (energy, fyield, edge_jump))

        References:
           Elam, Ravel, and Sieber.
        """
        elem = self.symbol(element)
        ltab = self.tables['xray_levels']
        out = {}
        for r in self.query(ltab).filter(ltab.c.element == elem).all():
            out[str(r.iupac_symbol)] = XrayEdge(r.absorption_edge,
                                                r.fluorescence_yield,
                                                r.jump_ratio)
        return out

    def xray_edge(self, element, edge):
        """
        returns XrayEdge for an element and edge

        Parameters:
            element (string or int): atomic number or symbol
            edge (string):  X-ray edge

        Returns:
            XrayEdge:  namedtuple of (energy, fyield, edge_jump))

        Example:
            >>> xdb = XrayDB()
            >>> xdb.xray_edge('Co', 'K')
            XrayEdge(edge=7709.0, fyield=0.381903, jump_ratio=7.796)

        References:
           Elam, Ravel, and Sieber.
        """
        return self.xray_edges(element).get(edge.title(), None)

    def xray_lines(self, element, initial_level=None, excitation_energy=None):
        """
        returns dictionary of X-ray emission lines of an element, with

        Parameters:
            initial_level (string or list/tuple of string):  initial level(s) to
                 limit output.
            excitation_energy (float): energy of excitation, limit output those
                 excited by X-rays of this energy (in eV).

        Returns:
            dictionary: keys of lines (Siegbahn symbol), values of Xray Lines

        Notes:
            if both excitation_energy and initial_level are given, excitation_level
            will limit output

        Example:
            >>> xdb = XrayDB()
            >>> for key, val in xdb.xray_lines('Ga', 'K').items():
            >>>      print(key, val)
            'Ka3', XrayLine(energy=9068.0, intensity=0.000326203,
                            initial_level=u'K', final_level=u'L1')
            'Ka2', XrayLine(energy=9223.8, intensity=0.294438,
                            initial_level=u'K', final_level=u'L2')
            'Ka1', XrayLine(energy=9250.6, intensity=0.57501,
                            initial_level=u'K', final_level=u'L3')
            'Kb3', XrayLine(energy=10263.5, intensity=0.0441511,
                            initial_level=u'K', final_level=u'M2')
            'Kb1', XrayLine(energy=10267.0, intensity=0.0852337,
                            initial_level=u'K', final_level=u'M3')
            'Kb5', XrayLine(energy=10348.3, intensity=0.000841354,
                            initial_level=u'K', final_level=u'M4,5')

        References:
           Elam, Ravel, and Sieber.
        """
        elem = self.symbol(element)
        ttab = self.tables['xray_transitions']
        row = self.query(ttab).filter(ttab.c.element == elem)
        if excitation_energy is not None:
            initial_level = []
            for ilevel, dat in self.xray_edges(elem).items():
                if dat[0] < excitation_energy:
                    initial_level.append(ilevel.title())

        if initial_level is not None:
            if isinstance(initial_level, (list, tuple)):
                row = row.filter(ttab.c.initial_level.in_(initial_level))
            else:
                row = row.filter(ttab.c.initial_level == initial_level.title())
        out = {}
        for r in row.all():
            out[str(r.siegbahn_symbol)] = XrayLine(r.emission_energy, r.intensity,
                                                   r.initial_level, r.final_level)
        return out

    def xray_line_strengths(self, element, excitation_energy=None):
        """
        return the absolute line strength in cm^2/gr for all available lines

        Parameters:
            element (string or int): Atomic symbol or number for element
            excitation_energy (float): incident energy, in eV

        Returns:
            dictionary: elemental line with fluorescence cross section in cm2/gr.

        References:
           Elam, Ravel, and Sieber.
        """
        out = {}
        lines = self.xray_lines(element, excitation_energy=excitation_energy)
        for label, eline in lines.items():
            edge = self.xray_edge(element, eline.initial_level)
            if edge is None and ',' in eline.initial_level:
                ilevel, _ = eline.initial_level.split(',')
                edge = self.xray_edge(element, ilevel)
            if edge is not None:
                mu = self.mu_elam(element, [edge.energy*(0.999),
                                            edge.energy*(1.001)], kind='photo')
                out[label] = (mu[1]-mu[0]) * eline.intensity * edge.fyield
        return out

    def ck_probability(self, element, initial, final, total=True):
        """
        return Coster-Kronig transition probability for an element and
        initial/final levels

        Parameters:
            element (string or int): Atomic symbol or number for element
            initial (string):  initial level
            final (string):  final level
            total (bool): whether to return total or partial probability

        Returns:
            float: transition probability

        Example:
            >>> xdb = XrayDB()
            >>> xdb.ck_probability('Cu', 'L1', 'L3', total=True)
            0.681

        References:
           Elam, Ravel, and Sieber.
        """
        elem = self.symbol(element)
        ctab = self.tables['Coster_Kronig']

        row = self.query(ctab).filter(ctab.c.element == elem)
        row = row.filter(ctab.c.initial_level == initial.title())
        row = row.filter(ctab.c.final_level == final.title()).all()
        out = 0.0
        if len(row) > 0:
            row = row[0]
            out = row.transition_probability
            if total:
                out = row.total_transition_probability
        return out

    def corehole_width(self, element, edge=None, use_keski=False):
        """
        returns core hole width for an element and edge

        Parameters:
            element (string, integer): atomic number or symbol for element
            edge (string or None): edge for hole, return all if None
            use_keski (bool) : force use of KeskiRahkonen and Krause table for all data.

        Returns:
            float: corehole width in eV.

        Notes:
            Uses Krause and Oliver where data is available (K, L lines Z > 10)
            Uses Keski-Rahkonen and Krause otherwise

        References:
            Krause and Oliver, 1979
            Keski-Rahkonen and Krause, 1974

        """
        version_qy = self.tables['Version'].select().order_by('date')
        version_id = version_qy.execute().fetchall()[-1].id

        ctab = self.tables['corelevel_widths']
        if version_id < 4 or use_keski:
            ctab = self.tables['KeskiRahkonen_Krause']

        rows = self.query(ctab).filter(ctab.c.element == self.symbol(element))
        if edge is not None:
            rows = rows.filter(ctab.c.edge == edge.title())
        result = rows.all()
        if len(result) == 1:
            result = result[0].width
        else:
            result = [(r.edge, r.width) for r in result]
        return result

    def cross_section_elam(self, element, energies, kind='photo'):
        """
        returns Elam Cross Section values for an element and energies

        Parameters:
            element (string or int):  atomic number or symbol for element
            energies (float or ndarray): energies (in eV) to calculate cross-sections
            kind (string):  one of 'photo', 'coh', and 'incoh' for photo-absorption,
                  coherent scattering, and incoherent scattering cross sections,
                  respectively. Default is 'photo'.

        Returns:
            ndarray of scattering data

        References:
            Elam, Ravel, and Sieber.
        """
        elem = self.symbol(element)
        kind = kind.lower()
        if kind not in ('coh', 'incoh', 'photo'):
            raise ValueError('unknown cross section kind=%s' % kind)

        stab = self.tables['scattering']
        if kind == 'photo':
            stab = self. tables['photoabsorption']

        row = self.query(stab).filter(stab.c.element == elem).all()[0]

        tab_lne = np.array(json.loads(row.log_energy))
        if kind.startswith('coh'):
            tab_val = np.array(json.loads(row.log_coherent_scatter))
            tab_spl = np.array(json.loads(row.log_coherent_scatter_spline))
        elif kind.startswith('incoh'):
            tab_val = np.array(json.loads(row.log_incoherent_scatter))
            tab_spl = np.array(json.loads(row.log_incoherent_scatter_spline))
        else:
            tab_val = np.array(json.loads(row.log_photoabsorption))
            tab_spl = np.array(json.loads(row.log_photoabsorption_spline))

        en = 1.0*as_ndarray(energies)
        emin_tab = 10*int(0.102*np.exp(tab_lne[0]))
        en[np.where(en < emin_tab)] = emin_tab
        out = np.exp(elam_spline(tab_lne, tab_val, tab_spl, np.log(en)))
        if len(out) == 1:
            return out[0]
        return out

    def mu_elam(self, element, energies, kind='total'):
        """
        returns attenuation cross section for an element at energies (in eV)

        Parameters:
            element (string or int):  atomic number or symbol for element
            energies (float or ndarray): energies (in eV) to calculate cross-sections
            kind (string):  one of 'photo' or 'total' for photo-electric or
                  total attenuation, respectively.  Default is 'total'.

        Returns:
           ndarray of scattering values in units of cm^2/gr

        References:
            Elam, Ravel, and Sieber.
        """
        calc = self.cross_section_elam
        kind = kind.lower()
        if kind.startswith('tot'):
            xsec = calc(element, energies, kind='photo')
            xsec += calc(element, energies, kind='coh')
            xsec += calc(element, energies, kind='incoh')
        elif kind.startswith('photo'):
            xsec = calc(element, energies, kind='photo')
        elif kind.lower().startswith('coh'):
            xsec = calc(element, energies, kind='coh')
        elif kind.lower().startswith('incoh'):
            xsec = calc(element, energies, kind='incoh')
        else:
            raise ValueError('unknown cross section kind=%s' % kind)
        return xsec

    def coherent_cross_section_elam(self, element, energies):
        """returns coherenet scattering crossrxr section for an element
        at energies (in eV)

        returns values in units of cm^2 / gr

        arguments
        ---------
        element:  atomic number, atomic symbol for element
        energies: energies in eV to calculate cross-sections

        Data from Elam, Ravel, and Sieber.
        """
        return self.cross_section_elam(element, energies, kind='coh')

    def incoherent_cross_section_elam(self, element, energies):
        """return incoherenet scattering cross section for an element
        at energies (in eV)

        returns values in units of cm^2 / gr

        arguments
        ---------
        element:  atomic number, atomic symbol for element
        energies: energies in eV to calculate cross-sections

        Data from Elam, Ravel, and Sieber.
        """
        return self.cross_section_elam(element, energies, kind='incoh')

    def ionization_potential(self, gas):
        """return effective ionization potential for a gas,
        as appropriate for ionization chambers in the linear
        regime (not in the 'proportional counter' regime)

        returns values in units of eV

        argument
        ---------
        gas (string):  name of gas

        notes
        -----
        Data from G. F. Knoll, Radiation Detection and Measurement, Table 5-1.
        """
        itab = self.tables['ionization_potentials']
        out = self.query(itab).filter(itab.c.gas == gas).all()
        if len(out) != 1:
            raise ValueError('unknown gas for ionization potential: %s' % gas)
        return float(out[0].potential)


# -*- coding: utf-8 -*-
# pylint: disable=wrong-import-position, range-builtin-not-iterating

"""
Physical constants and constant-for-us values
"""

from __future__ import absolute_import, division, print_function

__all__ = [
    'omkeys_to_sd_indices', 'get_sd_idx', 'get_string_dom_pair',

    # Constants
    'PI', 'TWO_PI', 'PI_BY_TWO', 'SPEED_OF_LIGHT_M_PER_NS',

    # Pre-calculated values
    'COS_CKV', 'THETA_CKV', 'SIN_CKV',
    'TRACK_M_PER_GEV', 'TRACK_PHOTONS_PER_M', 'CASCADE_PHOTONS_PER_GEV',
    'IC_DOM_JITTER_NS', 'DC_DOM_JITTER_NS', 'POL_TABLE_DCOSTHETA',
    'POL_TABLE_DRPWR', 'POL_TABLE_DT', 'POL_TABLE_RPWR', 'POL_TABLE_RMAX',
    'POL_TABLE_NTBINS', 'POL_TABLE_NRBINS', 'POL_TABLE_NTHETABINS',
    'IC_DOM_QUANT_EFF', 'DC_DOM_QUANT_EFF',

    # Particle naming conventions
    'ABS_FLAV_STR', 'ABS_FLAV_TEX', 'BAR_NOBAR_STR', 'BAR_NOBAR_TEX',
    'INT_TYPE_STR', 'INT_TYPE_TEX', 'PDG_STR', 'PDG_TEX', 'PDG_INTER_STR',
    'PDG_INTER_TEX', 'STR_TO_PDG_INTER',

    # "Enum"-like things
    'STR_ALL', 'STR_IC', 'STR_DC', 'AGG_STR_NONE', 'AGG_STR_ALL',
    'AGG_STR_SUBDET', 'DOM_ALL',

    'NUM_STRINGS', 'NUM_DOMS_PER_STRING', 'NUM_DOMS_TOT',

    'IC_STRS', 'DC_STRS', 'DC_IC_STRS', 'DC_ALL_STRS', 'DC_SUBDUST_DOMS',
    'IC_SUBDUST_DOMS', 'DC_SUBDUST_STRS_DOMS', 'DC_IC_SUBDUST_STRS_DOMS',
    'DC_ALL_SUBDUST_STRS_DOMS', 'ALL_STRS', 'ALL_DOMS', 'ALL_STRS_DOMS',
    'ALL_STRS_DOMS_SET', 'DC_ALL_STRS_DOMS',

    'EMPTY_HITS', 'EMPTY_SOURCES',
    'SRC_OMNI', 'SRC_CKV_BETA1',

    'PARAM_NAMES', 'PEGLEG_PARAM_NAMES', 'SCALING_PARAM_NAMES',
]

__author__ = 'P. Eller, J.L. Lanfranchi'
__license__ = '''Copyright 2017 Philipp Eller and Justin L. Lanfranchi

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

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

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

from itertools import product
from os.path import abspath, dirname
import sys

import numpy as np

if __name__ == '__main__' and __package__ is None:
    RETRO_DIR = dirname(dirname(dirname(abspath(__file__))))
    if RETRO_DIR not in sys.path:
        sys.path.append(RETRO_DIR)
from retro import FTYPE
from retro import retro_types

def omkeys_to_sd_indices(omkeys):
    """Get a single integer index from OMKeys.

    Parameters
    ----------
    omkeys : array of dtype OMKEY_T
        The dtype `OMKEY_T` must contain "string" and "dom" and can optionally
        include "pmt".

    Returns
    -------
    sd_idx : array of np.uint32

    """
    if 'pmt' in omkeys.dtype.names:
        raise NotImplementedError("OMKey field 'pmt' not implemented")
    return get_sd_idx(string=omkeys['string'], dom=omkeys['dom'])

def get_sd_idx(string, dom, pmt=0):
    """Get a single integer index from an IceCube string number (from 1 to 86)
    and DOM number (from 1 to 60).

    Parameters
    ----------
    string : int in [1, 60]
        String number
    dom : int in [1, 60]
        DOM number
    pmt : int
        PMT number in the DOM; if == 0, then this is ignored.

    Returns
    -------
    sd_idx : int

    """
    if pmt > 0:
        raise NotImplementedError('PMT != 0 is not implemented')
    return (dom - 1) * NUM_STRINGS + (string - 1)

def get_string_dom_pair(sd_idx):
    """Get an IceCube string number (1 to 86) and a DOM number (1 to 60) from
    the single-integer index (sd_idx).

    Parameters
    ----------
    sd_idx : int in [0, 5159]

    Returns
    -------
    string : int in [1, 86]
    dom : int in [1, 60]

    """
    dom_idx, string_idx = divmod(sd_idx, NUM_STRINGS)
    string = string_idx + 1
    dom = dom_idx + 1
    return string, dom

# -- Physical / mathematical constants -- #

PI = FTYPE(np.pi)
"""pi"""

TWO_PI = FTYPE(2*np.pi)
"""2 * pi"""

PI_BY_TWO = FTYPE(np.pi / 2)
"""pi / 2"""

SPEED_OF_LIGHT_M_PER_NS = FTYPE(299792458 / 1e9)
"""Speed of light in units of m/ns"""

# -- Pre-calculated values -- #

COS_CKV = 0.764540803152
"""Cosine of the Cherenkov angle for beta ~1 and IceCube phase index as used"""

THETA_CKV = np.arccos(0.764540803152)
"""Cosine of the Cherenkov angle for beta ~1 and IceCube phase index as used"""

SIN_CKV = np.sin(THETA_CKV)
"""Cosine of the Cherenkov angle for beta ~1 and IceCube phase index as used"""

TRACK_M_PER_GEV = FTYPE(15 / 3.3)
"""Track length per energy, in units of m/GeV"""

TRACK_PHOTONS_PER_M = FTYPE(2451.4544553)
"""Track photons per length, in units of 1/m (see ``nphotons.py``)"""

CASCADE_PHOTONS_PER_GEV = FTYPE(12805.3383311)
"""Cascade photons per energy, in units of 1/GeV (see ``nphotons.py``)"""

# TODO: Is jitter same (or close enough to the same) for all DOMs? Is it
#       different for DeepCore vs. non-DeepCore DOMs? Didn't see as much in
#       section 3.3. of arXiv:1612.05093v2 so assuming same for now.

# See arXiv:1612.05093v2, section 3.3
IC_DOM_JITTER_NS = 1.7
"""Timing jitter (stddev) for string 0-79 DOMs, in units of ns"""

# See arXiv:1612.05093v2, section 3.3
DC_DOM_JITTER_NS = 1.7
"""Timing jitter (stddev) for DeepCore (strings 80-86) DOMs, in units of ns"""

# TODO: figure these out from the tables rather than defining as constants
POL_TABLE_RMAX = 400 # m
POL_TABLE_DT = 10 # ns
POL_TABLE_RPWR = 2
POL_TABLE_DRPWR = 0.1
POL_TABLE_DCOSTHETA = -0.05
POL_TABLE_NTBINS = 300
POL_TABLE_NRBINS = 200
POL_TABLE_NTHETABINS = 40

#IC_DOM_QUANT_EFF = 0.25
IC_DOM_QUANT_EFF = 1.
"""scalar in [0, 1] : (Very rough approximation!) IceCube (i.e. non-DeepCore)
DOM quantum efficiency. Multiplies the tabulated detection probabilities to
yield the actual probabilitiy that a photon is detected."""
#DC_DOM_QUANT_EFF = 0.35
DC_DOM_QUANT_EFF = 1.
"""scalar in [0, 1] : (Very rough approximation!) DeepCore DOM quantum
efficiency. Multiplies the tabulated detection probabilities to yield the
actual probabilitiy that a photon is detected."""

# -- Particle / interaction type naming conventions -- #

ABS_FLAV_STR = {12: 'nue', 13: 'numu', 14: 'nutau'}
ABS_FLAV_TEX = {12: r'\nu_e', 13: r'\nu_\mu', 14: r'\nu_\tau'}

BAR_NOBAR_STR = {-1: 'bar', 1: ''}
BAR_NOBAR_TEX = {-1: r'\bar', 1: ''}

INT_TYPE_STR = {1: 'cc', 2: 'nc'}
INT_TYPE_TEX = {1: r'\, {\rm CC}', 2: r'\, {\rm NC}'}

PDG_STR = {}
PDG_TEX = {}
for _bnb, _abs_code in product(BAR_NOBAR_STR.keys(), ABS_FLAV_STR.keys()):
    PDG_STR[_abs_code*_bnb] = ABS_FLAV_STR[_abs_code] + BAR_NOBAR_STR[_bnb]
    PDG_TEX[_abs_code*_bnb] = BAR_NOBAR_TEX[_bnb] + ABS_FLAV_TEX[_abs_code]

PDG_INTER_STR = {}
PDG_INTER_TEX = {}
for _pdg, _it in product(PDG_STR.keys(), INT_TYPE_STR.keys()):
    PDG_INTER_STR[(_pdg, _it)] = '%s_%s' % (PDG_STR[_pdg], INT_TYPE_STR[_it])
    PDG_INTER_TEX[(_pdg, _it)] = '%s %s' % (PDG_TEX[_pdg], INT_TYPE_TEX[_it])

STR_TO_PDG_INTER = {v: k for k, v in PDG_INTER_STR.items()}

# -- "enums" -- #
STR_ALL, STR_IC, STR_DC = -1, -2, -3
AGG_STR_NONE, AGG_STR_ALL, AGG_STR_SUBDET = 0, 1, 2
DOM_ALL = -1

# -- geom constants --- #

NUM_STRINGS = 86
NUM_DOMS_PER_STRING = 60
NUM_DOMS_TOT = NUM_STRINGS * NUM_DOMS_PER_STRING

IC_STRS = np.array(range(1, 78+1), dtype=np.uint8)
DC_STRS = np.array(range(79, 86+1), dtype=np.uint8)
DC_IC_STRS = np.array([26, 27, 35, 36, 37, 45, 46], dtype=np.uint8)
DC_ALL_STRS = np.concatenate([DC_STRS, DC_IC_STRS], axis=0)

DC_SUBDUST_DOMS = np.array(range(11, 60+1), dtype=np.uint8)
IC_SUBDUST_DOMS = np.array(range(25, 60+1), dtype=np.uint8)

DC_SUBDUST_STRS_DOMS = np.array(
    [get_sd_idx(s, d) for s, d in product(DC_STRS, DC_SUBDUST_DOMS)]
)
DC_IC_SUBDUST_STRS_DOMS = np.array(
    [get_sd_idx(s, d) for s, d in product(DC_IC_STRS, IC_SUBDUST_DOMS)]
)

DC_ALL_SUBDUST_STRS_DOMS = np.concatenate(
    (DC_SUBDUST_STRS_DOMS, DC_IC_SUBDUST_STRS_DOMS)
)

ALL_STRS = list(range(1, 86+1))
ALL_DOMS = list(range(1, 60+1))
ALL_STRS_DOMS = np.array([get_sd_idx(s, d) for s, d in product(ALL_STRS, ALL_DOMS)])
ALL_STRS_DOMS_SET = set(ALL_STRS_DOMS)
DC_ALL_STRS_DOMS = np.array([get_sd_idx(s, d) for s, d in product(DC_STRS, ALL_DOMS)])

EMPTY_HITS = np.empty(shape=0, dtype=retro_types.HIT_T)

EMPTY_SOURCES = np.empty(shape=0, dtype=retro_types.SRC_T)

SRC_OMNI = np.uint32(0)
"""Source kind designator for a point emitting omnidirectional light"""

SRC_CKV_BETA1 = np.uint32(1)
"""Source kind designator for a point emitting Cherenkov light with beta ~ 1"""

PARAM_NAMES = [
    'time', 'x', 'y', 'z', 'track_azimuth', 'track_zenith', 'cascade_azimuth',
    'cascade_zenith', 'track_energy', 'cascade_energy', 'cascade_d_zenith',
    'cascade_d_azimuth'
]
"""All possible hypothesis param names"""

PEGLEG_PARAM_NAMES = ['track_energy']
"""Hypothesis param names handled by pegleg, if it's used"""

SCALING_PARAM_NAMES = ['cascade_energy']
"""Hypothesis param names handled by scaling, if it's used"""


import numpy as np
import scipy
if tuple(map(int, scipy.__version__.split('.'))) < (1, 0, 0):
    from scipy.misc import logsumexp
else:
    from scipy.special import logsumexp
import math
import time
from datetime import date
import os
import pickle

def normalize_features(features):
	'''features: n by d matrix'''
	assert(len(features.shape)==2)
	norma=np.sqrt(np.sum(features ** 2, axis=1).reshape(-1, 1))+1e-6
	return features/norma

class vMFMM:
	def __init__(self, cls_num, init_method = 'random', tmp_dir = '/tmp/'):
		self.cls_num = cls_num
		self.init_method = init_method

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

		self.tmp_file = os.path.join(tmp_dir, str(date.today())+'.pickle')

	def fit(self, features, kappa, max_it=300, tol = 5e-5, normalized=False, verbose=True):
		self.features = features
		if not normalized:
			self.features = normalize_features(features)

		self.n, self.d = self.features.shape
		self.kappa = kappa

		self.pi = np.random.random(self.cls_num)
		self.pi /= np.sum(self.pi)
		if self.init_method =='random':
			self.mu = np.random.random((self.cls_num, self.d))
			# self.mu = np.array([[1,0],[0,1]])
			self.mu = normalize_features(self.mu)
		elif self.init_method =='k++':
			#print('start k++')
			centers = []
			centers_i = []

			if self.n > 50000:
				rdn_index = np.random.choice(self.n, size=(50000,), replace=False)
			else:
				rdn_index = np.array(range(self.n), dtype=int)

			cos_dis = 1-np.dot(self.features[rdn_index], self.features[rdn_index].T)

			#print('finish cos_dis')
			centers_i.append(np.random.choice(rdn_index))
			centers.append(self.features[centers_i[0]])
			for i in range(self.cls_num-1):
				#if i%10==0:
			#		print('k++ center {0}'.format(i))

				cdisidx = [np.where(rdn_index==cci)[0][0] for cci in centers_i]
				prob = np.min(cos_dis[:,cdisidx], axis=1)**2
				prob /= np.sum(prob)
				centers_i.append(np.random.choice(rdn_index, p=prob))
				centers.append(self.features[centers_i[-1]])

			self.mu = np.array(centers)
			del(cos_dis)
			#print('finish k++')

		self.mllk_rec = []
		for itt in range(max_it):
			_st = time.time()
			self.e_step()
			self.m_step()
			_et = time.time()
			#if verbose and itt%1==0:
				#print("iter {0}: {1}, time: {2}".format(itt, self.mllk, (_et-_st)/60))

			if itt%20==0:
				with open(self.tmp_file, 'wb') as fh:
					pickle.dump(self.mu, fh)

				bins = 4
				per_bin = self.cls_num//bins+1
				for bb in range(bins):
					with open(self.tmp_file.replace('.pickle','_p{}.pickle'.format(bb)), 'wb') as fh:
						pickle.dump(self.p[:,bb*per_bin:(bb+1)*per_bin], fh)

			self.mllk_rec.append(self.mllk)
			if len(self.mllk_rec)>1 and self.mllk - self.mllk_rec[-2] < tol:
				print("early stop at iter {0}, llk {1}".format(itt, self.mllk))
				break

	def fit_soft(self, features, p, mu, pi, kappa, max_it=300, tol = 1e-6, normalized=False, verbose=True):
		self.features = features
		if not normalized:
			self.features = normalize_features(features)

		self.p = p
		self.mu = mu
		self.pi = pi
		self.kappa = kappa

		self.n, self.d = self.features.shape

		for itt in range(max_it):
			self.e_step()
			self.m_step()
			#if verbose and itt%20==0:
				#print("iter {0}: {1}".format(itt, self.mllk))

			self.mllk_rec.append(self.mllk)
			if len(self.mllk_rec)>1 and self.mllk - self.mllk_rec[-2] < tol:
				print("early stop at iter {0}, llk {1}".format(itt, self.mllk))
				break

	def e_step(self):
		# update p
		logP = np.dot(self.features, self.mu.T)*self.kappa + np.log(self.pi).reshape(1,-1)  # n by k
		logP_norm = logP - logsumexp(logP, axis=1).reshape(-1,1)
		self.p = np.exp(logP_norm)
		self.mllk = np.mean(logsumexp(logP, axis=1))

	def m_step(self):
		# update pi and mu
		self.pi = np.sum(self.p, axis=0)/self.n

		# fast version, requires more memory
		self.mu = np.dot(self.p.T, self.features)/np.sum(self.p, axis=0).reshape(-1,1)

#         d_cut = 52
#         bnum = int(math.ceil(self.d/d_cut))

#         for dd_i in range(bnum):
#             dd_start = dd_i*d_cut
#             dd_end = min((dd_i+1)*d_cut, self.d)
#             self.mu[:,dd_start:dd_end] = np.sum(np.tile(self.features.reshape(self.n,1,self.d)[:,:,dd_start:dd_end],(1,self.cls_num,1))*self.p.reshape(self.n,self.cls_num,1),axis=0)/np.sum(self.p, axis=0).reshape(-1,1)

		# for cc in range(self.cls_num):
		#     self.mu[cc] = np.sum(self.p[:,cc].reshape(-1,1) * self.features, axis=0)/np.sum(self.p[:,cc])

		# r = np.mean(np.sqrt(np.sum(self.mu**2, axis=1))*self.pi)
		# r = np.mean(np.sqrt(np.sum(self.mu**2, axis=1))/(self.n*self.pi))
		# self.kappa2 = (r*self.d-r**3)/(1-r**2)

		self.mu = normalize_features(self.mu)


# coding: utf-8
# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
# Distributed under the terms of "New BSD License", see the LICENSE file.

from __future__ import print_function
from collections import OrderedDict
import numpy as np
from pyiron.base.generic.parameters import GenericParameters
import decimal as dec

try:
    from ase.calculators.lammps import Prism

except ImportError:
    try:
        from ase.calculators.lammpsrun import Prism
    except ImportError:
        from ase.calculators.lammpsrun import prism as Prism

__author__ = "Joerg Neugebauer, Sudarsan Surendralal, Yury Lysogorskiy, Jan Janssen, Markus Tautschnig"
__copyright__ = "Copyright 2019, Max-Planck-Institut für Eisenforschung GmbH - " \
                "Computational Materials Design (CM) Department"
__version__ = "1.0"
__maintainer__ = "Sudarsan Surendralal"
__email__ = "surendralal@mpie.de"
__status__ = "production"
__date__ = "Sep 1, 2017"

class UnfoldingPrism(Prism):
    """
    Create a lammps-style triclinic prism object from a cell

    The main purpose of the prism-object is to create suitable
    string representations of prism limits and atom positions
    within the prism.
    When creating the object, the digits parameter (default set to 10)
    specify the precision to use.
    lammps is picky about stuff being within semi-open intervals,
    e.g. for atom positions (when using create_atom in the in-file),
    x must be within [xlo, xhi).

    Args:
        cell:
        pbc:
        digits:
    """
    def __init__(self, cell, pbc=(True, True, True), digits=10):
        # Temporary fix. Since the arguments for the constructor have changed, try to see if it is compatible with
        # the latest ase. If not, revert to the old __init__ parameters.
        try:
            super(UnfoldingPrism, self).__init__(cell, pbc=pbc, tolerance=float('1e-{}'.format(digits)))
        except TypeError:
            super(UnfoldingPrism, self).__init__(cell, pbc=pbc, digits=digits)
        a, b, c = cell
        an, bn, cn = [np.linalg.norm(v) for v in cell]

        alpha = np.arccos(np.dot(b, c) / (bn * cn))
        beta = np.arccos(np.dot(a, c) / (an * cn))
        gamma = np.arccos(np.dot(a, b) / (an * bn))

        xhi = an
        xyp = np.cos(gamma) * bn
        yhi = np.sin(gamma) * bn
        xzp = np.cos(beta) * cn
        yzp = (bn * cn * np.cos(alpha) - xyp * xzp) / yhi
        zhi = np.sqrt(cn ** 2 - xzp ** 2 - yzp ** 2)

        # Set precision
        self.car_prec = dec.Decimal('10.0') ** \
                        int(np.floor(np.log10(max((xhi, yhi, zhi)))) - digits)
        self.dir_prec = dec.Decimal('10.0') ** (-digits)
        self.acc = float(self.car_prec)
        self.eps = np.finfo(xhi).eps

        # For rotating positions from ase to lammps
        apre = np.array(((xhi, 0, 0),
                         (xyp, yhi, 0),
                         (xzp, yzp, zhi)))
        self.R = np.dot(np.linalg.inv(cell), apre)

        # Actual lammps cell may be different from what is used to create R
        eps = 1e-10

        def fold(vec, pvec, i):
            p = pvec[i]
            x = vec[i] + 0.5 * p
            n = (np.mod(x, p) - x) / p
            # print ('prism: ', n, x, p)
            return [float(self.f2qdec(vec_a)) for vec_a in (vec + n * pvec)], n

        apre[1, :], n1 = fold(apre[1, :], apre[0, :], 0)
        if np.abs(apre[1, 0] / apre[0, 0]) > 0.5:
            apre[1, 0] -= np.sign(n1) * apre[0, 0]
            n1 -= np.sign(n1)

        apre[2, :], n2 = fold(apre[2, :], apre[1, :], 1)
        if np.abs(apre[2, 1] / apre[1, 1]) > 0.5:
            apre[2, 1] -= np.sign(n2) * apre[1, 1]
            n2 -= np.sign(n2)

        apre[2, :], n3 = fold(apre[2, :], apre[0, :], 0)
        if np.abs(apre[2, 0] / apre[0, 0]) > 0.5:
            apre[2, 0] -= np.sign(n3) * apre[0, 0]
            n3 -= np.sign(n3)
        self.ns = [n1, n2, n3]

        d_a = apre[0, 0]/2 - apre[1,0]
        if np.abs(d_a) < eps:
            if d_a < 0:
                print ('debug: apply shift')
                apre[1, 0] += 2 * d_a
                apre[2, 0] += 2 * d_a
        self.A = apre
        self.Ainv = np.linalg.inv(self.A)
        self.prism = None

        if self.is_skewed() and \
                (not (pbc[0] and pbc[1] and pbc[2])):
            raise RuntimeError('Skewed lammps cells MUST have '
                               'PBC == True in all directions!')

    def unfold_cell(self, cell):
        """
        Unfold LAMMPS cell to original

        Args:
            cell: LAMMPS cell,

        Returns:
            unfolded cell
        """
        a = cell[0]
        bp = cell[1]
        cpp = cell[2]
        (n1, n2, n3) = self.ns
        b = bp - n1 * a
        c = cpp - n2 * bp - n3 * a
        return np.array([a, b, c])

    def pos_to_lammps(self, position):
        """
        Rotate an ase-cell position to the lammps cell orientation

        Args:
            position:

        Returns:
            tuple of float.
        """
        return tuple([x for x in np.dot(position, self.R)])

    def f2qdec(self, f):
        return dec.Decimal(repr(f)).quantize(self.car_prec, dec.ROUND_DOWN)

    def f2s(self, f):
        return str(dec.Decimal(repr(f)).quantize(self.car_prec, dec.ROUND_HALF_EVEN))

    def get_lammps_prism_str(self):
        """Return a tuple of strings"""
        p = self.get_lammps_prism()
        return tuple([self.f2s(x) for x in p])

class LammpsStructure(GenericParameters):
    """

    Args:
        input_file_name:
    """
    def __init__(self, input_file_name=None):
        super(LammpsStructure, self).__init__(input_file_name=input_file_name,
                                              table_name="structure_inp",
                                              comment_char="#",
                                              val_only=True)
        self._structure = None
        self._potential = None
        self._el_eam_lst = []
        self.atom_type = None
        self.cutoff_radius = None
        self.digits = 10

    @property
    def potential(self):
        return self._potential

    @potential.setter
    def potential(self, val):
        self._potential = val

    @property
    def structure(self):
        """

        Returns:

        """
        return self._structure

    @structure.setter
    def structure(self, structure):
        """

        Args:
            structure:

        Returns:

        """
        self._structure = structure
        if self.atom_type == 'full':
            input_str = self.structure_full()
        elif self.atom_type == 'bond':
            input_str = self.structure_bond()
        elif self.atom_type == 'charge':
            input_str = self.structure_charge()
        else:  # self.atom_type == 'atomic'
            input_str = self.structure_atomic()
        self.load_string(input_str)

    @property
    def el_eam_lst(self):
        """

        Returns:

        """
        return self._el_eam_lst

    @el_eam_lst.setter
    def el_eam_lst(self, el_eam_lst):
        """

        Args:
            el_eam_lst:

        Returns:

        """
        self._el_eam_lst = el_eam_lst

    def load_default(self):
        """

        Returns:

        """
        input_str = ''
        self.load_string(input_str)

    # def f2s(self, f):
    #     return str(dec.Decimal(repr(f)).quantize(self.car_prec,
    #                                              dec.ROUND_HALF_EVEN))

    def simulation_cell(self):
        """

        Returns:

        """
        # dim = self._structure.dimension
        # amat = self._structure.cell

        self.prism = UnfoldingPrism(self._structure.cell, digits=15)
        xhi, yhi, zhi, xy, xz, yz = self.prism.get_lammps_prism_str()
        # Please, be carefull and not round xhi, yhi,..., otherwise you will get too skew cell from LAMMPS.
        # These values are already checked in UnfoldingPrism to fullfill LAMMPS skewness criteria
        simulation_cell = '0. {} xlo xhi\n'.format(xhi) + \
                          '0. {} ylo yhi\n'.format(yhi) + \
                          '0. {} zlo zhi\n'.format(zhi)

        if self.prism.is_skewed():
            simulation_cell += '{0} {1} {2} xy xz yz\n'.format(xy, xz, yz)

        # if s.VERBOSE():
        #     s.warning("triclinic cells not supported for test purposes",
        #               module="lammps.lammps.LammpsStructure.simulation_cell")
        return simulation_cell

    def structure_bond(self):
        """

        Returns:

        """
        # analyze structure to get molecule_ids, bonds, angles etc
        coords = self.rotate_positions(self._structure)

        elements = self._structure.get_chemical_symbols()
        el_list = self._structure.get_species_symbols()
        el_dict = OrderedDict()
        for object_id, el in enumerate(el_list):
            el_dict[el] = object_id

        n_s = len(el_list)
        bond_type = np.ones([n_s, n_s], dtype=np.int)
        count = 0
        for i in range(n_s):
            for j in range(i, n_s):
                count += 1
                bond_type[i, j] = count
                bond_type[j, i] = count
        # print "bond_type: ", bond_type

        if self.structure.bonds is None:
            if self.cutoff_radius is None:
                bonds_lst = self.structure.get_bonds(max_shells=1)
            else:
                bonds_lst = self.structure.get_bonds(radius=self.cutoff_radius)
            bonds = []
            # id_mol = 0
            for ia, i_bonds in enumerate(bonds_lst):
                el_i = el_dict[elements[ia]]
                for el_j, b_lst in i_bonds.items():
                    b_type = bond_type[el_i][el_dict[el_j]]
                    for i_shell, ib_shell_lst in enumerate(b_lst):
                        for ib in np.unique(ib_shell_lst):
                            if ia < ib:  # avoid double counting of bonds
                                bonds.append([ia + 1, ib + 1, b_type])

            self.structure.bonds = np.array(bonds)
        bonds = self.structure.bonds

        atomtypes = ' Start File for LAMMPS \n' + \
                    '{0:d} atoms'.format(len(self._structure)) + ' \n' + \
                    '{0:d} bonds'.format(len(bonds)) + ' \n' + \
                    '{0} atom types'.format(self._structure.get_number_of_species()) + ' \n' + \
                    '{0} bond types'.format(np.max(bond_type)) + ' \n'

        cell_dimesions = self.simulation_cell()

        masses = 'Masses \n\n'
        el_obj_list = self._structure.get_species_objects()
        for object_id, el in enumerate(el_obj_list):
            masses += '{0:3d} {1:f}'.format(object_id + 1, el.AtomicMass) + '\n'

        atoms = 'Atoms \n\n'

        # atom_style bond
        # format: atom-ID, molecule-ID, atom_type, x, y, z
        format_str = '{0:d} {1:d} {2:d} {3:f} {4:f} {5:f} '
        if self._structure.dimension == 3:
            for id_atom, (x, y, z) in enumerate(coords):
                id_mol, id_species = 1, el_dict[elements[id_atom]]  # elList[id_atom]
                # print id_atom + 1, id_mol, id_species + 1, x, y, z
                atoms += format_str.format(id_atom + 1, id_mol, id_species + 1, x, y, z) + '\n'
        elif self._structure.dimension == 2:
            for id_atom, (x, y) in enumerate(coords):
                id_mol, id_species = 1, el_dict[elements[id_atom]]  # elList[id_atom]
                # print id_atom + 1, id_mol, id_species + 1, x, y, z
                atoms += format_str.format(id_atom + 1, id_mol, id_species + 1, x, y, 0.) + '\n'
        else:
            raise ValueError("dimension 1 not yet implemented")

        bonds_str = 'Bonds \n\n'
        for i_bond, (i_a, i_b, b_type) in enumerate(bonds):
            bonds_str += '{0:d} {1:d} {2:d} {3:d}'.format(i_bond + 1, b_type, i_a, i_b) + '\n'

        return atomtypes + '\n' + cell_dimesions + '\n' + masses + '\n' + atoms + '\n' + bonds_str + '\n'

    def structure_full(self):
        """
        Write routine to create atom structure static file for atom_type='full' that can be loaded by LAMMPS

        Returns:

        """
        coords = self.rotate_positions(self._structure)

        # extract electric charges from potential file
        q_dict = {}
        for el in self._structure.get_species_symbols():
            q_dict[el] = float(self.potential.get("set group {} charge".format(el)))

        species_translate_list = list()
        sorted_species_list = self._structure.get_species_symbols()
        for el in self._structure.species:
            ind = np.argwhere(sorted_species_list == el.Abbreviation).flatten()[-1]
            species_translate_list.append(ind)

        # analyze structure to get molecule_ids, bonds, angles etc
        molecule_lst, bonds_lst, angles_lst = [], [], []

        # species_lst = structure.get_species_objects()
        # for id_el, el in enumerate(structure.species):
        #     el.id = id
        # el_lst = structure.get_chemical_elements()

        num_atoms_in_molecule = 3
        neighbors = self._structure.get_neighbors(num_neighbors=num_atoms_in_molecule + 2)
        # print "neighbors: ", neighbors.distances
        id_mol = 0
        indices = self._structure.indices
        for id_el, id_species in enumerate(indices):
            el = self._structure.species[id_species]
            # print "id: ", id, el.Abbreviation, neighbors.indices[id][0:2]
            if el.Abbreviation in ["O"]:
                # print "id_mol: ", id_mol
                id_mol += 1
                molecule_lst.append([id_el, id_mol, id_species])
                # Just to ensure that the attached atoms are indeed H atoms
                # id_n1, id_n2 = np.intersect1d(neighbors.indices[id_el], self._structure.select_index("H"))[0:2]
                id_n1, id_n2 = neighbors.indices[id_el][0:2]
                # print "id: ", id, id_n1, len(el_lst), el_lst[1].id
                molecule_lst.append([id_n1, id_mol, species_translate_list[indices[id_n1]]])
                molecule_lst.append([id_n2, id_mol, species_translate_list[indices[id_n2]]])

                bonds_lst.append([id_el + 1, id_n1 + 1])
                bonds_lst.append([id_el + 1, id_n2 + 1])

                angles_lst.append([id_n1 + 1, id_el + 1, id_n2 + 1])
            elif el.Abbreviation not in ["H"]:  # non-bonded ions
                id_mol += 1
                molecule_lst.append([id_el, id_mol, id_species])

        m_lst = np.array(molecule_lst)
        molecule_lst = m_lst[m_lst[:, 0].argsort()]
        # print "m_lst: ", m_lst
        # print "mol: ", molecule_lst

        atomtypes = ' Start File for LAMMPS \n' + \
                    '{0:d} atoms'.format(len(self._structure)) + ' \n' + \
                    '{0:d} bonds'.format(len(bonds_lst)) + ' \n' + \
                    '{0:d} angles'.format(len(angles_lst)) + ' \n' + \
                    '{0} atom types'.format(self._structure.get_number_of_species()) + ' \n' + \
                    '{0} bond types'.format(1) + ' \n' + \
                    '{0} angle types'.format(1) + ' \n'

        cell_dimensions = self.simulation_cell()

        masses = 'Masses' + '\n\n'
        el_obj_list = self._structure.get_species_objects()
        for object_id, el in enumerate(el_obj_list):
            masses += '{0:3d} {1:f}'.format(object_id + 1, el.AtomicMass) + '\n'

        atoms = 'Atoms \n\n'

        # format: atom-ID, molecule-ID, atom_type, q, x, y, z
        format_str = '{0:d} {1:d} {2:d} {3:f} {4:f} {5:f} {6:f}'
        for atom in molecule_lst:
            id_atom, id_mol, id_species = atom
            # print id_atom, id_mol, id_species
            x, y, z = coords[id_atom]
            el_id = self._structure.species[id_species].Abbreviation
            atoms += format_str.format(id_atom + 1, id_mol, id_species + 1, q_dict[el_id], x, y, z) + '\n'

        if len(bonds_lst) > 0:
            bonds_str = 'Bonds \n\n'
            for i_bond, id_vec in enumerate(bonds_lst):
                bonds_str += '{0:d} {1:d} {2:d} {3:d}'.format(i_bond + 1, 1, id_vec[0], id_vec[1]) + '\n'
        else:
            bonds_str = "\n"

        if len(angles_lst) > 0:
            angles_str = 'Angles \n\n'
            for i_angle, id_vec in enumerate(angles_lst):
                # print "id: ", i_angle, id_vec
                angles_str += '{0:d} {1:d} {2:d} {3:d} {4:d}'.format(i_angle + 1, 1, id_vec[0], id_vec[1], id_vec[2]) \
                              + '\n'
        else:
            angles_str = "\n"
        return atomtypes + '\n' + cell_dimensions + '\n' + masses + '\n' + atoms + '\n' \
               + bonds_str + '\n' + angles_str + '\n'

    def structure_charge(self):
        """
        Create atom structure including the atom charges.

        By convention the LAMMPS atom type numbers are chose alphabetically for the chemical species.

        Returns: LAMMPS readable structure.

        """
        atomtypes = 'Start File for LAMMPS \n' + \
                    '{0:d} atoms'.format(self._structure.get_number_of_atoms()) + ' \n' + \
                    '{0} atom types'.format(self._structure.get_number_of_species()) + ' \n'

        cell_dimesions = self.simulation_cell()

        masses = 'Masses\n\n'

        for ind, obj in enumerate(self._structure.get_species_objects()):
            masses += '{0:3d} {1:f}'.format(ind+1, obj.AtomicMass) + '\n'

        atoms = 'Atoms\n\n'

        coords = self.rotate_positions(self._structure)

        el_charge_lst = self._structure.charge
        el_lst = self._structure.get_chemical_symbols()
        el_alphabet_dict = {}
        for ind,el in enumerate(self._structure.get_species_symbols()):
            el_alphabet_dict[el] = ind+1
        for id_atom, (el, coord) in enumerate(zip(el_lst, coords)):
            id_el = el_alphabet_dict[el]
            dim = self._structure.dimension
            c = np.zeros(3)
            c[:dim] = coord
            atoms += '{0:d} {1:d} {2:f} {3:.15f} {4:.15f} {5:.15f}'.format(
                id_atom + 1, id_el, el_charge_lst[id_atom], c[0], c[1], c[2]) + '\n'
        return atomtypes + '\n' + cell_dimesions + '\n' + masses + '\n' + atoms + '\n'

    def structure_atomic(self):
        """
        Write routine to create atom structure static file that can be loaded by LAMMPS

        Returns:

        """
        atomtypes = 'Start File for LAMMPS \n' + \
                    '{0:d} atoms'.format(len(self._structure)) + ' \n' + \
                    '{0} atom types'.format(len(
                        self._el_eam_lst)) + ' \n'  # '{0} atom types'.format(structure.get_number_of_species()) + ' \n'

        cell_dimesions = self.simulation_cell()

        masses = 'Masses\n\n'

        el_struct_lst = self._structure.get_species_symbols()
        el_obj_lst = self._structure.get_species_objects()
        # el_struct_lst = structure.get_chemical_symbols()
        # el_obj_lst = structure.get_chemical_elements()
        el_dict = {}
        for id_eam, el_eam in enumerate(self._el_eam_lst):
            if el_eam in el_struct_lst:
                id_el = list(el_struct_lst).index(el_eam)
                el = el_obj_lst[id_el]
                el_dict[el] = id_eam + 1
                masses += '{0:3d} {1:f}'.format(id_eam + 1, el.AtomicMass) + '\n'
            else:
                # element in EAM file but not used in structure, use dummy for atomic mass
                masses += '{0:3d} {1:f}'.format(id_eam + 1, 1.00) + '\n'

        atoms = 'Atoms\n\n'

        coords = self.rotate_positions(self._structure)

        el_lst = self._structure.get_chemical_elements()
        for id_atom, (el, coord) in enumerate(zip(el_lst, coords)):
            id_el = el_dict[el]
            dim = self._structure.dimension
            c = np.zeros(3)
            c[:dim] = coord
            atoms += '{0:d} {1:d} {2:.15f} {3:.15f} {4:.15f}'.format(id_atom + 1, id_el, c[0], c[1], c[2]) + '\n'
        return atomtypes + '\n' + cell_dimesions + '\n' + masses + '\n' + atoms + '\n'

    def rotate_positions(self, structure):
        """
        Rotate all atomic positions in given structure according to new Prism cell

        Args:
            structure: Atoms-like object. Should has .positions attribute

        Returns:
            (list): List of rotated coordinates
        """
        prism = UnfoldingPrism(self._structure.cell)
        coords = [prism.pos_to_lammps(position) for position in structure.positions]
        return coords

def write_lammps_datafile(structure, file_name='lammps.data', cwd=None):
    lammps_str = LammpsStructure()
    lammps_str.el_eam_lst = structure.get_species_symbols()
    lammps_str.structure = structure
    lammps_str.write_file(file_name=file_name, cwd=cwd)


# Copyright [2017] [Ronald Fowler]
#
# 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.

""" Utilities for beam hardening correction method. Classes are provided
to load the various data types that are needed by the algorithm.
classes:

    specData  object to load and store X-ray spectrum for specified case
    carousel  object to load and store description of test carousel
    carouselCalibrationData  object to load and store the calibration data
    fitData  object with methods and data related to the fitting process

"""
# just for reading 16 bit images and conversion to log(I0/I)
from __future__ import division
from __future__ import print_function

import sys
import os
import logging
import pdb
import pkg_resources
try:
    import numpy as np
except ImportError:
    sys.exit("Error: cant find numpy")
try:
    import matplotlib.pyplot as plt
except ImportError:
    sys.exit("Error: cant find matplotlib")

def resourcepath(filename):
    if os.path.isfile(filename):
        return filename
    try:
        pkgfilename = pkg_resources.resource_filename(__name__,'/data/'+filename)
        return pkgfilename
    except:
        pass
    return filename

class specData(object):
    """
    Spectral data for given condition

    Data referes to the initial values of target, voltage and angle.
    This object contains the spectral data for the given
    target material, applied voltage and take off angle.
    Currently this is only available for Tungsten in the range
    1-150KeV from data obtained from Spekcalc at fixed set of values.
    Interpolation may be used for intermeidate points.
    The resolution is as set in the data files.
    Note that only one spectrum is loaded as we do not vary
    the voltage in this fitting procedure.

    Attributes

    target : string
        X-ray tube target material, chemical symbol e.g. W
    voltage : float
        applied voltage to the X-ray tube in KeV
    angle : float
        take off angle in degrees for the X-ray beam. This is kept constant

    Methods

    getE()
        Get array of energies over which the spectrum is defined (KeV)
    getS()
        Get array of spectral intensities corresponding to energies of getE()
    getEnergyRes()
        Return the step size between spectral points, assumed constant.
    getMaxEnergy()
        Return the applied voltage (KeV)
    isValid()
        Flag to indicate if object correctly set up.
    """

    def __init__(self, target, voltage, angle):
        """ set up spectrum from file for given state """
        self.target = target
        self.voltage = voltage
        self.angle = angle
        # initialise these for pylint
        self.en = None
        self.amp = None
        self.valid = False
        self.readData(target, voltage, angle)

    def readData(self, target, voltage, angle):
        """ Read the spectra for this voltage/angle pair from file.
            If file not found, return spec=0 """
        filename = resourcepath("spectra/%s/%03d%03d.spc" % (target, voltage, angle*10))
        if not os.path.isfile(filename):
            print("File not found: ", filename, " in spectra")
            print("Only fitting with function spectra possible")
            # just set energy array 0 to voltage, in 0.5KeV steps
            estep=0.5
            self.en = np.arange(voltage/estep+estep*0.5)*estep
        else:
            with open(filename, 'r') as fl:
                self.en, self.amp = np.loadtxt(fl, unpack=True)
        self.valid = True

    def getE(self):
        """ return energy array; should remove """
        return self.en

    def getS(self):
        """ return intensity of spectra at energy points"""
        return self.amp

    def getEnergyRes(self):
        """ return the (assumed constant) energy resolution"""
        return self.en[1]-self.en[0]

    def getMaxEnergy(self):
        """ voltage is max energy"""
        return self.voltage

    def isValid(self):
        """ has object been set up OK"""
        return self.valid

class materialAtt(object):
    """Attenuation for a material expressed as formula

    For each material of interest will create this object to represent attenuation
    as a function of energy. Material is idenitified by chemical symbol and data is
    read from file produced by xcom.f.
    """

    def __init__(self, formula, density):
        self.name = formula
        # for pylint
        self.mu = None
        self.energy = None
        self.valid = False
        if len(formula) != 0:
            self.readFile(formula, density)

    def readFile(self, formula, density):
        """ load the mu data for this material from a simple ascii file in ./xcom
             mu is mulitplied by density """
        filename = resourcepath("xcom/%s.txt" % (formula))
        if os.path.isfile(filename):
            with open(filename, 'r') as fl:
                self.energy, self.mu = np.loadtxt(fl, unpack=True)
                self.mu = self.mu*density
            self.valid = True
        else:
            print("failed to find attenuation file: ", filename)
            self.valid = False

    def getE(self):
        """ return array of energy data"""
        return self.energy

    def getMu(self):
        """ attenuation array"""
        return self.mu

    def getEnergyRes(self):
        """ energy resolution, assumed constant"""
        return self.energy[1]-self.energy[0]

    def getMaxEnergy(self):
        """ highest energy value recorded"""
        return self.energy[-1]

    def getMuByE(self,energyVal):
        """ return attenuation for the given energy """
        if energyVal<0:
            print("error: bad value in getMuByE")
            return -1.
        for i in range(len(self.energy)):
            if energyVal <= self.energy[i]:
                if energyVal == self.energy[i]:
                    return self.mu[i]
                else:
                    frac = (energyVal-self.energy[i-1])/(self.energy[i]-self.energy[i-1])
                    muVal = self.mu[i-1]+(self.mu[i]-self.mu[i-1])*frac
                    return muVal
        print("error: failed to match energy in getMuByE")
        return -1.

    def isValid(self):
        """ if object ok"""
        return self.valid

class carousel(object):
    """ class for data describing the test carousel """

    def __init__(self, defFile):
        self.defFile = resourcepath(defFile)
        self.info = None
        self.density = None
        self.numSamples = None
        self.sampWidth = None
        self.materialTypes = None
        self.mask = None
        self.valid = False
        self.filterAtt = {}
        self.__readFile(self.defFile)

    def __readFile(self, defFile):
        """ read a simple ascii file describing carousel from file
        """
        if os.path.isfile(defFile):
            with open(defFile, 'r') as fl:
                self.info = fl.readline()
                self.numSamples = int(fl.readline())
                self.materialTypes = []
                mline = fl.readline()
                self.materialTypes = mline.split(',')
                self.density = np.fromstring(fl.readline(), dtype = float, sep=',')
                self.sampWidth = np.fromstring(fl.readline(), dtype = float,
                                               sep=',')
                self.filterAtt = {}
                # assume last sample is labeled "Nothing"
                for i in range(self.numSamples - 1):
                    try:
                        self.filterAtt[i] = materialAtt(self.materialTypes[i],self.density[i])
                    except:
                        print("** failed to set carousel attenuation for ",self.materialTypes[i])
                self.mask = np.zeros((self.numSamples),dtype=bool)
            self.valid = True
        else:
            print("failed to find carousel file: ", defFile)
            self.valid = False

    def getSamples(self):
        """ elements in the carousel """
        return self.numSamples

    def isValid(self):
        """ is object ok"""
        return self.valid

class carouselCalibrationData(object):
    """ This class reads the calibration data from a carousel calibration file.

    The calibration file contains information about the X-ray voltage, take-off angle,
    target material and image resolution.
    It also gives the name & possibly the format of the image data file for the carousel.
    """
    def __init__(self, calFile, carouselInfo):
        self.calFile = resourcepath(calFile)
        self.samples = carouselInfo.getSamples()-1 # note that last sample "Nothing" has no image
        self.voltage = None
        self.targeMat = None
        self.rows = 0
        self.lines = 0
        self.angle = 0
        self.filterDensity = 0
        self.valid = False

        self.filterMaterial = {}
        self.filterWidth = {}
        self.filterDensity = {}
        self.filterAtten = {}

        self.detectorMaterial = ''
        self.detectorWidth = 0.
        self.detectorDensity = 0.
        self.detectorAtten = {}

        self.targetMat = ''
        self.targetDensity = 0.
        self.targetAtten = {}

        self.info = ''
        self.imageFileFormat = ''
        self.imageFile = ''
        self.image = {}
        self.spec = {}
        self.filterCount = 0
        self.__centre = None

        self.__readCalFile(self.calFile)
        if self.valid:
            try:
                self.whiteLevel = 0
                self.__readImageFile(self.imageFile)
                self.__setAverages()
                self.width = 100
                self.__cacheAveSet = False
                self.__cacheAve = np.zeros(shape=(self.samples, self.lines))
                logging.debug('initialised cacheAve')
            except:
                self.valid = False
                logging.debug('reading data failed')

    def __readCalFile(self, calFile):
        """ read calibration data file, and from that the actual image data"""
        if os.path.isfile(calFile):
            with open(calFile, 'r') as fl:
                try:
                    self.info = self.__readLineStrip(fl)
                    self.voltage = float(self.__readLineStrip(fl))
                    self.angle = float(self.__readLineStrip(fl))
                    self.targetMat = self.__readLineStrip(fl).rstrip()
                    self.targetDensity = float(self.__readLineStrip(fl).rstrip())
                    if self.targetMat != "W":
                        print("Warning: only W (Tungsten) target supported at present: not '", self.targetMat, "'")
                    try:
                        self.targetAtten = materialAtt(self.targetMat,self.targetDensity)
                    except:
                        print("** failed to set target attenuation for ",self.targetMat)
                    try:
                        self.spec = specData(self.targetMat,self.voltage,self.angle)
                    except:
                        print("** failed to load spectra for ",self.targetMat,self.voltage,self.angle)
                    self.rows = int(self.__readLineStrip(fl))
                    self.lines = int(self.__readLineStrip(fl))
                    self.imageFile = self.__readLineStrip(fl)
                    self.imageFileFormat = self.__readLineStrip(fl)
                    self.filterCount = int(self.__readLineStrip(fl))
                    for i in range(self.filterCount):
                        self.filterMaterial[i] = self.__readLineStrip(fl)
                        self.filterWidth[i] = float(self.__readLineStrip(fl))
                        self.filterDensity[i] = float(self.__readLineStrip(fl))
                        try:
                            self.filterAtten[i] = materialAtt(self.filterMaterial[i],self.filterDensity[i])
                        except:
                            print("** failed to set attenuation for ",self.filterMaterial[i])

                    self.detectorMaterial = self.__readLineStrip(fl)
                    self.detectorWidth = float(self.__readLineStrip(fl))
                    self.detectorDensity = float(self.__readLineStrip(fl))
                    try:
                        self.detectorAtten = materialAtt(self.detectorMaterial,self.detectorDensity)
                    except:
                        print("** failed to set detector attenuation for ",self.detectorMaterial)
                    self.valid = True
                except (ValueError, IOError):
                    print("Read Calibration file failed")
                    self.valid = False
        else:
            print("failed to find calibration file: ", calFile)
            self.valid = False

    def __readLineStrip(self, fl):
        """ read line: strip newline and anything beyond #, if present
        """
        strng = fl.readline()
        if strng.find("#") == -1:
            return strng.rstrip()
        else:
            return strng[:strng.find("#")].rstrip()

    def __readImageFile(self, imageFile):
        """ read the image data based on specificed format """
        if os.path.isfile(imageFile):
            if self.imageFileFormat == "uint16":
                # if raw uint16 data, assume first image is flat field and normalise
                # rest of images by this and take log(I0/I)
                nimages = self.samples+1
                with open(imageFile, 'rb') as fl:
                    tmpimage = np.fromfile(fl, dtype = self.imageFileFormat,
                                           count = self.rows*self.lines*nimages)
                    tmpimage = tmpimage.reshape(nimages, self.lines, self.rows)
                self.image = np.zeros(self.rows*self.lines*self.samples,
                                      dtype=float).reshape(self.samples, self.lines, self.rows)
                # note - imported division from _future_ to avoid int div

                # assume that the first image is the white level, I0, a constant
                # over the image. To impose this assumption we take the average
                # value over the first (flat field, shading corrected) image.
                # This is only set for uint16; it is not known for float32
                whiteLev = np.average(tmpimage[0,:,:])
                self.whiteLevel = whiteLev
                for i in range(nimages-1):
                    # catch zero data:
                    if np.min(tmpimage[i+1,:,:])==0:
                        imt = tmpimage[i+1,:,:]
                        imt[imt==0] = np.mean(imt)
                        logging.warning('zero data replaced in image %d',i+1)
                    self.image[i,:,:] = np.log( whiteLev / tmpimage[i+1,:,:] )

            elif self.imageFileFormat == "uint16_65535":
                # if raw uint16_65535 data, whitelevel set as 65535
                # Question: Should we allow other values of whitelevel?
                nimages = self.samples
                with open(imageFile, 'rb') as fl:
                    tmpimage = np.fromfile(fl, dtype = "uint16",
                                           count = self.rows*self.lines*nimages)
                    tmpimage = tmpimage.reshape(nimages, self.lines, self.rows)
                self.image = np.zeros(self.rows*self.lines*self.samples,
                                      dtype=float).reshape(self.samples, self.lines, self.rows)
                whiteLev = 65535
                self.whiteLevel = whiteLev
                for i in range(nimages):
                    # catch zero data:
                    if np.min(tmpimage[i,:,:])==0:
                        imt = tmpimage[i,:,:]
                        imt[imt==0] = np.mean(imt) + 1 # in case mean < 1.0 (imt is uint16)
                        logging.warning('zero data replaced in image %d',i)
                    self.image[i,:,:] = np.log( whiteLev / tmpimage[i,:,:] )

            elif self.imageFileFormat == "float32":
                # assume float data already transformed by I0 and log
                with open(imageFile, 'rb') as fl:
                    try:
                        self.image = np.fromfile(fl, dtype = self.imageFileFormat,
                                 count = self.rows*self.lines*self.samples).reshape(self.samples, self.lines, self.rows)
                    except:
                        print("Failed in reading/converting image file: check data")
            else:
                print("** error: image format name ",self.imageFileFormat," not recognised")

        else:
            print("Image file not found!: ", imageFile)

    def printImageStats(self, carInf):
        """ print out some data for each frame in the set of images"""
        max1 = int(self.voltage*2) # number of 0.5 KeV steps
        min10 = int(max1/10)
        for i in range(self.samples):
            nancount = np.count_nonzero(np.isnan(self.image[i,:,:]))
            if nancount>0:
                print("*** img ", i, " contains: ", nancount, " NaNs (", nancount*100./(self.rows*self.lines), "%)")
                maskedimage = np.ma.array(self.image[i,:,:],mask = np.isnan(self.image[i,:,:]))
                ave = np.ma.average(maskedimage)
                print("    average(masked)= ", ave, "  max= ", np.ma.max(maskedimage))
            else:
                ave = np.average(self.image[i,:,:])
                print("img ", i, " ", carInf.materialTypes[i]," ", carInf.sampWidth[i], " average= ",ave)
            minmu = np.min(carInf.filterAtt[i].getMu()[min10:max1])
            maxmu = np.max(carInf.filterAtt[i].getMu()[min10:max1])
            maxatt = maxmu*carInf.sampWidth[i]
            minatt = minmu*carInf.sampWidth[i]
            print("    minatt = ",minatt,"  maxatt = ",maxatt)
            if ave<minatt or ave>maxatt:
                print("    *** Warning: Attenuation of sample outside expected bounds!")

    def __setAverages(self):
        """ pre compute mean centre of each image row
        """
        self.__centre = np.zeros(self.samples*self.lines).reshape(self.samples,
                                 self.lines)
        for samp in range(self.samples):
            for li in range(self.lines):
                self.__centre[samp, li] = self.__getCentrePos(li, samp)

    def __getCentrePos(self, line, sample):
        """ do a weighted average of position by signal to find the centre of
            the signal, in pixels"""
        vals = self.image[sample,line,:]
        meanpt = np.sum(np.arange(len(vals))*vals)/np.sum(vals)
        if meanpt>0.75*np.size(vals) or meanpt<0.25*np.size(vals):
            meanpt = np.size(vals)*0.5
            logging.info('reset getCentrePos to %f for sample %d',meanpt,sample)
        return meanpt

    def getCentrePos(self, line, sample):
        """ get centre point"""
        return int(self.__centre[sample, line])

    def getLines(self):
        """ return the number of lines of the data"""
        return self.lines

    def getRows(self):
        """ return the number of rows of the data"""
        return self.rows

    def getImage(self, imgNum):
        """ access a given image by number"""
        if imgNum >= 0 and imgNum <self.samples:
            return self.image[imgNum,:,:]
        else:
            return 0

    def getAvAtten(self, line, sample):
        """ average over a named range of rows """
        if self.__cacheAveSet:
            return self.__cacheAve[sample, line]
        logging.debug('calc cache values of Ave')
        self.__cacheAveSet =True
        for s in range(self.samples):
            for l in range(self.lines):
                logging.debug('s= %d l=%d c=%f', s, l, self.__centre[s, l])
                rowStart = int(self.__centre[s, l]-self.width)
                rowEnd = int(rowStart+2*self.width)
                self.__cacheAve[s, l] = np.average(self.image[s, l, rowStart:rowEnd])
        return self.__cacheAve[sample, line]

    def setWidthAve(self, width):
        """ set the (half) width to be used when calculating the average
            attenuation along a row
            This is measured either size of the centre point. All points should
            lie on the sample image.
        """
        self.width = width
        self.__cacheAveSet = False

    def isValid(self):
        """ is object ok"""
        return self.valid

    def plotCalData(self, showplt, markWidth):
        """ plot calibration images in one window using matplotlib. Only
            return when window closed. Mark image with width of "averaging",
            markWidth """
        nsample = self.samples
        if nsample<5:
            pltgrid=2
        elif nsample<10:
            pltgrid=3
        else:
            pltgrid=4
            nsample = min(nsample, 15)
        for isam in range(nsample):
            # deep copy image so can modify
            z = np.copy(self.getImage(isam))
            zzmax = 0
            # mark width of each line used for average
            for li in range(self.lines):
                cen = self.getCentrePos(li, isam)
                if markWidth>0:
                    z[li, int(round(cen-markWidth))] = 0
                    z[li, int(round(cen+markWidth))] = 0
            maskedimage = np.ma.array(z, mask = np.isnan(z))
            zmax = np.max(maskedimage)
            zzmax = max(zmax, zzmax)
            plt.subplot(pltgrid, pltgrid, isam+1, facecolor='y')
            plt.imshow(maskedimage, cmap='RdBu', vmin = 0, vmax = zzmax,
                       aspect='auto')
            plt.draw()
        if showplt:
            plt.show()

    def plotCalProf(self):
        """ will be used to plot profile along a row; not yet complete"""
        nsample = self.samples
        if self.samples < 5:
            pltgrid = 2
        elif self.samples < 10:
            pltgrid = 3
        else:
            pltgrid = 4
            nsample = max(self.samples)
        for isam in range(nsample):
            plt.subplot(pltgrid, pltgrid, isam+1, facecolor = 'y')

class fitData(object):
    """ This object contains fit related data and functions
    """

    def __init__(self, carInfo, carCal, defMat):
        if not ( carInfo.isValid and carCal.isValid ):
            self.isValid = False
            return
        self.carInfo = carInfo
        self.carCal = carCal
        #
        # set default fitting parameters; values are order of polynomials
        # in line number. Hence "0" is a constant, not dependent on line number
        # default is for 3rd polys for detector and target width, with global
        # width for Cu filters. Also Gaussian for spectra, if no pre-defined values.
        #
        self.vary_target = 1
        self.vary_detector = 1
        self.vary_filter = 0
        self.vary_energy = -1
        self.vary_epk = -1
        self.vary_ewidlow = -1
        self.vary_ewidhigh = -1
        # make space for default values to use if vary==-1
        self.defaults = np.zeros(4)
        #
        # since there may be several filters, define which one should vary
        self.vary_filter_name="Cu"
        self.verbose = False
        # define the fitting parameters and their global/local status
        self.varFilter = -1
        # check if we have filter of material defMat
        for i in carCal.filterMaterial:
            if defMat == carCal.filterMaterial[i]:
                self.varFilter = i
                self.defFilterMat = defMat
        # if defMat not found, use first material in list, if any
        if self.varFilter==-1:
            print("default filter: ", defMat," not found")
            if len(carCal.filterMaterial)>0:
                self.defFilterMat=carCal.filterMaterial[0]
                self.varFilter = 0
                print("using material = ",self.defFilterMat," for fitting")
            else:
                print("no filter material present, hence cannot fit this")
        self.atten = np.zeros([self.carCal.lines,self.carInfo.getSamples()])
        self.objFnCalls = 0
        self.nlines = 0
        self.lineStep = 1
        self.linestep = 1
        self.bounds = False
        self.boundsValues = {}
        self.solver = "old"

    def calcWidths(self,x0,nlines,xe):
        """ Function to return the 3 widths for the target,
            the detector and the filter from the set of fitting
            variables. Each is assumed to be a polynomial of some
            order in the line number. Global varables are zero
            order polynomials.
            Also returns the energy array which can be a polynomial:
            E+aE**2+... ; this should be constrained >=0, not done at
            present."""
        lines=np.array(range(nlines),dtype="double")
        nt = self.vary_target+1
        nd = self.vary_detector+1
        nf = self.vary_filter+1
        ne = self.vary_energy+1
        ns = self.vary_epk+self.vary_ewidlow+self.vary_ewidhigh+3
        if len(x0)<nt+nd+nf+ne+ns:
            print("** calcWidthd called with too few values in x0")
            sys.exit(1)
        # Polynomial expressions: highest order term is first in the array.
        if nt>0:
            twidth = np.polyval(x0[:nt],lines)
        else:
            twidth = np.polyval(self.defaults[0:1],lines)
        #twidth=np.polyval(x0[nt-1:0:-1],lines)
        if nd>0:
            dwidth = np.polyval(x0[nt:nt+nd],lines)
        else:
            dwidth = np.polyval(self.defaults[1:2],lines)
        #dwidth=np.polyval(x0[nt+nd-1:nt:-1],lines)
        dwidth = np.exp( dwidth ) # force >0 by working in log space
        if nf>0:
            fwidth = np.polyval(x0[nt+nd:nt+nd+nf],lines)
        else:
            fwidth = np.polyval(self.defaults[2:3],lines)
        if ne>0:
            # This term should be constrained as >=0 for all xe but is not at present.
            # -Ve values will give errors in output stage.
            ecoeffs = xe + xe*xe*(np.polyval(x0[nt+nd+nf:nt+nd+nf+ne],xe))
        else:
            ecoeffs = xe
        if ns>0:
            i0 = nt+nd+nf+ne
            indarr = xe>x0[i0:i0+1]
            spectra = xe-x0[i0:i0+1]
            spectra[indarr] = spectra[indarr]*x0[i0+1:i0+2]
            spectra[~indarr] = spectra[~indarr]*x0[i0+2:i0+3]
            spectra = np.exp(-spectra**2)
            # mask out lowest 10% of spectra, as in spekCalc
            masklen = int(len(spectra)*0.1)
            spectra[0:masklen] = 0.
        else:
            spectra = 0.
        return twidth,dwidth,fwidth,ecoeffs,spectra

    def dofit(self,nlines,lstep,xin):
        """ perform fit """
        got=0
        try:
            # from scipy.optimize import minimize
            from scipy.optimize import leastsq
            got=1
            from scipy.optimize import least_squares
        except:
            # if found old lib try and continue
            if got==0:
                print("** cannot find scipy leastsq or least_squares - check python has scipy")
                return

        if self.verbose:
            pdb.set_trace()
        x = xin
        self.nlines = nlines
        self.lineStep = lstep

        # use either old or new solver interface from scipy for least squares
        if self.solver=="old":
            res = leastsq(self.objFunSq, x, full_output = True)
        else:
            if self.bounds:
                resobj = least_squares(self.objFunSq, x, verbose = 1, bounds=self.boundsValues)
            else:
                resobj = least_squares(self.objFunSq, x, verbose = 1, method='lm')
            infodict = {"nfev":resobj.nfev}
            cov = [0]
            res = (resobj.x,cov,infodict,resobj.message,resobj.status)

        print("Line 0 atten=",self.atten[0,:])
        expt = np.zeros(self.carCal.samples+1)
        for i in range(self.carCal.samples):
            expt[i] = self.carCal.getAvAtten(0,i)
        print("Line 0 expt=",expt)
        # Do final calulation on all lines:
        self.lineStep = 1
        if self.solver=="old":
            self.objFunSq(res[0])
        else:
            self.objFunSq(resobj.x)
        #
        return res

    def objFunSq(self,x):
        """ The function to minimize; returns the squared error for every point on each
            selected line  """
        # Get the 3 widths: target(e.g. W), detector(e.g. CsI), global filter(e.g. Cu)
        # target and detector widths depend on line number, filter is a global value
        # for flexiblity all 3 are dimesioned by nlines
        #
        # mask out low en spectral points which may get undue weight if -ve filter widths occur
        minpt = int(0.1*len(self.carCal.spec.getE()))
        #
        xe = self.carCal.spec.getE()
        tw,dw,fw,ec,spectra = self.calcWidths(x,self.nlines,xe)
        nsamples = self.carInfo.numSamples - 1 # ignore null sample
        ans = np.zeros(nsamples*self.nlines)
        tarAtt = self.carCal.targetAtten
        if isinstance(spectra,np.ndarray):
            se = spectra
        else:
            se = self.carCal.spec.getS()
        #
        for line in range(0,self.nlines,self.lineStep):
            # compute filter attenuation and hence i0 as signal level with no sample for this line
            attDet = dw[line]*self.carCal.detectorAtten.getMu()[:len(xe)]
            attSum = np.zeros(len(xe))
            for filt in range(self.carCal.filterCount):
                if filt == self.varFilter:
                    fwid = fw[line]
                else:
                    fwid = self.carCal.filterWidth[filt]
                attSum = attSum + fwid*self.carCal.filterAtten[filt].getMu()[:len(xe)]
            # this is the key integral done as a simple sum. Can ignore width of each value
            # as constant energy steps, so cancels in I/I0
            at_se = se*np.exp(-attSum-tarAtt.getMu()[:len(xe)]*tw[line])*ec*(1-np.exp(-attDet))
            # drop low energy terms below 10%
            at_se_t = at_se[minpt:]
            # remove nan's - why are nan's present? exp overflow gives inf, multiply by 0 gives nan
            # in most cases nans are OK to ignore.
            at_se_finite = at_se_t[np.logical_not(np.isnan(at_se_t))]
            i0 = np.sum(at_se_finite)
            #
            for sample in range(nsamples):
                # skip masked samples
                if self.carInfo.mask[sample]:
                    continue
                widSam = self.carInfo.sampWidth[sample]
                #
                attSam = widSam*self.carInfo.filterAtt[sample].getMu()[:len(xe)]
                #
                # drop low energy terms < 10%
                at_se_sample = (at_se * np.exp(-attSam))[minpt:]
                #
                # remove nan's - see above
                at_se_sample_finite = at_se_sample[np.logical_not(np.isnan(at_se_sample))]
                i_sample = np.sum(at_se_sample_finite)
                if i0==0. and self.verbose:
                    print("warn: i0 zero at ",line)
                    i0=1.
                if i_sample == 0 and self.verbose:
                    print("i_sample=0")
                if i_sample < 0. and self.verbose:
                    print("i_sample<0",at_se_sample[:8])
                #sumSq = sumSq + ( (i_sample/i0) - self.carCal.getAvAtten(line,sample) ) ** 2
                ans[line*nsamples+sample] = ( np.log(i0/i_sample) - self.carCal.getAvAtten(line,sample) ) ** 2
                self.atten[line,sample] = np.log(i0/i_sample)
        if self.verbose:
            print("tw,dw,fw,sumSq: ",tw[0],dw[0],fw[0],np.sum(ans))
            # for debugging we print the normalised detector response for the current parameters
            plotFreq=10
            if self.objFnCalls%plotFreq == 0:
                plt.figure('NormRespone')
                plt.xlabel('energy')
                plt.ylabel('Normalised response')
                # map nans to zero for plotting
                at_se[np.isnan(at_se)] = 0.
                plt.plot(xe,at_se/i0)
                plt.draw()
                plt.show(block=False)
        self.objFnCalls=self.objFnCalls+1

        #
        # return vector of squared errors: length=samples*lines
        return ans

    def linesPolyFit(self,soln,corMat,corEn,npoints,attrange):
        """ Function to calculate the attenuation over "npoints" for attenuation up to "attrange" for each line.
            Uses the fitted parameters for attenuation in "soln". Having calculated apparent attenuation for the correction
            material "corMat", map the observed attenuation to the true attenuation at the corEn energy. Then fit a
            polynomial to the data and save the coefficients.
            """
        # Get the 3 widths: target(e.g. W), detector(e.g. CsI), global filter(e.g. Cu)
        # target and detector widths depend on line number, filter is a global value
        # for flexiblity all 3 are dimesioned by nlines. Also the energy parameter, if fitted.
        xe = self.carCal.spec.getE()
        tw,dw,fw,ec,spectra = self.calcWidths(soln,self.nlines,xe)
        #nsamples = self.carInfo.numSamples
        # allocate space to store all calculated points and polynomials fitted to them
        attout = np.zeros(shape=(self.nlines,npoints+1))
        # set order of polynomial fits to use
        # for xtek a 4th order polynomial can be used by the reconstruction - note
        # that as constant term is forced to zero, 3 gives 4th order
        odpoly = 8
        xtekodpoly = 3
        # determine if the solution varies with line number; if not only one fit required
        vary_line = (self.vary_target>0 or self.vary_detector>0 or self.vary_filter>0)
        #
        # find the actual attenuation of the correction material at the correction energy
        corrAtt = corMat.getMuByE(corEn)

        tarAtt = self.carCal.targetAtten
        #
        if isinstance(spectra,np.ndarray):
            se = spectra
        else:
            se = self.carCal.spec.getS()
        #

        # generate points to evaluate attenuation at.
        mulist = np.arange(npoints+1,dtype='float')*attrange/(npoints*corrAtt)
        # attin is the observed attenuation; for each line want to find corresponding attout
        # the "true" attenuation at energy corEn
        attin = mulist*corrAtt
        #
        # for each line generate npoints values of attenuation from fit data
        #
        if vary_line:
            nlines = self.nlines
        else:
            nlines = 1
        #
        polyfit = np.zeros(shape=(nlines,odpoly+2))
        xpolyfit = np.zeros(shape=(nlines,xtekodpoly+2))
        #
        for line in range(nlines):
            # compute filter attenuation and hence i0 as signal level with no sample for this line
            attDet = dw[line]*self.carCal.detectorAtten.getMu()[:len(xe)]
            attSum = np.zeros(len(xe))
            for filt in range(self.carCal.filterCount):
                if filt == self.varFilter:
                    fwid = fw[line]
                else:
                    fwid = self.carCal.filterWidth[filt]
                attSum = attSum + fwid*self.carCal.filterAtten[filt].getMu()[:len(xe)]
            # this is the key integral done as a simple sum. Can ignore width of each value
            # as constant energy steps, so cancels in I/I0
            at_se = se*np.exp(-attSum-tarAtt.getMu()[:len(xe)]*tw[line])*ec*(1-np.exp(-attDet))
            # remove nan's - why are nan's present? exp overflow gives inf, multiply by 0 gives nan
            # in most cases nans are OK to ignore.
            at_se_finite = at_se[np.logical_not(np.isnan(at_se))]
            i0 = np.sum(at_se_finite)
            #
            count = 0
            # loop over required attenuation values
            for muwid in mulist:
                attSam = muwid*corMat.getMu()[:len(xe)]
                at_se_sample = at_se * np.exp(-attSam)
                #
                # remove nan's - see above
                #  should not be needed here if fit is OK
                at_se_sample_finite = at_se_sample[np.logical_not(np.isnan(at_se_sample))]
                i_sample = np.sum(at_se_sample_finite)
                if i0==0. and self.verbose:
                    print("warn: i0 zero at ",line)
                    i0=1.
                if i_sample == 0 and self.verbose:
                    print("i_sample=0")
                if i_sample < 0. and self.verbose:
                    print("i_sample<0",at_se_sample[:8])
                #sumSq = sumSq + ( (i_sample/i0) - self.carCal.getAvAtten(line,sample) ) ** 2
                attout[line,count] = np.log(i0/i_sample)
                count = count+1
            #
            try:
                polyfit[line,0:odpoly+1] = np.polyfit(attout[line,1:],attin[1:]/attout[line,1:],odpoly)
                xpolyfit[line,0:xtekodpoly+1] = np.polyfit(attout[line,1:],attin[1:]/attout[line,1:],xtekodpoly)
            except:
                print("*** Polynomial fit of result failed")
        #
        # following carousel.pro, attout is the apparent attenuation or the x-axis of our correction
        # graph. the y-axis should be the actual attenuation at monochromatic energy corEn for the
        # correction material. Since the density of the correction material is unknown, which is the
        # main point of this code, width and density are not used here. For fitting of a polynomial
        # through (0,0) can divide y values by x values, ignoring origin (first point in this case).
        #

        return attout,attin,polyfit,xpolyfit