#!/usr/bin/python

'''

The header is defined as a class and the attributes
of header are all input parameters.

'''

import numpy as np

class header(object):

    def __init__(self, path):
        self.path = path
        file = open(path,'r')
        #self.file = file = file.readlines()
        self.file = file.readlines()

        line_num = 0
        for line in self.file:
            ## Determine lines defining the qmmm block ##
            if '&qmmm' in line:
                qmmm_block_start = line_num
            if '&endqmmm' in line:
                qmmm_block_end  = line_num
            ## Determine lines defining the moldyn block ##
            if '&moldyn' in line:
                moldyn_block_start = line_num
            if '&endmoldyn' in line:
                moldyn_block_end = line_num
            line_num += 1

        line_num = 0
        for line in self.file:
            '''
            ## Start geometry optimization ##
            if 'maxcyc=' in line and 'dav_maxcyc=' not in line:
                self.maxcyc = np.int(line.split()[0][len('maxcyc='):-1])
            if 'ntpr=' in line:
                self.ntpr = np.int(line.split()[0][len('ntpr='):-1])
            if 'grms_tol=' in line:
                self.grms_tol = np.float(line.split()[0][len('grms_tol='):-1])
            ## End geometry optimization ##

            ## Start ground-state and output parameters ##
            if 'qm_theory=' in line:
                self.qm_theory = np.str(line.split()[0][len('qm_theory='):-1])
            if 'scfconv=' in line:
                self.scfconv = np.float(line.split()[0][len('scfconv='):-1])
            if 'verbosity=' in line:
                self.verbosity = np.int(line.split()[0][len('verbosity='):-1])
            if 'printdipole=' in line:
                self.printdipole = np.int(line.split()[0][len('printdipole='):-1])
            if 'printbondorders=' in line:
                self.printbondorders = np.int(line.split()[0][len('printbondorders='):-1])
            if 'density_predict=' in line:
                self.density_predict = np.int(line.split()[0][len('density_predict='):-1])
            if 'itrmax=' in line:
                self.itrmax = np.int(line.split()[0][len('itrmax='):-1])
            ## End ground-state and output parameters ##

            ## Start excited-state parameters ##
            if 'exst_method=' in line:
                self.exst_method = np.int(line.split()[0][len('exst_method='):-1])
            if 'dav_guess=' in line:
                self.dav_guess = np.int(line.split()[0][len('dav_guess='):-1])
            if 'ftol0=' in  line:
                self.ftol0 = np.float(line.split()[0][len('ftol0='):-1])
            if 'ftol1=' in line:
                self.ftol1 = np.float(line.split()[0][len('ftol1='):-1])
            if 'dav_maxcyc=' in line:
                self.dav_maxcyc = np.int(line.split()[0][len('dav_maxcyc='):-1])
            if 'printcharges=' in line:
                self.printcharges = np.int(line.split()[0][len('printcharges='):-1])
            if 'calcxdens=' in line:
                self.calcxdens = np.str(line.split()[0][len('calcxdens='):-1])
            ## End excited-state parameters ##

            ## Start solvent models and external electric fields ##
            if 'solvent_model=' in line:
                self.solvent_model = np.int(line.split()[0][len('solvent_model='):-1])
            if 'potential_type=' in line:
                self.potential_type = np.int(line.split()[0][len('potential_type='):-1])
            if 'onsager_radius=' in line:
                self.onsager_radius = np.float(line.split()[0][len('onsager_radius='):-1])
            if 'ceps=' in line:
                self.ceps = np.float(line.split()[0][len('ceps='):-1])
            if 'linmixparam=' in line:
                self.linmixparam = np.float(line.split()[0][len('linmixparam='):-1])
            if 'cosmo_scf_ftol=' in line:
                self.cosmo_scf_ftol = np.float(line.split()[0][len('cosmo_scf_ftol='):-1])
            if 'doZ=' in line:
                self.doZ = np.str(line.split()[0][len('doZ='):-1])
            if 'index_of_refraction=' in line:
                self.index_of_refraction = np.float(line.split()[0][len('index_of_refraction='):-1])
            if 'EF=' in line:
                self.EF = np.int(line.split()[0][len('EF='):-1])
            if 'Ex=' in line:
                self.Ex = np.float(line.split()[0][len('Ex='):-1])
            if 'Ey=' in line:
                self.Ey = np.float(line.split()[0][len('Ey='):-1])
            if 'Ez=' in line:
                self.Ez = np.float(line.split()[0][len('Ez='):-1])
            ## End solvent models and external electric fields ##
            '''
            ## Start general parameters ##
            if 'natoms=' in line:
                self.natoms = np.int(line.split()[0][len('natoms='):-1])
            '''
            if 'rnd_seed=' in line:
                self.rnd_seed = np.int(line.split()[0][len('rnd_seed='):-1])
            '''
            if 'bo_dynamics_flag=' in line:
                self.bo_dynamics_flag = np.int(line.split()[0][len('bo_dynamics_flag='):-1])
            if 'exc_state_init=' in line:
                self.exc_state_init = np.int(line.split()[0][len('exc_state_init='):-1])
            if 'n_exc_states_propagate=' in line:
                self.n_exc_states_propagate = np.int(line.split()[0][len('n_exc_states_propagate='):-1])
            ## End general parameters ###

            ## Start dynamics parameters ##
            if 'time_init=' in line:
                self.time_init = np.float(line.split()[0][len('time_init='):-1])
            if 'time_step=' in line:
                self.time_step = np.float(line.split()[0][len('time_step='):-1])
            if 'n_class_steps=' in line:
                self.n_class_steps = np.int(line.split()[0][len('n_class_steps='):-1])
            if 'n_quant_steps=' in line:
                self.n_quant_steps = np.int(line.split()[0][len('n_quant_steps='):-1])
            '''
            if 'moldyn_deriv_flag=' in line:
                self.moldyn_deriv_flag = np.int(line.split()[0][len('moldyn_deriv_flag='):-1])
            if 'num_deriv_step=' in line:
                self.num_deriv_step = np.float(line.split()[0][len('num_deriv_step='):-1])
            if 'rk_tolerance=' in line:
                self.rk_tolerance = np.float(line.split()[0][len('rk_tolerance='):-1])
            ## End dynamics parameters ##

            ## Start Nonadiabatic parameters ##
            if 'decoher_type=' in line:
                self.decoher_type = np.int(line.split()[0][len('decoher_type='):-1])
            if 'decoher_e0=' in line:
                self.decoher_e0 = np.float(line.split()[0][len('decoher_e0='):-1])
            if 'decoher_c=' in line:
                self.decoher_c = np.float(line.split()[0][len('decoher_c='):-1])
            if 'dotrivial=' in line:
                self.dotrivial = np.int(line.split()[0][len('dotrivial='):-1])
            if 'quant_step_reduction_factor=' in line:
                self.quant_step_reduction_factor = np.float(line.split()[0][len('quant_step_reduction_factor='):-1])
            ## End Nonadiabatic parameters ##

            ## Start thermostat parameters ##
            if 'therm_type=' in line:
                self.therm_type = np.int(line.split()[0][len('therm_type='):-1])
            if 'therm_temperature=' in line:
                self.therm_temperature = np.float(line.split()[0][len('therm_temperature='):-1])
            if 'therm_friction=' in line:
                self.therm_friction = np.float(line.split()[0][len('therm_friction='):-1])
            if 'berendsen_relax_const=' in line:
                self.berendsen_relax_const = np.float(line.split()[0][len('berendsen_relax_const='):-1])
            if 'heating=' in line:
                self.heating = np.int(line.split()[0][len('heating='):-1])
            if 'heating_steps_per_degree=' in line:
                self.heating_steps_per_degree = np.int(line.split()[0][len('heating_steps_per_degree='):-1])
            ## End thermostat parameters ##
            '''
            ## Start output and log parameters ##
            if 'verbosity=' in line and moldyn_block_start < line_num < moldyn_block_end:
                self.moldyn_verbosity = np.int(line.split()[0][len('verbosity='):-1])
            if 'out_data_steps=' in line:
                self.out_data_steps = np.int(line.split()[0][len('out_data_steps='):-1])
            if 'out_coords_steps=' in line:
                self.out_coords_steps = np.int(line.split()[0][len('out_coords_steps='):-1])
            '''
            if 'out_data_cube=' in line:
                self.out_data_cube = np.int(line.split()[0][len('out_data_cube='):-1])
            if 'out_count_init=' in line:
                self.out_count_init = np.int(line.split()[0][len('out_count_init='):-1])
            ## End output and log parameters ##
            '''
            ## Check if coefficients are set ##
            if 'quant_amp_phase_flag' in line:
                header.quant_amp_phase_flag = 'The quant_amp_phase_flag is in the header.'

            line_num += 1


# coding=utf-8
# !/usr/bin/python3.6 ## Please use python 3.6
"""
__synopsis__    : Builds a matching network, the training and evaluation ops as well as data_loader augmentation routines.

__description__ : Builds a matching network, the training and evaluation ops as well as data_loader augmentation routines.
__project__     : MNXC
__author__      : Samujjwal Ghosh <cs16resch01001@iith.ac.in>
__version__     : "0.1"
__date__        : "08-11-2018"
__copyright__   : "Copyright (c) 2019"
__license__     : This source code is licensed under the MIT-style license found in the LICENSE file in the root
                  directory of this source tree.

__classes__     : MatchingNetwork

__variables__   :

__methods__     :
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from models import Attn
from models import PairCosineSim
from logger.logger import logger
from models import BiLSTM
from models import EmbedText

from config import configuration as config

class MatchingNetwork(nn.Module):
    """Builds a matching network, the training and evaluation ops as well as data_loader augmentation routines."""

    def __init__(self,num_channels: int,layer_size: int,fce: bool = config["model"]["fce"]) -> None:
        """
        Builds a matching network, the training and evaluation ops as well as data_loader augmentation routines.

        :param dropout: A placeholder of type torch.float32 denotes the amount of dropout to be used
        :param batch_size: The batch size for the experiment
        :param num_channels: Number of channels of the samples.
        :param is_training: Flag indicating whether we are training or evaluating
        :param rotate_flag: Flag indicating whether to rotate the samples.
        :param fce: Flag indicating whether to use full context embeddings (i.e. apply an LSTM on the CNNText embeddings)
        # :param num_classes_per_set: Integer indicating the number of classes per set
        # :param num_samples_per_class: Integer indicating the number of samples per class
        :param classify_count: total number of classes. It changes the text_lstm size of the classifier g with a final FC layer.
        :param sample_input: size of the input sample. It is needed in case we want to create the last FC classification.
        """
        super(MatchingNetwork,self).__init__()
        self.fce = fce

        self.attn = Attn()
        self.cosine_dis = PairCosineSim.PairCosineSim()
        self.g = EmbedText(num_channels,layer_size)
        if self.fce:
            self.lstm = BiLSTM(input_size=self.g.output_size + 1)

    def forward(self,supports_x: torch.Tensor,supports_hots: torch.Tensor,targets_x: torch.Tensor,
                targets_hots: torch.Tensor,target_cat_indices: torch.Tensor,requires_grad: bool = True,
                batch_size: int = config["sampling"]["batch_size"],
                dropout_external: float = config["model"]["dropout_external"]) -> [torch.Tensor,torch.Tensor]:
        """
        Builds graph for Matching Networks, produces losses and summary statistics.

        :param target_cat_indices:
        :param requires_grad:
        :param dropout_external:
        :param batch_size:
        :param supports_x: A tensor containing the support set samples.
            [batch_size, sequence_size, n_channels, 28]
            torch.Size([32, 25, 1, 28])
        :param supports_hots: A tensor containing the support set labels.
            [batch_size, sequence_size, n_classes]
            torch.Size([32, 25, 5])
        :param targets_x: A tensor containing the target sample (sample to produce label for).
            [batch_size, n_channels, 28]
            torch.Size([32, 5, 1, 28])
        :param targets_hots: A tensor containing the target label.
            [batch_size, 1]
            torch.Size([32, 5])
        :return:
        """
        ## Convert target indices to Pytorch multi-label loss format.
        # target_y_mlml = self.create_mlml_data(target_cat_indices, output_shape=targets_hots.shape)

        logger.debug("targets: \n{}".format(targets_x))
        logger.debug("supports: \n{}".format(supports_x))
        logger.debug(
            "targets X targets: \n{}".format(self.cosine_dis(supports=targets_x,targets=targets_x,normalize=False)))
        logger.debug(
            "supports X supports: \n{}".format(self.cosine_dis(supports=supports_x,targets=supports_x,normalize=False)))
        logger.debug(
            "supports X targets: \n{}".format(self.cosine_dis(supports=supports_x,targets=targets_x,normalize=False)))
        ## Encode supports
        supports_x = self.g(supports_x,dropout_external=dropout_external)
        # logger.debug("supports_x [{}] output: {}".format(supports_x.shape,supports_x))

        ## Encode targets
        # targets_x = self.g(targets_x, dropout_external=dropout_external)
        # logger.debug("targets_x[{}] output: {}".format(targets_x.shape,targets_x))
        # logger.debug("emb_targets: \n{}".format(targets_x))
        logger.debug("emb_supports: \n{}".format(supports_x))

        # logger.debug("emb_targets X emb_targets: \n{}".format(self.cosine_dis(supports=targets_x, targets=targets_x, normalize=False)))
        logger.debug("emb_supports X emb_supports: \n{}".format(
            self.cosine_dis(supports=supports_x,targets=supports_x,normalize=False)))
        logger.debug("emb_supports X emb_targets: \n{}".format(
            self.cosine_dis(supports=supports_x,targets=targets_x,normalize=False)))
        if self.fce:
            # supports_x, _ = self.lstm(supports_x, requires_grad=requires_grad)
            targets_x,_ = self.lstm(targets_x,requires_grad=requires_grad)
            # logger.debug("FCE supports_x: \n{}".format(supports_x))
            logger.debug("FCE targets_x: \n{}".format(targets_x))
            logger.debug("FCE_targets X FCE_targets: \n{}".format(
                self.cosine_dis(supports=targets_x,targets=targets_x,normalize=False)))
            # logger.debug("FCE supports_x X FCE supports_x: \n{}".format(self.cosine_dis(supports=supports_x, targets=supports_x, normalize=False)))

        ## Calculate similarity between encoded supports and targets
        similarities = self.cosine_dis(supports=supports_x,targets=targets_x,normalize=True)

        logger.debug("FCE supports_x X FCE targets_x: \n{}".format(similarities))

        ## Produce predictions for target probabilities. targets_preds.shape = batch_size x # classes
        targets_preds = self.attn(similarities,supports_hots=supports_hots.float())
        # logger.debug("target_cat_indices: {}".format(target_cat_indices))
        # logger.debug("targets_preds output: {}".format(targets_preds))

        loss = 0
        # multilabel_batch_loss = F.multilabel_margin_loss(targets_preds, target_y_mlml.long())

        ## Calculate loss, need to calculate loss for each sample but for whole batch.
        for j in np.arange(targets_preds.size(1)):
            # logger.debug((targets_preds[:, j, :].shape, target_y_mlml.long()[:, j, :].shape))
            # logger.debug(targets_preds[:, j, :])
            # logger.debug(target_y_mlml.long()[:, j, :])
            # logger.debug(targets_preds[:, j, :][0])
            # logger.debug(target_y_mlml.long()[:, j, :][0])
            # loss += F.multilabel_margin_loss(targets_preds[:, j, :], target_y_mlml.long()[:, j, :])
            loss += F.cross_entropy(targets_preds[:,j,:],target_cat_indices[:,j].long())
        multilabel_batch_loss = loss / targets_x.size(1)

        return multilabel_batch_loss,targets_preds

    @staticmethod
    def create_mlml_data(target_cat_indices: list,output_shape: tuple) -> torch.Tensor:
        """
        Generates true labels in proper format for Pytorch Multilabel_Margin_Loss.

        Converts target indices to Pytorch 'multilabel_margin_loss' format. Takes class indices at the beginning and
        rest should be filled with -1.
        Link 1: https://gist.github.com/bartolsthoorn/36c813a4becec1b260392f5353c8b7cc#gistcomment-2742606
        Link 2: https://gist.github.com/bartolsthoorn/36c813a4becec1b260392f5353c8b7cc#gistcomment-2840045

        :param output_shape: Shape of the output = batch_size, target count, # labels.
        :param target_cat_indices: List of categories.
        """
        target_y_mlml = torch.full(output_shape,-1)  ## Createing a tensor filled with -1.
        ## Replacing -1 with label indices at the begining of the tensor row.
        for i in np.arange(target_y_mlml.size(0)):
            # logger.debug(target_cat_indices[i])
            for j in np.arange(target_y_mlml.size(1)):
                # logger.debug(type(target_cat_indices[i][j]))
                # logger.debug((i,j,target_cat_indices[i][j]))
                for k in range(len(target_cat_indices[i][j])):
                    # logger.debug(target_cat_indices[i][j][k])
                    target_y_mlml[i,j,k] = torch.tensor(target_cat_indices[i][j][k],dtype=torch.int32)

        return target_y_mlml.long()

if __name__ == '__main__':
    import torch

    support_set = torch.rand(4,2,4)  # [batch_size, sequence_size, input_size]
    support_set_hot = torch.zeros(4,2,1)  # [batch_size, n_classes]
    x_hat = torch.rand(4,2,4)
    x_hat_hot = torch.zeros(4,2,1)
    logger.debug(support_set_hot)
    logger.debug(x_hat_hot)

    # support_set = torch.tensor([[[1., 0.4],
    #                              [1., 1.]],
    #                             [[1., 0.4],
    #                              [0., 1.5]],
    #                             [[1., 0.4],
    #                              [1., 1.5]]])
    #
    support_set_hot = torch.tensor([[[1.,0.],
                                     [0.,1.]],

                                    [[1.,0.],
                                     [0.,1.]],

                                    [[1.,0.],
                                     [0.,1.]],

                                    [[1.,0.],
                                     [0.,1.]]])
    #
    # x_hat = torch.tensor([[[1., 0.4],
    #                        [0., 1.5]],
    #                       [[1., 0.4],
    #                        [1., 1.5]]])
    #
    x_hat_hot = torch.tensor([[[1.,0.],
                               [0.,1.]],

                              [[1.,0.],
                               [0.,1.]],

                              [[1.,0.],
                               [0.,1.]],

                              [[1.,0.],
                               [0.,1.]]])

    logger.debug(support_set.shape)
    logger.debug(support_set_hot.shape)
    logger.debug(x_hat.shape)
    logger.debug(x_hat_hot.shape)

    cls = MatchingNetwork(input_size=4,hid_size=4,classify_count=5,use_cuda=False)
    logger.debug(cls)
    sim = cls.forward(support_set,support_set_hot,x_hat,x_hat_hot,batch_size=4)
    logger.debug(sim)


import numpy as np

try:
    import cPickle as pickle
except ImportError:
    import pickle

from src.SpectralAnalysis import utils
from src.SpectralAnalysis import LightCurve
from src.SpectralAnalysis import powerspectrum
from src.SpectralAnalysis import bayes
from src.SpectralAnalysis import mcmc
from src.SpectralAnalysis.parametricmodels import pl, bpl

class Burst(object):
    def __init__(self, bstart, blength,
                 energies=None,
                 photons=None,
                 events=None,
                 filename=None,
                 instrument="gbm",
                 fnyquist=4096.0,
                 norm='leahy',
                 fluence=None,
                 epeak=None,
                 ttrig=None,
                 addfrac=0.2):

        ### which instrument was used to record this data
        ### note: this makes a difference in terms of file formats
        ### and general data structure
        self.instrument = instrument

        if ":" in str(bstart):
            self.bst = self.convert_time(bstart)
        else:
            self.bst = bstart

        self.blen = blength + 2. * addfrac * blength

        start = self.bst - addfrac * blength
        length = (1. + 2. * addfrac) * self.blen
        end = start + length

        ### assume burst length is in seconds
        self.fluence = fluence
        self.epeak = epeak
        self.ttrig = ttrig

        ### data is in form of Photon objects such that time/energy filtering
        ### becomes easy
        if photons is None and filename:
            self.read_data(filename)

        elif not photons is None:
            self.photons = photons

        else:
            raise Exception("Data missing! You must specify either a photon "
                            "object or a file name from which to read the data!")

        if not events is None:
            self.energies = events
        startind = photons.searchsorted(start)
        endind = photons.searchsorted(end)
        self.photons = self.photons[startind:endind]
        self.energies = self.energies[startind:endind]

        ### filter for energy selection, if this is specified
        # if energies:
        #    gt.Data.filterenergy(self, energies[0], energies[1])
        if energies:
            self.photons = np.array([s for s, e in zip(self.photons, self.energies) if energies[0] <= e <= energies[1]])
            self.energies = np.array(
                [e for s, e in zip(self.photons, self.energies) if energies[0] <= e <= energies[1]])

        #### filter for burst times
        # gt.Data.filterburst([self.bst-0.1*self.blen, self.bend+0.1*self.blen])

        ### make a light curve
        self.time = self.photons
        # print("length time: " + str(len(self.time)))
        # print("tseg: " + str(self.time[-1] - self.time[0]))

        # self.time = np.array([s.time for s in self.photons])
        self.lc = LightCurve.LightCurve(self.time, timestep=0.5 / fnyquist, tseg=self.blen, tstart=self.bst)

        # print("length lc: " + str(len(self.lc.time)))

        ### make a periodogram
        self.ps = powerspectrum.PowerSpectrum(self.lc, norm=norm)

        return

    #### CONVERT TIME ##############
    #
    # convert from HH:MM:SS to seconds
    #
    def convert_time(self, time):
        hours = float(time[:2]) * 3600.0
        minutes = float(time[3:5]) * 60.0
        seconds = float(time[6:])
        tnew = hours + minutes + seconds
        return tnew

    #### READ DATA FROM FILE
    #
    # read time-tagged event data from file
    # keyword 'type' can be either one of
    #                - "ascii" = read ascii data
    #                - "pickle" = read data from python pickle file
    #                   that contains a list of photon objects
    def read_data(self, filename, type="ascii"):
        if type in ["a", "ascii"]:
            data = utils.conversion(filename)
            time = np.array([float(t) for t in data[0]])
            events = np.array([float(e) for e in data[1]])
            self.photons = [utils.Photon(t, e) for t, e in zip(time, events)]
        elif type in ["p", "pickle"]:
            self.photons = utils.getpickle(filename)
        else:
            raise Exception("File type not recognized! Must be one of 'pickle' or 'ascii'!")
        return

    def bayesian_analysis(self, namestr='test', nchain=500, niter=100, nsim=1000, m=1, fitmethod='bfgs'):

        btest = bayes.Bayes(self.ps, namestr=namestr, m=m)
        psfit, fakeper, self.model_summary = btest.choose_noise_model(pl, [2, 3, 0.5],
                                                                      bpl, [1, 3, 2, 3, 0.5],
                                                                      nchain=nchain, niter=niter, nsim=nsim,
                                                                      fitmethod=fitmethod)

        if not psfit:
            print("Analysis of burst " + str(namestr) + " failed. Returning ...")
            return

        else:
            if self.model_summary["p_lrt"][0] < 0.05:
                print("Model not adequately fit by a power law! Using broken power law instead!")
                self.model = bpl
                self.psfit = getattr(psfit, str(self.model).split()[1] + "fit")
            else:
                self.model = pl
                self.psfit = getattr(psfit, str(self.model).split()[1] + "fit")

            self.per_summary = btest.find_periodicity(self.model, self.psfit["popt"], nchain=nchain, niter=niter,
                                                      nsim=nsim, fitmethod=fitmethod)

            self.mcmc = self.per_summary["mcobs"]
            return

    #### MAKE AN MCMC SAMPLE #######################
    #
    # Runs the MarkovChainMonteCarlo code to make
    # an MCMC sample independent of the Bayes routines.
    #
    def mcmc_sample(self, func, pars, cov, nchain=500, niter=100, nsim=1000):

        mcobs = mcmc.MarkovChainMonteCarlo(self.ps, func=func, topt=pars, tcov=cov, nchain=nchain, niter=niter,
                                           nsim=nsim)

        return mcobs

    def save_burst(self, filename):

        burstfile = open(filename, 'w')
        pickle.dump(self, burstfile)
        burstfile.close()
        return

#####################################################
#####################################################
#####################################################

#### BURST SUBCLASS FOR GBM DATA ###################
#
# Subclass of class Burst to deal with GBM data
# and GBM specific issues
#
#
class GBMBurst(Burst):
    def __init__(self, bid, bstart, blength,
                 energies=None,
                 photons=None,
                 events=None,
                 filename=None,
                 instrument="gbm",
                 fnyquist=4096.0,
                 norm='leahy',
                 fluence=None,
                 epeak=None,
                 ttrig=None,
                 addfrac=0.2):
        ### set burst ID
        self.bid = bid

        ### if photons and filename aren't given, then data comes from procdata file
        if photons is None and filename is None:
            filename = "tte_bn" + str(bid) + "_procdata.dat"

        Burst.__init__(self, bstart, blength, energies,
                       photons, events, filename,
                       instrument, fnyquist, norm,
                       fluence, epeak, ttrig, addfrac)
        # super(self.__class__, self).__init__(bstart, blength, energies,
        #                               photons, events, filename,
        #                               instrument, fnyquist, norm,
        #                               fluence, epeak, ttrig, addfrac)

        return

    def read_data(self, filename, filetype='ascii', det="combined"):
        Burst.read_data(self, filename, type=filetype)

        evt = self.photons[det]
        self.photons = np.array([x.time for x in evt.photons])
        self.energies = np.array([x.energy for x in evt.photons])
        return


import sys
import os.path
from numpy import logical_and

# print sys.path
sys.path.append("~/git/py-mmcif")
from mmcif.io.PdbxReader import PdbxReader
from NEFTranslator import NEFTranslator
from math import sqrt, acos
import numpy
import pynmrstar
from monkeypatch import patch_parser
import xml.etree.ElementTree as ET
from operator import itemgetter

patch_parser(pynmrstar)

import json

class Validate_NMR_Restraints:
    """
    NMR restraints validation module
    """
    __version__ = "v0.8"

    def __init__(self,cif_file, star_file):
        self.run_validation(cif_file,star_file)

    def run_validation(self,cif_file,star_file):
        #nt = NEFTranslator.NEFTranslator()
        pdb, atom_ids = self.get_coordinates(cif_file)
        max_models = len(pdb.keys())
        distance, angle, chain_dict = self.get_restraints(star_file)
        dv = self.calculate_distance_violations(pdb, distance)
        av = self.calculate_angle_violations(pdb, angle)
        self.write_xml(dv, av, distance, angle, atom_ids)
        # self.write_xml_simple(dv, av)
        dist_viol_stat, dist_viol = self.calculate_violation_statistics(dv)
        ang_viol_stat, ang_viol = self.calculate_violation_statistics(av)
        self.bin_distance_violations(dv)
        sorted_dist_viol_stat = sorted(dist_viol_stat, reverse=True, key=itemgetter(0, 3))
        sorted_dist_viol = sorted(dist_viol, reverse=True, key=itemgetter(0))
        sorted_ang_viol_stat = sorted(ang_viol_stat, reverse=True, key=itemgetter(0, 3))
        sorted_ang_viol = sorted(ang_viol, reverse=True, key=itemgetter(0))
        type_stat_dist = self.restraints_type_statistics(dist_viol_stat, max_models)
        type_stat_ang = self.restraints_type_statistics(ang_viol_stat, max_models)
        json_data = self.generate_json(type_stat_dist, type_stat_ang, sorted_dist_viol_stat, sorted_dist_viol,
                                       sorted_ang_viol_stat, sorted_ang_viol)
        with open('data_json.json', 'w') as write_file:
            json.dump(json_data, write_file)
        return json_data

    @staticmethod
    def generate_json(type_stat_dist, type_stat_ang, dist_viol_stat, dist_viol, ang_viol_stat, ang_viol):
        restraints_validation = {}
        distance = {}
        distance['summary'] = type_stat_dist[0]
        distance['violated'] = type_stat_dist[2]
        distance['consistently_violated'] = type_stat_dist[1]
        distance['sorted_average_violations'] = dist_viol_stat
        distance['sorted_violations'] = dist_viol
        angle = {}
        angle['summary'] = type_stat_ang[0]
        angle['violated'] = type_stat_ang[2]
        angle['consistently_violated'] = type_stat_ang[1]
        angle['sorted_average_violations'] = ang_viol_stat
        angle['sorted_violations'] = ang_viol
        restraints_validation['distance'] = distance
        restraints_validation['angle'] = angle
        return restraints_validation

    @staticmethod
    def restraints_type_statistics(violation_statistics, max_models):
        """
        Counts the number of restraints and violations in each restraints type(distance: intraresidue, sequential,
        medium, long; angle : PHI, PSI, etc.. )
        :param violation_statistics: output from calculate_violation_statisitcs
        :param max_models: Number of models in the ensemble (required to estimate the consistently violated restraitns)
        :return: three statistics in a dictionary format (General summary, Consistently violated, Violated)
        """
        rest_type_v = [i[5] for i in violation_statistics if i[3] != 0]  # Violated at lease in one model
        rest_type_cv = [i[5] for i in violation_statistics if i[3] == max_models]  # Violated in all modes
        rest_type_all = [i[5] for i in violation_statistics]
        uniq_type = list(set(rest_type_all))
        type_count_all = {}
        type_count_v = {}
        type_count_cv = {}
        total_all = len(violation_statistics)
        type_count_all['total'] = (total_all, 100.00)  # Count and percentage
        for i in uniq_type:
            type_count_all[i] = (
                rest_type_all.count(i), round(((float(rest_type_all.count(i)) / float(total_all)) * 100.00), 2))
            type_count_v[i] = (
                rest_type_v.count(i),
                round(((float(rest_type_v.count(i)) / float(rest_type_all.count(i))) * 100.00), 2))
            type_count_cv[i] = (
                rest_type_cv.count(i),
                round(((float(rest_type_cv.count(i)) / float(rest_type_all.count(i))) * 100.00), 2))
        return type_count_all, type_count_cv, type_count_v

    @staticmethod
    def get_coordinates(cif_file):
        """
        Extract coordinate information from cif file as a dictionary
        {model_id : {(seq_id,chain_id,res_id,atom_id) : array[x,y,x],...},...}
        :param cif_file: Input coordinate file
        :return: dictionary
        """
        cif_data = []
        ifh = open(cif_file, 'r')
        pRd = PdbxReader(ifh)
        pRd.read(cif_data)
        ifh.close()
        c0 = cif_data[0]
        atom_site = c0.getObj('atom_site')
        max_models = int(atom_site.getValue('pdbx_PDB_model_num', -1))
        col_names = atom_site.getAttributeList()
        model_id = col_names.index('pdbx_PDB_model_num')
        x_id = col_names.index('Cartn_x')
        y_id = col_names.index('Cartn_y')
        z_id = col_names.index('Cartn_z')
        atom_id = col_names.index('label_atom_id')
        comp_id = col_names.index('label_comp_id')
        asym_id = col_names.index('label_asym_id')
        entity_id = col_names.index('label_entity_id')
        seq_id = col_names.index('label_seq_id')
        icode_id = col_names.index('pdbx_PDB_ins_code')
        alt_id = col_names.index('label_alt_id')
        aut_seq_id = col_names.index('auth_seq_id')
        aut_aym_id = col_names.index('auth_asym_id')
        pdb_models = {}
        atom_ids = {}
        for model in range(1, max_models + 1):
            pdb = {}
            aid = {}
            for dat in atom_site.getRowList():
                if int(dat[model_id]) == model:
                    aid[(dat[seq_id], dat[asym_id], dat[comp_id], dat[atom_id])] = \
                        (dat[entity_id], dat[asym_id], dat[comp_id], dat[seq_id], dat[aut_seq_id],
                         dat[alt_id], dat[icode_id], dat[aut_aym_id])
                    pdb[(dat[seq_id], dat[asym_id], dat[comp_id], dat[atom_id])] = \
                        numpy.array([float(dat[x_id]), float(dat[y_id]), float(dat[z_id])])
            pdb_models[model] = pdb
            atom_ids[model] = aid
        return pdb_models, atom_ids

    @staticmethod
    def get_restraints(star_file):
        """
        Extracts restraints from NMR-STAR file
        :param star_file: NMR-STAR file
        :return: distance restraints in dictionary format, angle restraints in dictionary format, chain map as dictionary
        """
        lp_flg = 0
        dat_flg = 1
        try:
            star_data = pynmrstar.Entry.from_file(star_file)
            dist_sf = star_data.get_saveframes_by_category('general_distance_constraints')
        except ValueError:
            try:
                sf_data = pynmrstar.Saveframe.from_file(star_file)
                if sf_data.get_tag('Sf_category')[0] == 'general_distance_constraints':
                    dist_sf = [sf_data]
                else:
                    print("Error : file doesn't have restraints data")
                    dat_flg = 0
            except ValueError:
                try:
                    lp_data = pynmrstar.Loop.from_file(star_file)
                    lp_flg = 1
                    if lp_data.category == '_Gen_dist_constraint':
                        dist_sf = [lp_data]
                    else:
                        print("Error : file doesn't have restraints data")
                        dat_flg = 0
                except ValueError:
                    print("ERROR : File contains no valid saveframe or loop")
                    dat_flg = 0
        except IOError:
            print("Error file not found")
            dat_flg = 0

        try:
            entity_assembly = star_data.get_loops_by_category('_Entity_assembly')[0]
            col_names = entity_assembly.get_tag_names()
            id_index = col_names.index('_Entity_assembly.ID')
            asym_index = col_names.index('_Entity_assembly.Asym_ID')
            chain_dict = {}
            for row in entity_assembly:
                chain_map[row[id_index]] = row[asym_index]
        except IndexError:
            print("Entity_assembly loop not found, No chain mapping created")
            chain_dict = {}
        if dat_flg:
            dist_dict2 = {}
            dist_dict = {}
            for sf in dist_sf:
                if lp_flg:
                    dat = sf
                else:
                    dat = sf.get_loop_by_category('_Gen_dist_constraint')
                col_names = dat.get_tag_names()
                rest_id = col_names.index('_Gen_dist_constraint.ID')
                seq_id_1 = col_names.index('_Gen_dist_constraint.Comp_index_ID_1')
                asym_id_1 = col_names.index('_Gen_dist_constraint.Auth_asym_ID_1')
                entity_id_1 = col_names.index('_Gen_dist_constraint.Entity_assembly_ID_1')
                comp_id_1 = col_names.index('_Gen_dist_constraint.Comp_ID_1')
                atom_id_1 = col_names.index('_Gen_dist_constraint.Atom_ID_1')
                seq_id_2 = col_names.index('_Gen_dist_constraint.Comp_index_ID_2')
                asym_id_2 = col_names.index('_Gen_dist_constraint.Auth_asym_ID_2')
                entity_id_2 = col_names.index('_Gen_dist_constraint.Entity_assembly_ID_2')
                comp_id_2 = col_names.index('_Gen_dist_constraint.Comp_ID_2')
                atom_id_2 = col_names.index('_Gen_dist_constraint.Atom_ID_2')
                lb_id = col_names.index('_Gen_dist_constraint.Distance_lower_bound_val')
                ub_id = col_names.index('_Gen_dist_constraint.Distance_upper_bound_val')
                list_id = col_names.index('_Gen_dist_constraint.Gen_dist_constraint_list_ID')
                r_dict = {}
                for rest in dat:
                    if rest[rest_id] not in r_dict.keys():
                        r_dict[rest[rest_id]] = []
                    if (rest[list_id], rest[rest_id]) not in dist_dict2.keys():
                        dist_dict2[(rest[list_id], rest[rest_id])] = []
                    if rest[asym_id_1] == '.':
                        eid1 = chain_dict[rest[entity_id_1]]
                    else:
                        eid1 = rest[asym_id_1]
                    if rest[asym_id_2] == '.':
                        eid2 = chain_dict[rest[entity_id_2]]
                    else:
                        eid2 = rest[asym_id_2]

                    atom1 = (rest[seq_id_1], eid1, rest[comp_id_1], rest[atom_id_1])
                    atom2 = (rest[seq_id_2], eid2, rest[comp_id_2], rest[atom_id_2])
                    if atom1[1] != atom2[1]:
                        cat = 'long'
                    elif abs(int(atom1[0]) - int(atom2[0])) == 0:
                        cat = 'intraresidue'
                    elif abs(int(atom1[0]) - int(atom2[0])) == 1:
                        cat = 'sequential'
                    elif 1 < abs(int(atom1[0]) - int(atom2[0])) < 5:
                        cat = 'medium'
                    elif abs(int(atom1[0]) - int(atom2[0])) >= 5:
                        cat = 'long'
                    try:
                        lb = float(rest[lb_id])
                    except ValueError:
                        lb = -999.9
                    try:
                        ub = float(rest[ub_id])
                    except ValueError:
                        ub = 999.9
                    if lb == -999.9 and ub == 999.9:
                        print("Error: Distance restraint value not readable for restraint id {}; "
                              "for atoms {},{}".format(rest[rest_id], atom1, atom2))
                    else:
                        r_dict[rest[rest_id]].append([atom1, atom2, lb, ub, cat])
                        dist_dict2[(rest[list_id], rest[rest_id])].append([atom1, atom2, cat, lb, ub])
                if lp_flg:
                    dist_dict['distance_restraints'] = r_dict
                else:
                    dist_dict[sf.name] = r_dict
        else:
            dist_dict2 = None
        dat_flg = 1
        lp_flg = 0
        try:
            star_data = pynmrstar.Entry.from_file(star_file)
            ang_sf = star_data.get_saveframes_by_category('torsion_angle_constraints')
        except ValueError:
            try:
                sf_data = pynmrstar.Saveframe.from_file(star_file)
                if sf_data.get_tag('Sf_category')[0] == 'torsion_angle_constraints':
                    ang_sf = [sf_data]
                else:
                    print("Error : file doesn't have restraints data")
                    dat_flg = 0
            except ValueError:
                try:
                    lp_data = pynmrstar.Loop.from_file(star_file)
                    lp_flg = 1
                    if lp_data.category == '_Torsion_angle_constraint':
                        ang_sf = [lp_data]
                    else:
                        print("Error : file doesn't have restraints data")
                        dat_flg = 0
                except ValueError:
                    print("ERROR : File contains no valid saveframe or loop")
                    dat_flg = 0
        except IOError:
            print("Error file not found")
            dat_flg = 0
        if dat_flg:
            angle_dict = {}
            angle_dict2 = {}
            for sf in ang_sf:
                if lp_flg:
                    dat = sf
                else:
                    dat = sf.get_loop_by_category('_Torsion_angle_constraint')
                col_names = dat.get_tag_names()
                rest_id = col_names.index('_Torsion_angle_constraint.ID')
                rest_name_id = col_names.index("_Torsion_angle_constraint.Torsion_angle_name")
                seq_id_1 = col_names.index('_Torsion_angle_constraint.Comp_index_ID_1')
                asym_id_1 = col_names.index('_Torsion_angle_constraint.Auth_asym_ID_1')
                entity_id_1 = col_names.index('_Torsion_angle_constraint.Entity_assembly_ID_1')
                comp_id_1 = col_names.index('_Torsion_angle_constraint.Comp_ID_1')
                atom_id_1 = col_names.index('_Torsion_angle_constraint.Atom_ID_1')
                seq_id_2 = col_names.index('_Torsion_angle_constraint.Comp_index_ID_2')
                asym_id_2 = col_names.index('_Torsion_angle_constraint.Auth_asym_ID_2')
                entity_id_2 = col_names.index('_Torsion_angle_constraint.Entity_assembly_ID_2')
                comp_id_2 = col_names.index('_Torsion_angle_constraint.Comp_ID_2')
                atom_id_2 = col_names.index('_Torsion_angle_constraint.Atom_ID_2')
                seq_id_3 = col_names.index('_Torsion_angle_constraint.Comp_index_ID_3')
                asym_id_3 = col_names.index('_Torsion_angle_constraint.Auth_asym_ID_3')
                entity_id_3 = col_names.index('_Torsion_angle_constraint.Entity_assembly_ID_3')
                comp_id_3 = col_names.index('_Torsion_angle_constraint.Comp_ID_3')
                atom_id_3 = col_names.index('_Torsion_angle_constraint.Atom_ID_3')
                seq_id_4 = col_names.index('_Torsion_angle_constraint.Comp_index_ID_4')
                asym_id_4 = col_names.index('_Torsion_angle_constraint.Auth_asym_ID_4')
                entity_id_4 = col_names.index('_Torsion_angle_constraint.Entity_assembly_ID_4')
                comp_id_4 = col_names.index('_Torsion_angle_constraint.Comp_ID_4')
                atom_id_4 = col_names.index('_Torsion_angle_constraint.Atom_ID_4')
                lb_id = col_names.index('_Torsion_angle_constraint.Angle_lower_bound_val')
                ub_id = col_names.index('_Torsion_angle_constraint.Angle_upper_bound_val')
                list_id = col_names.index('_Torsion_angle_constraint.Torsion_angle_constraint_list_ID')
                r_dict = {}
                for rest in dat:
                    if rest[rest_id] not in r_dict.keys():
                        r_dict[rest[rest_id]] = []
                    if (rest[list_id], rest[rest_id]) not in angle_dict2.keys():
                        angle_dict2[(rest[list_id], rest[rest_id])] = []

                    if rest[asym_id_1] == '.':
                        eid1 = chain_dict[rest[entity_id_1]]
                    else:
                        eid1 = rest[asym_id_1]
                    if rest[asym_id_2] == '.':
                        eid2 = chain_dict[rest[entity_id_2]]
                    else:
                        eid2 = rest[asym_id_2]
                    if rest[asym_id_3] == '.':
                        eid3 = chain_dict[rest[entity_id_3]]
                    else:
                        eid3 = rest[asym_id_3]
                    if rest[asym_id_4] == '.':
                        eid4 = chain_dict[rest[entity_id_4]]
                    else:
                        eid4 = rest[asym_id_4]

                    atom1 = (rest[seq_id_1], eid1, rest[comp_id_1], rest[atom_id_1])
                    atom2 = (rest[seq_id_2], eid2, rest[comp_id_2], rest[atom_id_2])
                    atom3 = (rest[seq_id_3], eid3, rest[comp_id_3], rest[atom_id_3])
                    atom4 = (rest[seq_id_4], eid4, rest[comp_id_4], rest[atom_id_4])
                    rest_name = rest[rest_name_id]
                    try:
                        lb = float(rest[lb_id])
                    except ValueError:
                        lb = -999.9
                    try:
                        ub = float(rest[ub_id])
                    except ValueError:
                        ub = 999.9
                    if lb == -999.9 and ub == 999.9:
                        print("Error: Distance restraint value not readable for restraint id {}; "
                              "for atoms {},{}".format(rest[rest_id], atom1, atom2))
                    else:
                        r_dict[rest[rest_id]].append([atom1, atom2, atom3, atom4, rest_name, lb, ub])
                        angle_dict2[(rest[list_id], rest[rest_id])].append(
                            [atom1, atom2, atom3, atom4, rest_name, lb, ub])
                if lp_flg:
                    angle_dict['angle_restraints'] = r_dict
                else:
                    angle_dict[sf.name] = r_dict
        else:
            angle_dict = None
            angle_dict2 = None
        return dist_dict2, angle_dict2, chain_dict

    @staticmethod
    def get_distance(c1, c2):
        """
        Calculates the distance between two coordinate points
        :param c1: array of x,y,z
        :param c2: array of x,y,z
        :return: distance between two ponts
        """
        return numpy.linalg.norm(c1 - c2)

    @staticmethod
    def get_dihedral_angle(c1, c2, c3, c4):
        """
        Calculates the dihedral angle from the given set of four coordinate values
        :param c1: array of x,y,z
        :param c2: array of x,y,z
        :param c3: array of x,y,z
        :param c4: array of x,y,z
        :return: angle in degrees
        """
        bv12 = c1 - c2
        bv32 = c3 - c2
        bv43 = c4 - c3
        pv13 = numpy.cross(bv12, bv32)
        pv24 = numpy.cross(bv43, bv32)
        pro = numpy.dot(pv13, pv24)
        sqdist13 = numpy.dot(pv13, pv13)
        sqdist24 = numpy.dot(pv24, pv24)
        cosin = pro / sqrt(sqdist13 * sqdist24)
        cosin - min(1.0, max(-1.0, cosin))
        angle = acos(cosin)

        if numpy.dot(pv13, numpy.cross(pv24, bv32)) < 0:
            angle = -angle
        return round(numpy.degrees(angle), 4)

    @staticmethod
    def r6sum(dist_list):
        """
        Calculates 1/r^6 sum for ambiguous restraints as recommended by NMR VTF
        :param dist_list: list of distances
        :return: r6 sum
        """
        return (sum([i ** (-6.) for i in dist_list])) ** (-1. / 6.)

    def calculate_distance_violations(self, coordinates, restraints):
        """
        Calculates violation for each restraint
        :param coordinates:  output from get_coordinates
        :param restraints:  output from get_restraints
        :return: dictionary { rest identifier : { model no : (value, 'type') }} example {('1', '1210'):
        {1: (0.22874009681108554, 'long'), 2: (0.084672752498432757, 'long')...}}
        """
        violations = {}
        for rest_id in restraints.keys():
            m = {}
            for model in coordinates.keys():
                dist_list = []
                for rest in restraints[rest_id]:
                    atom_1 = rest[0]
                    atom_2 = rest[1]
                    cat = rest[2]
                    lb = rest[3]
                    ub = rest[4]
                    pos_1 = coordinates[model][atom_1]
                    pos_2 = coordinates[model][atom_2]
                    d = self.get_distance(pos_1, pos_2)
                    dist_list.append(d)
                r6dist = self.r6sum(dist_list)
                if lb <= r6dist <= ub:
                    err = 0.0
                elif r6dist < lb:
                    err = abs(r6dist - lb)
                else:
                    err = abs(r6dist - ub)
                m[model] = (err, cat)
            violations[rest_id] = m
        return violations

    def calculate_angle_violations(self, coordinates, restraints):
        """
        Calculates violation for each restraint
        :param coordinates: output from get_coordinates
        :param restraints:  output from get_restraints
        :return: dictionary { rest identifier : { model no : (value, 'type') }} example {('1', '121'):
        {1: (0.22874009681108554, 'PSI'), 2: (0.084672752498432757, 'PHI')...}}
        """
        violations = {}
        for rest_id in restraints.keys():
            m = {}
            for model in coordinates.keys():
                ang_list = []
                for rest in restraints[rest_id]:
                    atom_1 = rest[0]
                    atom_2 = rest[1]
                    atom_3 = rest[2]
                    atom_4 = rest[3]
                    cat = rest[4]
                    lb = rest[5]
                    ub = rest[6]
                    pos_1 = coordinates[model][atom_1]
                    pos_2 = coordinates[model][atom_2]
                    pos_3 = coordinates[model][atom_3]
                    pos_4 = coordinates[model][atom_4]
                    ang = self.get_dihedral_angle(pos_1, pos_2, pos_3, pos_4)
                    ang_list.append(ang)
                avg_viol = numpy.mean(ang_list)
                if lb <= avg_viol <= ub:
                    err = 0.0
                elif avg_viol < lb:
                    err = abs(avg_viol - lb)
                else:
                    err = abs(avg_viol - ub)
                m[model] = (err, cat)
            violations[rest_id] = m
        return violations

    @staticmethod
    def write_xml(distance_violations, angle_violations, distance, angle, atom_ids):
        rest_ids = list(distance_violations.keys())
        model_ids = list(distance_violations[rest_ids[0]].keys())
        violations = ET.Element('RestraintsViolations')
        dist_viol = ET.SubElement(violations, 'DistanceViolations')
        for m_id in model_ids[:50]:
            models = ET.SubElement(dist_viol, 'Model')
            models.set('model', str(m_id))
            for r_id in rest_ids[:50]:
                for rest in distance[r_id]:
                    if distance_violations[r_id][m_id][0] > 0.0:
                        model = ET.SubElement(models, 'Violation')
                        model.set('rest_id', str(r_id[1]))
                        model.set('rest_list_id', str(r_id[0]))
                        model.set('model_1', str(m_id))
                        model.set('chain_1', atom_ids[m_id][rest[0]][7])
                        model.set('resnum_1', atom_ids[m_id][rest[0]][4])
                        model.set('ent_1', atom_ids[m_id][rest[0]][0])
                        model.set('altcode_1', atom_ids[m_id][rest[0]][5])
                        model.set('icode_1', atom_ids[m_id][rest[0]][6])
                        model.set('seq_1', atom_ids[m_id][rest[0]][3])
                        model.set('said_1', atom_ids[m_id][rest[0]][1])
                        model.set('resname_1', atom_ids[m_id][rest[0]][2])
                        model.set('model_2', str(m_id))
                        model.set('chain_2', atom_ids[m_id][rest[1]][7])
                        model.set('resnum_2', atom_ids[m_id][rest[1]][4])
                        model.set('ent_2', atom_ids[m_id][rest[1]][0])
                        model.set('altcode_2', atom_ids[m_id][rest[1]][5])
                        model.set('icode_2', atom_ids[m_id][rest[1]][6])
                        model.set('seq_2', atom_ids[m_id][rest[1]][3])
                        model.set('said_2', atom_ids[m_id][rest[1]][1])
                        model.set('resname_2', atom_ids[m_id][rest[1]][2])

                        # model.set('seq_1', str(rest[0][0]))
                        # model.set('chain_1', str(rest[0][1]))
                        # model.set('res_1', str(rest[0][2]))
                        # model.set('atom_1', str(rest[0][3]))
                        # model.set('seq_2', str(rest[1][0]))
                        # model.set('chain_2', str(rest[1][1]))
                        # model.set('res_2', str(rest[1][2]))
                        # model.set('atom_2', str(rest[1][3]))3
                        model.set('violation', str(distance_violations[r_id][m_id][0]))
                        model.set('rest_type', distance_violations[r_id][m_id][1])
        rest_ids = list(angle_violations.keys())
        ang_viol = ET.SubElement(violations, 'AngleViolations')
        for m_id in model_ids[:50]:
            models = ET.SubElement(ang_viol, 'Model')
            models.set('model', str(m_id))
            for r_id in rest_ids[:50]:
                for rest in angle[r_id]:
                    if angle_violations[r_id][m_id][0] > 0.0:
                        model = ET.SubElement(models, 'Violation')
                        model.set('rest_id', str(r_id[1]))
                        model.set('rest_list_id', str(r_id[0]))
                        model.set('model_1', str(m_id))
                        model.set('chain_1', atom_ids[m_id][rest[0]][7])
                        model.set('resnum_1', atom_ids[m_id][rest[0]][4])
                        model.set('ent_1', atom_ids[m_id][rest[0]][0])
                        model.set('altcode_1', atom_ids[m_id][rest[0]][5])
                        model.set('icode_1', atom_ids[m_id][rest[0]][6])
                        model.set('seq_1', atom_ids[m_id][rest[0]][3])
                        model.set('said_1', atom_ids[m_id][rest[0]][1])
                        model.set('resname_1', atom_ids[m_id][rest[0]][2])
                        model.set('model_2', str(m_id))
                        model.set('chain_2', atom_ids[m_id][rest[1]][7])
                        model.set('resnum_2', atom_ids[m_id][rest[1]][4])
                        model.set('ent_2', atom_ids[m_id][rest[1]][0])
                        model.set('altcode_2', atom_ids[m_id][rest[1]][5])
                        model.set('icode_2', atom_ids[m_id][rest[1]][6])
                        model.set('seq_2', atom_ids[m_id][rest[1]][3])
                        model.set('said_2', atom_ids[m_id][rest[1]][1])
                        model.set('resname_2', atom_ids[m_id][rest[1]][2])
                        model.set('model_3', str(m_id))
                        model.set('chain_3', atom_ids[m_id][rest[2]][7])
                        model.set('resnum_3', atom_ids[m_id][rest[2]][4])
                        model.set('ent_3', atom_ids[m_id][rest[2]][0])
                        model.set('altcode_3', atom_ids[m_id][rest[2]][5])
                        model.set('icode_3', atom_ids[m_id][rest[2]][6])
                        model.set('seq_3', atom_ids[m_id][rest[2]][3])
                        model.set('said_3', atom_ids[m_id][rest[2]][1])
                        model.set('resname_4', atom_ids[m_id][rest[2]][2])
                        model.set('model_4', str(m_id))
                        model.set('chain_4', atom_ids[m_id][rest[3]][7])
                        model.set('resnum_4', atom_ids[m_id][rest[3]][4])
                        model.set('ent_4', atom_ids[m_id][rest[3]][0])
                        model.set('altcode_4', atom_ids[m_id][rest[3]][5])
                        model.set('icode_4', atom_ids[m_id][rest[3]][6])
                        model.set('seq_4', atom_ids[m_id][rest[3]][3])
                        model.set('said_4', atom_ids[m_id][rest[3]][1])
                        model.set('resname_4', atom_ids[m_id][rest[3]][2])

                        # model.set('seq_1', str(rest[0][0]))
                        # model.set('chain_1', str(rest[0][1]))
                        # model.set('res_1', str(rest[0][2]))
                        # model.set('atom_1', str(rest[0][3]))
                        # model.set('seq_2', str(rest[1][0]))
                        # model.set('chain_2', str(rest[1][1]))
                        # model.set('res_2', str(rest[1][2]))
                        # model.set('atom_2', str(rest[1][3]))
                        # model.set('seq_3', str(rest[2][0]))
                        # model.set('chain_3', str(rest[2][1]))
                        # model.set('res_3', str(rest[2][2]))
                        # model.set('atom_3', str(rest[2][3]))
                        # model.set('seq_4', str(rest[3][0]))
                        # model.set('chain_4', str(rest[3][1]))
                        # model.set('res_4', str(rest[3][2]))
                        # model.set('atom_4', str(rest[3][3]))
                        model.set('violation', str(angle_violations[r_id][m_id][0]))
                        model.set('rest_type', angle_violations[r_id][m_id][1])
        ET.tostring(violations)
        t = ET.ElementTree(violations)
        t.write('full.xml')

    @staticmethod
    def calculate_violation_statistics(violations):
        """
        Calculates average violation value for each restraint for an ensemble
        :param violations: output from calculate_distance_violations or calculate_angle_violations
        :return: list of average violation for each restraint, list of violations
        example output [0.050748176870277231, 0.00147115297574274, 0.11772253375387987, 7, [2, 5, 7, 11, 14, 18, 19],
         'medium', ('1', '1')],[0.041170250870272262, 2, 'medium', ('1', '1')]
        """
        rest_list = list(violations.keys())
        models = list(violations[rest_list[0]].keys())
        avg_violations = {}
        viol = []
        for rest in rest_list:
            v = []
            m_id = []
            for m in models:
                cat = violations[rest][m][1]
                if violations[rest][m][0] > 0.0:
                    v.append(violations[rest][m][0])
                    viol.append([violations[rest][m][0], m, cat, rest])
                    m_id.append(m)

            if len(v) > 0:
                avg_violations[rest] = [numpy.mean(v), min(v), max(v), len(v), m_id, cat]
            else:
                avg_violations[rest] = [0.0, 0.0, 0.0, 0, m_id, cat]
        avg_viol_list = []
        for rest in rest_list:
            avg_viol_list.append(avg_violations[rest])
            avg_viol_list[-1].append(rest)
        return avg_viol_list, viol

    @staticmethod
    def bin_distance_violations(violations):
        """
        Count the number of violations in different bins (not violated)(0-0.2),(0.2-0.5),(0.5-1.0),(1.0-2.0),(2.0-5.0),(>5)
        :param violations: output from calculate_distance_violations
        :return: disctionary {model id : [count in each bin]} example {1: [1475, 19, 7, 12, 19, 7, 0],
         2: [1469, 30, 7, 9, 15, 9, 0], 3: [1473, 23, 7, 7, 20, 9, 0]...}
        """
        rest_list = list(violations.keys())
        models = list(violations[rest_list[0]].keys())
        stat = {}
        for m in models:
            c = [0, 0, 0, 0, 0, 0, 0]
            for rest in rest_list:
                if violations[rest][m][0] == 0.0:
                    c[0] += 1
                elif 0.0 < violations[rest][m][0] <= 0.2:
                    c[1] += 1
                elif 0.2 < violations[rest][m][0] <= 0.5:
                    c[2] += 1
                elif 0.5 < violations[rest][m][0] <= 1.0:
                    c[3] += 1
                elif 1.0 < violations[rest][m][0] <= 2.0:
                    c[4] += 1
                elif 2.0 < violations[rest][m][0] <= 5.0:
                    c[5] += 1
                elif 5.0 < violations[rest][m][0]:
                    c[6] += 1
                else:
                    print("Error in violation calculation")
            stat[m] = c
        return stat

    @staticmethod
    def bin_angle_violations(violations):
        """
        Count the number of violations in different bins (not violated)(0-5),(5-10),(10-20),(20-40),(40-80),(>80)
        :param violations: output from calculate_angle_violations
        :return: dictionary {model id : [count in each bin]} example {1: [1475, 19, 7, 12, 19, 7, 0],
         2: [1469, 30, 7, 9, 15, 9, 0], 3: [1473, 23, 7, 7, 20, 9, 0]...}
        """
        rest_list = list(violations.keys())
        models = list(violations[rest_list[0]].keys())
        stat = {}
        for m in models:
            c = [0, 0, 0, 0, 0, 0, 0]
            for rest in rest_list:
                if violations[rest][m][0] == 0.0:
                    c[0] += 1
                elif 0.0 < violations[rest][m][0] <= 5.0:
                    c[1] += 1
                elif 5.0 < violations[rest][m][0] <= 10.0:
                    c[2] += 1
                elif 10.0 < violations[rest][m][0] <= 20.0:
                    c[3] += 1
                elif 20.0 < violations[rest][m][0] <= 40.0:
                    c[4] += 1
                elif 40.0 < violations[rest][m][0] <= 80.0:
                    c[5] += 1
                elif 80.0 < violations[rest][m][0]:
                    c[6] += 1
                else:
                    print("Error in violation calculation")
            stat[m] = c
        return stat

if __name__ == "__main__":
    #cif  = sys.argv[1]
    #star = sys.argv[2]
    p = Validate_NMR_Restraints('nef_examples/2lci.cif', 'nef_examples/2lci.str')


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

import re
import os
import json
import numpy as np
import pandas as pd

from scipy.interpolate import interp1d

from .data import _get_connection
from .plotting import _init_plot, _draw_plot
from .element import ELEMENTS, Element

COMPOUND_LIST = list(map(str, pd.read_sql('SELECT compound FROM compounds', _get_connection('ziegler'))['compound']))

class Compound(object):
	"""Data and properties of atomic compounds

	The compound class provides the same functionality as the
	element class, but for compounds of atomic elements rather
	than the individual atomic elements.  The compound is described
	by a set of elements, a set of weights for each element, and a
	density.

	The weights can be either given as atom-weights, e.g.
	in H2O the atom weights are 0.67 for H and 0.33 for O, or as
	mass-weights, e.g. brass is 0.33 Zn by weight and 0.67 Cu by
	weight.  Some preset compounds are available; their names can be
	found by printing `ci.COMPOUND_LIST`.

	Parameters
	----------
	compound : str
		The name of the compound.  If the compound is in `ci.COMPOUND_LIST`,
		the weights and density will take preset values, if not given.
		If the name of the compound is given in chemical notation, e.g. NaCl
		or H2O2, the atom-weights can be inferred if not given explicitly.  Note
		that full chemical notation is not supported, so Ca3(PO4)2 must be written
		as Ca3P2O8.  Decimal weights are supported, e.g. C0.5O is equivalent to CO2.

	weights : dict, str or pd.DataFrame, optional
		The weights of each element in the compound.  Multiple formats are supported.
		If weights is a dict, it must be formatted as {'el1':wt1, 'el2':wt2}, where
		atom-weights are positive, and mass-weights are negative.

		If weights is a pandas DataFrame, it must contain an 'element' column,
		and one of 'weight', 'atom_weight', or 'mass_weight'.  If 'weight' is the column
		given, the convention of positive atom-weights and negative mass-weights is followed.

		If weights is a str, it can either be formatted as 'el1:wt1, el2:wt2', or it can be
		a path to a .csv, .json or .db file.  These files must contain the same information
		as the DataFrame option, and can contain weights for multiple compounds, if 'compound'
		is one of the columns/keys.

		If a .json file, it must follow the 'records' formatting convention (see pandas docs).

	density : float, optional
		Density of the compound in g/cm^3.  The density is required for the cm.attenuation(),
		cm.S(), cm.range() and cm.plot_range() functions, but is an optional argument in each
		of those functions if not provided at construction.  Can also be specified by using
		a 'density' column/key in the file/DataFrame for weights.

	Attributes
	----------
	name : str
		The name of the compound.

	weights : pd.DataFrame
		The weights for each element in the compound.  DataFrame columns are
		'element', 'Z', 'mass_weight', 'atom_weight'.

	density : float
		Density of the compound in g/cm^3. The density is used in
		calculations of charged particle dEdx and photon attenuation, so
		if the density was not explicitly given at construction,
		you can assign a new density using `cm.density = new_density` if
		needed, or using the `density` keyword in either of those functions.

	elements : list of ci.Element
		Elements in the compound.

	mass_coeff : pd.DataFrame
		Table of mass-attenuation coefficients as a function of photon
		energy, from the NIST XCOM database.  Energies are in keV, and
		mass-attenuation coefficients, or mu/rho, are given in cm^2/g.
		DataFrame columns are 'energy', 'mu' and 'mu_en' for the
		mass-energy absorption coefficient.

	Examples
	--------
	>>> print('Silicone' in ci.COMPOUND_LIST)
	True
	>>> cm = ci.Compound('Silicone') # preset compound
	>>> print(list(map(str, cm.elements)))
	['H', 'C', 'O', 'Si']
	>>> cm = ci.Compound('H2O', density=1.0)
	print(cm.weights)
	  element  Z  atom_weight  mass_weight
	0       H  1     0.666667     0.111907
	1       O  8     0.333333     0.888093
	>>> cm = ci.Compound('Brass', weights={'Zn':-33,'Cu':-66})
	>>> print(cm.weights)
	  element   Z  atom_weight  mass_weight
	0      Zn  30     0.327041     0.333333
	1      Cu  29     0.672959     0.666667
	>>> cm.saveas('brass.csv')

	"""

	def __init__(self, compound, weights=None, density=None):
		self.name = compound
		self.density = None

		if compound in COMPOUND_LIST:
			df = pd.read_sql('SELECT * FROM compounds WHERE compound="{}"'.format(compound), _get_connection('ziegler'))
			self.density = df['density'][0]

			if weights is None:
				wts = df['weights'][0].split(',')
				elements = [str(i.split(':')[0]) for i in wts]
				atom_weights = np.array([float(i.split(':')[1]) for i in wts])
				self._set_weights(elements, atom_weights=atom_weights)

		elif weights is None:
			elements = []
			for el_gp in [i for i in re.split('[0-9]+|\\.', compound) if i]:
				for s in [i for i in re.split('([A-Z])', el_gp) if i]:
					if s.upper()==s:
						elements.append(s)
					else:
						elements[-1] += s

			if all([e in ELEMENTS for e in elements]):
				wts = re.split('|'.join(sorted(elements, key=lambda i:-len(i))), compound)
				atom_weights = np.array([float(wts[n+1]) if wts[n+1] else 1.0 for n,e in enumerate(elements)])
				self._set_weights(elements, atom_weights=atom_weights)

		if weights is not None:

			if type(weights)==dict:
				elements = [e for e in weights]
				wts = np.array([weights[e] for e in elements], dtype=np.float64)
				weights = pd.DataFrame({'element':elements, 'weight':wts})

			elif type(weights)==str:
				if weights.endswith('.json'):
					weights = pd.read_json(weights, orient='records').fillna(method='ffill')
					weights.columns = map(str.lower, map(str, weights.columns))
					if 'compound' in weights.columns:
						weights = weights[weights['compound']==self.name]

				elif weights.endswith('.csv'):
					weights = pd.read_csv(weights, header=0).fillna(method='ffill')
					weights.columns = map(str.lower, map(str, weights.columns))
					if 'compound' in weights.columns:
						weights = weights[weights['compound']==self.name]

				elif weights.endswith('.db'):
					weights = pd.read_sql('SELECT * FROM compounds WHERE compound={}'.format(self.name), _get_connection(weights))
					weights.columns = map(str.lower, map(str, weights.columns))
					if 'compound' in weights.columns:
						weights = weights[weights['compound']==self.name]

				else:
					elements = [str(i.split(':')[0]).strip() for i in weights.split(',')]
					wts = np.array([float(i.split(':')[1].strip()) for i in weights.split(',')])

					if wts[0]>0:
						self._set_weights(elements, atom_weights=wts)
					else:
						self._set_weights(elements, mass_weights=np.abs(wts))

			if type(weights)==pd.DataFrame:
				weights.columns = map(str.lower, map(str, weights.columns))
				if 'density' in weights.columns:
					self.density = weights['density'].iloc[0]

				cols = ['element', 'Z', 'atom_weight', 'mass_weight']
				if all([i in weights.columns for i in cols]):
					self.weights = weights[cols]

				elif 'atom_weight' in weights.columns:
					self._set_weights(list(weights['element']), atom_weights=weights['atom_weight'].to_numpy())

				elif 'mass_weight' in weights.columns:
					self._set_weights(list(weights['element']), mass_weights=weights['mass_weight'].to_numpy())

				else:
					elements, wts = list(weights['element']), weights['weight'].to_numpy()

					if wts[0]>0:
						self._set_weights(elements, atom_weights=wts)
					else:
						self._set_weights(elements, mass_weights=np.abs(wts))

		if density is not None:
			self.density = density

		self.elements = [Element(el) for el in self.weights['element']]

		if self.density is None and len(self.weights)==1:
			self.density = self.elements[0].density

		E = np.unique(np.concatenate([el.mass_coeff['energy'].to_numpy() for el in self.elements]))
		mu = np.average([el.mu(E) for el in self.elements], weights=self.weights['mass_weight'], axis=0)
		mu_en = np.average([el.mu_en(E) for el in self.elements], weights=self.weights['mass_weight'], axis=0)
		self.mass_coeff = pd.DataFrame({'energy':E,'mu':mu,'mu_en':mu_en})
		self._mc_interp, self._mc_en_interp = None, None

	def _set_weights(self, elements, atom_weights=None, mass_weights=None):
		amu = pd.read_sql('SELECT * FROM weights', _get_connection('ziegler'))
		Zs = [ELEMENTS.index(el) for el in elements]

		if mass_weights is None:
			mass_weights = np.array([amu[amu['Z']==z]['amu'][z-1]*atom_weights[n] for n,z in enumerate(Zs)])
		elif atom_weights is None:
			atom_weights = np.array([mass_weights[n]/amu[amu['Z']==z]['amu'][z-1] for n,z in enumerate(Zs)])

		atom_weights, mass_weights = atom_weights/np.sum(atom_weights), mass_weights/np.sum(mass_weights)
		self.weights = pd.DataFrame({'element':elements, 'Z':Zs, 'atom_weight':atom_weights, 'mass_weight':mass_weights},
								columns=['element','Z','atom_weight','mass_weight'])

	def __str__(self):
		return self.name

	def saveas(self, filename, replace=False):
		"""Save the compound definition to a file

		The weights and density of the compound can be saved to one of
		the following file formats: .csv, .json, .db.  If the file exists,
		the data will be appended, unless `replace=True`, in which case
		the file will be replaced.  If a definition for the compound exists
		in the file already, it will be replaced.

		Parameters
		----------
		filename : str
			Filename where the compound will be saved.  Available formats
			are .csv, .json and .db.

		replace : bool, optional
			If `True`, replace the file if it exists.  Default `False`, which
			appends the data to the file.

		Examples
		--------
		>>> cm = ci.Compound('Brass', weights={'Zn':-33,'Cu':-66})
		>>> cm.saveas('brass.csv')
		>>> cm = ci.Compound('Water', weights={'H':2, 'O':1}, density=1.0)
		>>> cm.saveas('water.json')

		"""
		wts = self.weights.copy()
		wts['compound'] = self.name
		if self.density is not None:
			wts['density'] = self.density

		if filename.endswith('.csv'):
			if os.path.exists(filename) and not replace:
				df = pd.read_csv(filename, header=0)
				df = df[df['compound']!=self.name]
				df = pd.concat([df, wts])
				df.to_csv(filename, index=False)
			else:
				wts.to_csv(filename, index=False)

		if filename.endswith('.db'):
			if os.path.exists(filename) and not replace:
				con = _get_connection(filename)
				df = pd.read_sql('SELECT * FROM weights', con)
				df = df[df['compound']!=self.name]
				df = pd.concat([df, wts])
				df.to_sql('weights', con, if_exists='replace', index=False)
			else:
				wts.to_sql('weights', _get_connection(filename), if_exists='replace', index=False)

		if filename.endswith('.json'):
			if os.path.exists(filename) and not replace:
				df = pd.read_json(filename, orient='records')
				df = df[df['compound']!=self.name][wts.columns]
				df = pd.concat([df, wts])
				json.dump(json.loads(df.to_json(orient='records')), open(filename, 'w'), indent=4)
			else:
				json.dump(json.loads(wts.to_json(orient='records')), open(filename, 'w'), indent=4)

	def mu(self, energy):
		"""Mass-attenuation coefficient

		Interpolates the mass-attenuation coefficient, mu/rho,
		for the compound along the input energy grid.

		Parameters
		----------
		energy : array_like
			The incident photon energy, in keV.

		Returns
		-------
		mu : np.ndarray
			Mass attenuation coefficient, mu/rho, in cm^2/g.

		Examples
		--------
		>>> cm = ci.Compound('H2O')
		>>> print(cm.mu(200))
		0.13703928393005832

		"""

		if self._mc_interp is None:
			self._mc_interp = interp1d(np.log(self.mass_coeff['energy']), np.log(self.mass_coeff['mu']), bounds_error=False, fill_value='extrapolate')
		return np.exp(self._mc_interp(np.log(energy)))

	def mu_en(self, energy):
		"""Mass energy-absorption coefficient

		Interpolates the mass-energy absorption coefficient, mu_en/rho,
		for the compound along the input energy grid.

		Parameters
		----------
		energy : array_like
			The incident photon energy, in keV.

		Returns
		-------
		mu_en : np.ndarray
			Mass energy absorption coefficient, mu_en/rho, in cm^2/g.

		Examples
		--------
		>>> cm = ci.Compound('H2O')
		>>> print(cm.mu_en(200))
		0.029671598667776862

		"""

		if self._mc_en_interp is None:
			self._mc_en_interp = interp1d(np.log(self.mass_coeff['energy']), np.log(self.mass_coeff['mu_en']), bounds_error=False, fill_value='extrapolate')
		return np.exp(self._mc_en_interp(np.log(energy)))

	def attenuation(self, energy, x, density=None):
		"""Photon attenuation in matter

		Calculate the attenuation factor I(x)/I_0 = e^(-mu*x) for a given
		photon energy (in keV) and slab thickness (in cm).

		Parameters
		----------
		energy : array_like
			Incident photon energy in keV.

		x : float
			Thickness of slab of given compound, in cm.

		density : float, optional
			Density of the compound in g/cm^3.  Default behavior is to
			use `Compound.density`, which must be supplied at construction.

		Returns
		-------
		attenuation : numpy.ndarray
			The slab attenuation factor as an absolute number (i.e. from 0 to 1).
			E.g. if the incident intensity is I_0, the transmitted intensity I(x)
			is I_0 times the attenuation factor.

		Examples
		--------
		>>> cm = ci.Compound('SS_316') # preset compound for 316 Stainless
		>>> print(cm.attenuation(511, x=0.3))
		0.8199829388434694
		>>> print(cm.attenuation(300, x=1.0, density=5.0))
		0.5752140388004373

		"""

		energy = np.asarray(energy, dtype=np.float64)
		x = np.asarray(x, dtype=np.float64)
		if density is None:
			density = self.density

		return np.exp(-self.mu(energy)*x*density)

	def S(self, energy, particle='p', density=None):
		"""Charged particle stopping power in matter

		Calculate the stopping power, S=-dE/dx, for a given ion as a
		function of the ion energy in MeV.  Units of S are MeV/cm.  To return
		stopping power in units of MeV/(mg/cm^2), use option `density=1E-3`.
		The stopping power is calculated using the Element.S() methods for
		each element in cm.elements, added using Bragg's rule.

		Parameters
		----------
		energy : array_like
			Incident ion energy in MeV.

		particle : str, optional
			Incident ion.  For light ions, options are 'p' (default), 'd', 't', 'a' for proton,
			deuteron, triton and alpha, respectively.  Additionally, heavy ions can be
			specified either by element or isotope, e.g. 'Fe', '40CA', 'U', 'Bi-209'. For
			light ions, the charge state is assumed to be fully stripped. For heavy ions
			the charge state is handled by a Bohr/Northcliffe parameterization consistent
			with the Anderson-Ziegler formalism.

		density : float, optional
			Density of the compound in g/cm^3.  Default behavior is to use
			`Compound.density`.  To return stopping power in units of MeV/(mg/cm^2), i.e.
			the mass-stopping power, use `density=1E-3`.

		Returns
		-------
		stopping_power : numpy.ndarray
			Stopping power, S=-dE/dx, for a given ion as a function of the
			ion energy in MeV.  Units of S are MeV/cm.

		Examples
		--------
		>>> cm = ci.Compound('SrCO3', density=3.5)
		>>> print(cm.S(60.0))
		27.196387031247834
		>>> print(cm.S(55.0, density=1E-3)) ### S in MeV/(mg/cm^2)
		0.008307827781861116

		"""

		energy = np.asarray(energy, dtype=np.float64)
		if density is None:
			density = self.density

		return np.average([el.S(energy, particle=particle, density=1E-3) for el in self.elements], weights=self.weights['mass_weight'], axis=0)*1E3*density

	def range(self, energy, particle='p', density=None):
		"""Charged particle range in matter

		Calculates the charged particle range in the compound, in cm.  Incident
		energy should be in MeV, and the particle type definition is identical
		to `Compound.S()`.

		Parameters
		----------
		energy : array_like
			Incident ion energy in MeV.

		particle : str, optional
			Incident ion.  For light ions, options are 'p' (default), 'd', 't', 'a' for proton,
			deuteron, triton and alpha, respectively.  Additionally, heavy ions can be
			specified either by element or isotope, e.g. 'Fe', '40CA', 'U', 'Bi-209'. For
			light ions, the charge state is assumed to be fully stripped. For heavy ions
			the charge state is handled by a Bohr/Northcliffe parameterization consistent
			with the Anderson-Ziegler formalism.

		density : float, optional
			Density of the compound in g/cm^3.  Default behavior is to use
			`Compound.density`, which must be supplied at construction.

		Returns
		-------
		range : np.ndarray
			Charged particle range in the compound, in cm.

		Examples
		--------
		>>> cm = ci.Compound('Fe') # same behavior as element
		>>> print(cm.range(60.0))
		0.5858151125192633
		>>> cm = ci.Compound('SS_316') # preset compound
		>>> print(cm.range(60.0))
		0.5799450918147814

		"""

		energy = np.asarray(energy, dtype=np.float64)

		dE = np.max(energy)/1E3
		E_min = min((np.min(energy), 1.0))
		E_grid = np.arange(E_min, np.max(energy)+dE, dE)

		S = self.S(E_grid, particle=particle, density=density)
		x = np.cumsum((1.0/S)*dE)
		return interp1d(np.log(E_grid), x, bounds_error=None, fill_value='extrapolate')(np.log(energy))

	def plot_mass_coeff(self, energy=None, **kwargs):
		"""Plot the mass-attenuation coefficient in the compound

		Creates a plot of the mass-attenuation coefficient (in cm^2/g)
		as a function of photon energy in keV.

		Parameters
		----------
		energy : array_like, optional
			Energy grid on which to plot, replacing the default energy grid.
			Units are in keV.

		Other Parameters
		----------------
		**kwargs
			Optional keyword arguments for plotting.  See the
			plotting section of the curie API for a complete
			list of kwargs.

		Examples
		--------
		>>> cm = ci.Compound('Fe')
		>>> cm.plot_mass_coeff()
		>>> cm = ci.Compound('H2O')
		>>> cm.plot_mass_coeff(style='poster')

		"""

		if energy is None:
			energy, mu = self.mass_coeff['energy'], self.mass_coeff['mu']

		else:
			energy = np.asarray(energy, dtype=np.float64)
			mu = self.mu(energy)

		f,ax = _init_plot(**kwargs)

		ax.plot(energy, mu, label=r'$\mu/\rho$'+' ({})'.format(self.name))
		ax.set_xlabel('Photon Energy (keV)')
		ax.set_ylabel(r'Attenuation Coeff. (cm$^2$/g)')
		ax.set_xscale('log')
		ax.set_yscale('log')
		ax.legend()

		return _draw_plot(f, ax, **kwargs)

	def plot_mass_coeff_en(self, energy=None, **kwargs):
		"""Plot the mass energy-absorption coefficient in the compound

		Creates a plot of the mass energy-absorption coefficient (in cm^2/g)
		as a function of photon energy in keV.

		Parameters
		----------
		energy : array_like, optional
			Energy grid on which to plot, replacing the default energy grid.
			Units are in keV.

		Other Parameters
		----------------
		**kwargs
			Optional keyword arguments for plotting.  See the
			plotting section of the curie API for a complete
			list of kwargs.

		Examples
		--------
		>>> cm = ci.Compound('Silicone') # preset compound

		Example plotting the mass-attenuation coefficient together with the mass
		energy-absorption coefficient, on the same axes.

		>>> f,ax = cm.plot_mass_coeff(return_plot=True)
		>>> cm.plot_mass_coeff_en(f=f, ax=ax)

		"""

		if energy is None:
			energy, mu = self.mass_coeff['energy'], self.mass_coeff['mu_en']

		else:
			energy = np.asarray(energy, dtype=np.float64)
			mu = self.mu_en(energy)

		f,ax = _init_plot(**kwargs)

		ax.plot(energy, mu, label=r'$\mu_{en}/\rho$'+' ({})'.format(self.name))
		ax.set_xlabel('Photon Energy (keV)')
		ax.set_ylabel(r'Attenuation Coeff. (cm$^2$/g)')
		ax.set_xscale('log')
		ax.set_yscale('log')
		ax.legend()

		return _draw_plot(f, ax, **kwargs)

	def plot_S(self, particle='p', energy=None, **kwargs):
		"""Plot the stopping power in the compound

		Creates a plot of the charged particle stopping power (in MeV/(mg/cm^2))
		in the compound as a function of the incident ion energy (in MeV).

		Parameters
		----------
		particle : str
			Incident ion.  For light ions, options are 'p' (default), 'd', 't', 'a' for proton,
			deuteron, triton and alpha, respectively.  Additionally, heavy ions can be
			specified either by element or isotope, e.g. 'Fe', '40CA', 'U', 'Bi-209'.

		energy : array_like, optional
			Energy grid on which to plot, replacing the default energy grid.
			Units are in MeV.

		Other Parameters
		----------------
		**kwargs
			Optional keyword arguments for plotting.  See the
			plotting section of the curie API for a complete
			list of kwargs.

		Examples
		--------
		>>> cm = ci.Compound('He') # same as element
		>>> cm.plot_S(particle='a')
		>>> cm = ci.Compound('Kapton')
		>>> cm.plot_S(particle='d')

		"""

		if energy is None:
			energy = 10.0**np.arange(-1.5, 2.8, 0.05)

		f,ax = _init_plot(**kwargs)
		ax.plot(energy, self.S(energy, particle=particle, density=1E-3), label=r'$-\frac{dE}{dx}$ ('+self.name+')')

		ax.set_xlabel('Incident Energy (MeV)')
		ax.set_ylabel(r'Stopping Power (MeV/(mg/cm$^2$))')
		ax.set_xscale('log')
		ax.legend()

		return _draw_plot(f, ax, **kwargs)

	def plot_range(self, particle='p', energy=None, density=None, **kwargs):
		"""Plot the charged particle range in the compound

		Creates a plot of the charged particle range (in cm)
		in the compound as a function of the incident ion energy (in MeV).

		Parameters
		----------
		particle : str
			Incident ion.  For light ions, options are 'p' (default), 'd', 't', 'a' for proton,
			deuteron, triton and alpha, respectively.  Additionally, heavy ions can be
			specified either by element or isotope, e.g. 'Fe', '40CA', 'U', 'Bi-209'.

		energy : array_like, optional
			Energy grid on which to plot, replacing the default energy grid.
			Units are in MeV.

		density : float, optional
			Density of the compound in g/cm^3.  Default behavior is to use
			`Compound.density`.

		Other Parameters
		----------------
		**kwargs
			Optional keyword arguments for plotting.  See the
			plotting section of the curie API for a complete
			list of kwargs.

		Examples
		--------
		>>> cm = ci.Compound('Bronze', weights={'Cu':-80, 'Sn':-20}, density=8.9)
		>>> f,ax = cm.plot_range(return_plot=True)
		>>> cm.plot_range(particle='d', f=f, ax=ax)

		"""

		if energy is None:
			energy = 10.0**np.arange(-1.5, 2.8, 0.05)

		f,ax = _init_plot(**kwargs)
		ax.plot(energy, self.range(energy, particle=particle, density=density), label='Range ({})'.format(self.name))

		ax.set_xlabel('Incident Energy (MeV)')
		ax.set_ylabel('Range (cm)')
		ax.set_xscale('log')
		ax.set_yscale('log')
		ax.legend()

		return _draw_plot(f, ax, **kwargs)