#!/usr/bin/env python
#
# Author: Sandeep Sharma <sanshar@gmail.com>
#         Sheng Guo <shengg@princeton.edu>
#         Qiming Sun <osirpt.sun@gmail.com>
#

'''
DMRG solver for CASCI and CASSCF.
'''

import ctypes
import os
import sys
import struct
import time
import tempfile
from subprocess import check_call, check_output, STDOUT, CalledProcessError
import numpy
import pyscf.tools
import pyscf.lib
from pyscf.lib import logger
from pyscf.lib import chkfile
from pyscf import mcscf
from pyscf.dmrgscf import dmrg_sym

import pyscf.lib

libunpack = pyscf.lib.load_library('libicmpspt')

try:
    from pyscf.dmrgscf import settings
except ImportError:
    import sys
    sys.stderr.write('''settings.py not found.  Please create %s
''' % os.path.join(os.path.dirname(__file__), 'settings.py'))
    raise ImportError

class DMRGCI(pyscf.lib.StreamObject):
    '''Block program interface and the object to hold Block program input parameters.

    Attributes:
        outputlevel : int
            Noise level for Block program output.
        maxIter : int

        approx_maxIter : int
            To control the DMRG-CASSCF approximate DMRG solver accuracy.
        twodot_to_onedot : int
            When to switch from two-dot algroithm to one-dot algroithm.
        nroots : int

        weights : list of floats
            Use this attribute with "nroots" attribute to set state-average calculation.
        restart : bool
            To control whether to restart a DMRG calculation.
        tol : float
            DMRG convergence tolerence
        maxM : int
            Bond dimension
        scheduleSweeps, scheduleMaxMs, scheduleTols, scheduleNoises : list
            DMRG sweep scheduler.  See also Block documentation
        wfnsym : str or int
            Wave function irrep label or irrep ID
        orbsym : list of int
            irrep IDs of each orbital
        groupname : str
            groupname, orbsym together can control whether to employ symmetry in
            the calculation.  "groupname = None and orbsym = []" requires the
            Block program using C1 symmetry.

    Examples:

    >>> mol = gto.M(atom='C 0 0 0; C 0 0 1')
    >>> mf = scf.RHF(mol).run()
    >>> mc = mcscf.CASCI(mf, 4, 4)
    >>> mc.fcisolver = DMRGCI(mol)
    >>> mc.kernel()
    -74.379770619390698
    '''
    def __init__(self, mol, maxM=None, tol=None, num_thrds=1, memory=None):
        self.mol = mol
        self.verbose = mol.verbose
        self.stdout = mol.stdout
        self.outputlevel = 2

        self.executable = settings.BLOCKEXE
        self.scratchDirectory = os.path.abspath(settings.BLOCKSCRATCHDIR)
        self.mpiprefix = settings.MPIPREFIX
        self.memory = memory

        self.integralFile = "FCIDUMP"
        self.configFile = "dmrg.conf"
        self.outputFile = "dmrg.out"
        if hasattr(settings, 'BLOCKRUNTIMEDIR'):
            self.runtimeDir = settings.BLOCKRUNTIMEDIR
        else:
            self.runtimeDir = '.'
        self.maxIter = 20
        self.approx_maxIter = 4
        self.twodot_to_onedot = 15
        self.dmrg_switch_tol = 1e-3
        self.nroots = 1
        self.weights = []
        self.wfnsym = 1
        self.extraline = []

        if tol is None:
            self.tol = 1e-8
        else:
            self.tol = tol/10
        if maxM is None:
            self.maxM = 1000
        else:
            self.maxM = maxM

        self.num_thrds= num_thrds
        self.startM =  None
        self.restart = False
        self.nonspinAdapted = False
        self.scheduleSweeps = []
        self.scheduleMaxMs  = []
        self.scheduleTols   = []
        self.scheduleNoises = []
        self.onlywriteIntegral = False
        self.spin = 0

        self.orbsym = []
        if mol.symmetry:
            self.groupname = mol.groupname
        else:
            self.groupname = None

##################################################
# don't modify the following attributes, if you do not finish part of calculation, which can be reused.

        #DO NOT CHANGE these parameters, unless you know the code in details
        self.twopdm = True #By default, 2rdm is calculated after the calculations of wave function.
        self.block_extra_keyword = [] #For Block advanced user only.
        self.has_fourpdm = False
        self.has_threepdm = False
        self.has_nevpt = False
# This flag _restart is set by the program internally, to control when to make
# Block restart calculation.
        self._restart = False
        self.generate_schedule()

        self._keys = set(self.__dict__.keys())

    @property
    def threads(self):
        return self.num_thrds
    @threads.setter
    def threads(self, x):
        self.num_thrds = x

    def generate_schedule(self):

        if self.startM is None:
            if self.maxM < 200:
                self.startM = 50
            else:
                self.startM = 200
        if len(self.scheduleSweeps) == 0:
            startM = self.startM
            N_sweep = 0
            if self.restart or self._restart :
                Tol = self.tol / 10.0
            else:
                Tol = 1.0e-5
            Noise = Tol
            while startM < int(self.maxM):
                self.scheduleSweeps.append(N_sweep)
                N_sweep += 4
                self.scheduleMaxMs.append(startM)
                startM *= 2
                self.scheduleTols.append(Tol)
                self.scheduleNoises.append(Noise)
            while Tol > float(self.tol):
                self.scheduleSweeps.append(N_sweep)
                N_sweep += 2
                self.scheduleMaxMs.append(self.maxM)
                self.scheduleTols.append(Tol)
                Tol /= 10.0
                self.scheduleNoises.append(5.0e-5)
            self.scheduleSweeps.append(N_sweep)
            N_sweep += 2
            self.scheduleMaxMs.append(self.maxM)
            self.scheduleTols.append(self.tol)
            self.scheduleNoises.append(0.0)
            self.twodot_to_onedot = N_sweep + 2
            self.maxIter = self.twodot_to_onedot + 12
        return self

    def dump_flags(self, verbose=None):
        if verbose is None:
            verbose = self.verbose
        log = logger.Logger(self.stdout, verbose)
        log.info('******** Block flags ********')
        log.info('executable = %s', self.executable)
        log.info('Block version %s', block_version(self.executable))
        log.info('BLOCKEXE_COMPRESS_NEVPT = %s', settings.BLOCKEXE_COMPRESS_NEVPT)
        log.info('mpiprefix = %s', self.mpiprefix)
        log.info('scratchDirectory = %s', self.scratchDirectory)
        log.info('integralFile = %s', os.path.join(self.runtimeDir, self.integralFile))
        log.info('configFile = %s', os.path.join(self.runtimeDir, self.configFile))
        log.info('outputFile = %s', os.path.join(self.runtimeDir, self.outputFile))
        log.info('maxIter = %d', self.maxIter)
        log.info('scheduleSweeps = %s', str(self.scheduleSweeps))
        log.info('scheduleMaxMs = %s', str(self.scheduleMaxMs))
        log.info('scheduleTols = %s', str(self.scheduleTols))
        log.info('scheduleNoises = %s', str(self.scheduleNoises))
        log.info('twodot_to_onedot = %d', self.twodot_to_onedot)
        log.info('tol = %g', self.tol)
        log.info('maxM = %d', self.maxM)
        log.info('fullrestart = %s', str(self.restart or self._restart))
        log.info('dmrg switch tol =%s', self.dmrg_switch_tol)
        log.info('wfnsym = %s', self.wfnsym)
        log.info('num_thrds = %d', self.num_thrds)
        log.info('memory = %s', self.memory)
        return self

    # ABOUT RDMs AND INDEXES: -----------------------------------------------------------------------
    #   There is two ways to stored an RDM
    #   (the numbers help keep track of creation/annihilation that go together):
    #     E3[i1,j2,k3,l3,m2,n1] is the way BLOCK and STACKBLOCK outputs text and bin files
    #     E3[i1,j2,k3,l1,m2,n3] is the way the tensors need to be written for SQA and ICPT
    #
    #   --> See various remarks in the pertinent functions below.
    # -----------------------------------------------------------------------------------------------

    def make_rdm1s(self, state, norb, nelec, link_index=None, **kwargs):
        # Ref: IJQC, 109, 3552 Eq (3)
        if isinstance(nelec, (int, numpy.integer)):
            nelecb = (nelec-self.spin) // 2
            neleca = nelec - nelecb
        else :
            neleca, nelecb = nelec
        dm1, dm2 = DMRGCI.make_rdm12(self, state, norb, nelec, link_index, **kwargs)
        dm1n = (2-(neleca+nelecb)/2.) * dm1 - numpy.einsum('pkkq->pq', dm2)
        dm1n *= 1./(neleca-nelecb+1)
        dm1a, dm1b = (dm1+dm1n)*.5, (dm1-dm1n)*.5
        return dm1a, dm1b

    def make_rdm1(self, state, norb, nelec, link_index=None, **kwargs):
        # Avoid calling self.make_rdm12 because it may be overloaded
        return DMRGCI.make_rdm12(self, state, norb, nelec, link_index, **kwargs)[0]

    def make_rdm12(self, state, norb, nelec, link_index=None, **kwargs):
        nelectrons = 0
        if isinstance(nelec, (int, numpy.integer)):
            nelectrons = nelec
        else:
            nelectrons = nelec[0]+nelec[1]

        twopdm = numpy.zeros( (norb, norb, norb, norb) )
        file2pdm = "spatial_twopdm.%d.%d.txt" %(state, state)
        # The 2RDMs written by "save_spatial_twopdm_text" in BLOCK and STACKBLOCK
        # are written as E2[i1,j2,k2,l1] (right?)
        # and stored here as E2[i1,l1,j2,k2] (weird?)
        # This is NOT done with SQA in mind.
        with open(os.path.join(self.scratchDirectory, "node0", file2pdm), "r") as f:
            norb_read = int(f.readline().split()[0])
            assert(norb_read == norb)

            for line in f:
                linesp = line.split()
                i, k, l, j = [int(x) for x in linesp[:4]]
                twopdm[i,j,k,l] = 2.0 * float(linesp[4])

        # (this is coherent with previous statement about indexes) (right?)
        onepdm = numpy.einsum('ikjj->ik', twopdm)
        onepdm /= (nelectrons-1)
        return onepdm, twopdm

    def trans_rdm1s(self, statebra, stateket, norb, nelec, link_index=None, **kwargs):
        # Ref: IJQC, 109, 3552 Eq (3)
        if isinstance(nelec, (int, numpy.integer)):
            nelecb = (nelec-self.spin) // 2
            neleca = nelec - nelecb
        else :
            neleca, nelecb = nelec
        dm1, dm2 = DMRGCI.trans_rdm12(self, statebra, stateket, norb, nelec, link_index, **kwargs)
        dm1n = (2-(neleca+nelecb)/2.) * dm1 - numpy.einsum('pkkq->pq', dm2)
        dm1n *= 1./(neleca-nelecb+1)
        dm1a, dm1b = (dm1+dm1n)*.5, (dm1-dm1n)*.5
        return dm1a, dm1b

    def trans_rdm1(self, statebra, stateket, norb, nelec, link_index=None, **kwargs):
        return DMRGCI.trans_rdm12(self, statebra, stateket, norb, nelec, link_index, **kwargs)[0]

    def trans_rdm12(self, statebra, stateket, norb, nelec, link_index=None, **kwargs):
        nelectrons = 0
        if isinstance(nelec, (int, numpy.integer)):
            nelectrons = nelec
        else:
            nelectrons = nelec[0]+nelec[1]

        writeDMRGConfFile(self, nelec, True, with_2pdm=False,
                          extraline=['restart_tran_twopdm',
                                     'specificpdm %d %d' % (statebra, stateket)])
        executeBLOCK(self)

        twopdm = numpy.zeros( (norb, norb, norb, norb) )
        file2pdm = "spatial_twopdm.%d.%d.txt" %(statebra, stateket)
        # The 2RDMs written by "save_spatial_twopdm_text" in BLOCK and STACKBLOCK
        # are written as E2[i1,j2,k2,l1] (right?)
        # and stored here as E2[i1,l1,j2,k2] (weird?)
        # This is NOT done with SQA in mind.
        with open(os.path.join(self.scratchDirectory, "node0", file2pdm), "r") as f:
            norb_read = int(f.readline().split()[0])
            assert(norb_read == norb)

            for line in f:
                linesp = line.split()
                i, k, l, j = [int(x) for x in linesp[:4]]
                twopdm[i,j,k,l] = 2.0 * float(linesp[4])

        # (this is coherent with previous statement about indexes) (right?)
        onepdm = numpy.einsum('ikjj->ik', twopdm)
        onepdm /= (nelectrons-1)
        return onepdm, twopdm

    def make_rdm123(self, state, norb, nelec, link_index=None, **kwargs):
        if self.has_threepdm == False:
            writeDMRGConfFile(self, nelec, True,
                              with_2pdm=False, extraline=['restart_threepdm'])
            if self.verbose >= logger.DEBUG1:
                inFile = os.path.join(self.runtimeDir, self.configFile)
                logger.debug1(self, 'Block Input conf')
                logger.debug1(self, open(inFile, 'r').read())
            executeBLOCK(self)
            if self.verbose >= logger.DEBUG1:
                outFile = os.path.join(self.runtimeDir, self.outputFile)
                logger.debug1(self, open(outFile).read())
            self.has_threepdm = True

        nelectrons = 0
        if isinstance(nelec, (int, numpy.integer)):
            nelectrons = nelec
        else:
            nelectrons = nelec[0]+nelec[1]

        threepdm = numpy.zeros( (norb, norb, norb, norb, norb, norb) )
        file3pdm = "spatial_threepdm.%d.%d.txt" %(state, state)
        # The 3RDMs written by "Threepdm_container::save_spatial_npdm_text" in BLOCK and STACKBLOCK
        # are written as E3[i1,j2,k3,l3,m2,n1]
        # and are also stored here as E3[i1,j2,k3,l3,m2,n1]
        # This is NOT done with SQA in mind.
        with open(os.path.join(self.scratchDirectory, "node0", file3pdm), "r") as f:
            norb_read = int(f.readline().split()[0])
            assert(norb_read == norb)

            for line in f:
                linesp = line.split()
                i, j, k, l, m, n = [int(x) for x in linesp[:6]]
                threepdm[i,j,k,l,m,n] = float(linesp[6])

        # (this is coherent with previous statement about indexes)
        twopdm = numpy.einsum('ijkklm->ijlm',threepdm)
        twopdm /= (nelectrons-2)
        onepdm = numpy.einsum('ijjk->ik', twopdm)
        onepdm /= (nelectrons-1)
        return onepdm, twopdm, threepdm

    def _make_dm123(self, state, norb, nelec, link_index=None, **kwargs):
        r'''Note this function does NOT compute the standard density matrix.
        The density matrices are reordered to match the the fci.rdm.make_dm123
        function (used by NEVPT code).
        The returned "2pdm" is :math:`\langle p^\dagger q r^\dagger s\rangle`;
        The returned "3pdm" is :math:`\langle p^\dagger q r^\dagger s t^\dagger u\rangle`.
        '''
        onepdm, twopdm, threepdm = self.make_rdm123(state, norb, nelec, None, **kwargs)
        threepdm = numpy.einsum('mkijln->ijklmn',threepdm).copy()
        threepdm += numpy.einsum('jk,lm,in->ijklmn',numpy.identity(norb),numpy.identity(norb),onepdm)
        threepdm += numpy.einsum('jk,miln->ijklmn',numpy.identity(norb),twopdm)
        threepdm += numpy.einsum('lm,kijn->ijklmn',numpy.identity(norb),twopdm)
        threepdm += numpy.einsum('jm,kinl->ijklmn',numpy.identity(norb),twopdm)

        twopdm =(numpy.einsum('iklj->ijkl',twopdm)
               + numpy.einsum('il,jk->ijkl',onepdm,numpy.identity(norb)))

        return onepdm, twopdm, threepdm

    def make_rdm3(self, state, norb, nelec, dt=numpy.dtype('Float64'), filetype = "binary", link_index=None, **kwargs):
        import os

        if self.has_threepdm == False:
            self.twopdm = False
            self.extraline.append('threepdm\n')

            writeDMRGConfFile(self, nelec, False)
            if self.verbose >= logger.DEBUG1:
                inFile = self.configFile
                #inFile = os.path.join(self.scratchDirectory,self.configFile)
                logger.debug1(self, 'Block Input conf')
                logger.debug1(self, open(inFile, 'r').read())
            executeBLOCK(self)
            if self.verbose >= logger.DEBUG1:
                outFile = self.outputFile
                #outFile = os.path.join(self.scratchDirectory,self.outputFile)
                logger.debug1(self, open(outFile).read())
            self.has_threepdm = True
            self.extraline.pop()

        # The binary files coming from STACKBLOCK and BLOCK are different
        # - STACKBLOCK uses the 6-fold symmetry, this must be unpacked
        #   using "libunpack.unpackE3" (see lib/icmpspt/icmpspt.c)
        # - BLOCK just writes a list of all values, this is directly read
        #   using "unpackE3_BLOCK" (see below)
        if (filetype == "binary") :
            fname = os.path.join(self.scratchDirectory,"node0", "spatial_threepdm.%d.%d.bin" %(state, state))
            if 'stackblock' in settings.BLOCKEXE:
              print 'Reading binary 3RDM from STACKBLOCK'
              fnameout = os.path.join(self.scratchDirectory,"node0", "spatial_threepdm.%d.%d.bin.unpack" %(state, state))
              libunpack.unpackE3(ctypes.c_char_p(fname), ctypes.c_char_p(fnameout), ctypes.c_int(norb))
              E3 = numpy.fromfile(fnameout, dtype=numpy.dtype('Float64'))
              E3 = numpy.reshape(E3, (norb, norb, norb, norb, norb, norb), order='F')
            else:
              print 'Reading binary 3RDM from BLOCK'
              E3 = DMRGCI.unpackE3_BLOCK(self,fname,norb)

        # The 3RDMs written by "Threepdm_container::save_spatial_npdm_text" in BLOCK and STACKBLOCK
        # are written as E3[i1,j2,k3,l3,m2,n1]
        # and are stored here as E3[i1,j2,k3,n1,m2,l3]
        # This is done with SQA in mind.
        else:
            print 'Reading text-file 3RDM'
            fname = os.path.join(self.scratchDirectory,"node0", "spatial_threepdm.%d.%d.txt" %(state, state))
            f = open(fname, 'r')
            lines = f.readlines()
            E3 = numpy.zeros(shape=(norb, norb, norb, norb, norb, norb), dtype=dt, order='F')
            assert(int(lines[0])==norb)
            for line in lines[1:]:
              linesp = line.split()
              if (len(linesp) != 7) :
                  continue
              a, b, c, d, e, f, integral = int(linesp[0]), int(linesp[1]), int(linesp[2]), int(linesp[3]), int(linesp[4]), int(linesp[5]), float(linesp[6])
              if (False):
                E3[a,b,c, f,e,d] = integral
                E3[a,c,b, f,d,e] = integral
                E3[b,a,c, e,f,d] = integral
                E3[b,c,a, e,d,f] = integral
                E3[c,a,b, d,f,e] = integral
                E3[c,b,a, d,e,f] = integral
              else:
                self.populate(E3, [a,b,c,  f,e,d], integral)
        print ''
        return E3

    def make_rdm4(self, state, norb, nelec, dt=numpy.dtype('Float64'), filetype = "binary", link_index=None, **kwargs):
        import os

        if self.has_fourpdm == False:
            self.twopdm = False
            self.threepdm = False
            self.extraline.append('threepdm\n')
            self.extraline.append('fourpdm\n')

            writeDMRGConfFile(self, nelec, False)
            if self.verbose >= logger.DEBUG1:
                inFile = self.configFile
                #inFile = os.path.join(self.scratchDirectory,self.configFile)
                logger.debug1(self, 'Block Input conf')
                logger.debug1(self, open(inFile, 'r').read())
            executeBLOCK(self)
            if self.verbose >= logger.DEBUG1:
                outFile = self.outputFile
                #outFile = os.path.join(self.scratchDirectory,self.outputFile)
                logger.debug1(self, open(outFile).read())
            self.has_fourpdm = True
            self.has_threepdm = True
            self.extraline.pop()

        # The binary files coming from STACKBLOCK and BLOCK are different:
        # - STACKBLOCK does not have 4RDM
        #   If it had, it would probably come in a 8-fold symmetr which must unpacked
        #   using "libunpack.unpackE4" (see lib/icmpspt/icmpspt.c)
        # - BLOCK just writes a list of all values, this is directly read
        #   using "unpackE4_BLOCK" (see below)
        if (filetype == "binary") :
            fname = os.path.join(self.scratchDirectory,"node0", "spatial_fourpdm.%d.%d.bin" %(state, state))
            if 'stackblock' in settings.BLOCKEXE:
              print 'Reading binary 4RDM from STACKBLOCK'
              fnameout = os.path.join(self.scratchDirectory,"node0", "spatial_fourpdm.%d.%d.bin.unpack" %(state, state))
              libunpack.unpackE4(ctypes.c_char_p(fname), ctypes.c_char_p(fnameout), ctypes.c_int(norb))
              E4 = numpy.fromfile(fnameout, dtype=numpy.dtype('Float64'))
              E4 = numpy.reshape(E4, (norb, norb, norb, norb, norb, norb, norb, norb), order='F')
            else:
              print 'Reading binary 4RDM from BLOCK'
              E4 = DMRGCI.unpackE4_BLOCK(self,fname,norb)

        # The 4RDMs written by "Fourpdm_container::save_spatial_npdm_text" in BLOCK and STACKBLOCK
        # are written as E4[i1,j2,k3,l4,m4,n3,o2,p1]
        # and are stored here as E4[i1,j2,k3,l4,p1,o2,n3,m4]
        # This is done with SQA in mind.
        else:
            print 'Reading text-file 4RDM'
            fname = os.path.join(self.scratchDirectory,"node0", "spatial_fourpdm.%d.%d.txt" %(state, state))
            f = open(fname, 'r')
            lines = f.readlines()
            E4 = numpy.zeros(shape=(norb, norb, norb, norb, norb, norb, norb, norb), dtype=dt, order='F')
            assert(int(lines[0])==norb)
            for line in lines[1:]:
              linesp = line.split()
              if (len(linesp) != 9) :
                  continue
              a, b, c, d, e, f, g, h, integral = int(linesp[0]), int(linesp[1]), int(linesp[2]), int(linesp[3]), int(linesp[4]), int(linesp[5]), int(linesp[6]), int(linesp[7]), float(linesp[8])
              if (False):
                up_indexes=[a,b,c,d]
                dn_indexes=[h,g,f,e]
                for i in range(4):
                  for j in range(4):
                    if (i==j):
                      continue
                    for k in range(4):
                      if ((i==k)or(j==k)):
                        continue
                      for l in range(4):
                        if ((i==l)or(j==l)or(k==l)):
                          continue
                        E4[up_indexes[i],up_indexes[j],up_indexes[k],up_indexes[l],\
                           dn_indexes[i],dn_indexes[j],dn_indexes[k],dn_indexes[l]] = integral
              else:
                self.populate(E4, [a,b,c,d,  h,g,f,e], integral)
        print ''
        return E4

    def populate(self, array, list, value):
        dim=len(list)/2
        up=list[:dim]
        dn=list[dim:]
        import itertools
        for t in itertools.permutations(range(dim), dim):
          updn=[up[i] for i in t]+[dn[i] for i in t]
          array[tuple(updn)] = value

    def unpackE3_BLOCK(self,fname,norb):
        # The 3RDMs written by "Threepdm_container::save_spatial_npdm_binary" in BLOCK
        # are written as E3[i1,j2,k3,l3,m2,n1]
        # and are stored here as E3[i1,j2,k3,n1,m2,l3]
        # This is done with SQA in mind.
        E3=numpy.zeros((norb,norb,norb,norb,norb,norb), order='F')
        fil=open(fname,"rb")
        fil.seek(93)
        for a in range(norb):
          for b in range(norb):
            for c in range(norb):
              for d in range(norb):
                for e in range(norb):
                  for f in range(norb):
                    (value,)=struct.unpack('d',fil.read(8))
                    E3[a,b,c,  f,e,d]=value
        fil.close()
        return E3

    def unpackE4_BLOCK(self,fname,norb):
        # The 4RDMs written by "Fourpdm_container::save_spatial_npdm_binary" in BLOCK
        # are written as E4[i1,j2,k3,l4,m4,n3,o2,p1]
        # and are stored here as E4[i1,j2,k3,l4,p1,o2,n3,m4]
        # This is done with SQA in mind.
        E4=numpy.zeros((norb,norb,norb,norb,norb,norb,norb,norb), order='F')
        fil=open(fname,"rb")
        fil.seek(109)
        for a in range(norb):
          for b in range(norb):
            for c in range(norb):
              for d in range(norb):
                for e in range(norb):
                  for f in range(norb):
                    for g in range(norb):
                      for h in range(norb):
                        (value,)=struct.unpack('d',fil.read(8))
                        E4[a,b,c,d,  h,g,f,e]=value
        fil.close()
        return E4

    def clearSchedule(self):
        self.scheduleSweeps = []
        self.scheduleMaxMs = []
        self.scheduleTols = []
        self.scheduleNoises = []

    def nevpt_intermediate(self, tag, norb, nelec, state, **kwargs):

        if self.has_nevpt == False:
            writeDMRGConfFile(self, nelec, True,
                              with_2pdm=False, extraline=['restart_nevpt2_npdm'])
            if self.verbose >= logger.DEBUG1:
                inFile = os.path.join(self.runtimeDir, self.configFile)
                logger.debug1(self, 'Block Input conf')
                logger.debug1(self, open(inFile, 'r').read())
            executeBLOCK(self)
            if self.verbose >= logger.DEBUG1:
                outFile = os.path.join(self.runtimeDir, self.outputFile)
                logger.debug1(self, open(outFile).read())
            self.has_nevpt = True

        a16 = numpy.zeros( (norb, norb, norb, norb, norb, norb) )
        filename = "%s_matrix.%d.%d.txt" % (tag, state, state)
        with open(os.path.join(self.scratchDirectory, "node0", filename), "r") as f:
            norb_read = int(f.readline().split()[0])
            assert(norb_read == norb)

            for line in f:
                linesp = line.split()
                i, j, k, l, m, n = [int(x) for x in linesp[:6]]
                a16[i,j,k,l,m,n] = float(linesp[6])

        return a16

    def kernel(self, h1e, eri, norb, nelec, fciRestart=None, ecore=0, **kwargs):
        if self.nroots == 1:
            roots = 0
        else:
            roots = range(self.nroots)
        if fciRestart is None:
            fciRestart = self.restart or self._restart

        if 'orbsym' in kwargs:
            self.orbsym = kwargs['orbsym']
        writeIntegralFile(self, h1e, eri, norb, nelec, ecore)
        writeDMRGConfFile(self, nelec, fciRestart)
        if self.verbose >= logger.DEBUG1:
            inFile = os.path.join(self.runtimeDir, self.configFile)
            logger.debug1(self, 'Block Input conf')
            logger.debug1(self, open(inFile, 'r').read())
        if self.onlywriteIntegral:
            logger.info(self, 'Only write integral')
            try:
                calc_e = readEnergy(self)
            except IOError:
                if self.nroots == 1:
                    calc_e = 0.0
                else :
                    calc_e = [0.0] * self.nroots
            return calc_e, roots

        executeBLOCK(self)
        if self.verbose >= logger.DEBUG1:
            outFile = os.path.join(self.runtimeDir, self.outputFile)
            logger.debug1(self, open(outFile).read())
        calc_e = readEnergy(self)

        return calc_e, roots

    def approx_kernel(self, h1e, eri, norb, nelec, fciRestart=None, ecore=0, **kwargs):
        fciRestart = True

        if 'orbsym' in kwargs:
            self.orbsym = kwargs['orbsym']
        writeIntegralFile(self, h1e, eri, norb, nelec, ecore)
        writeDMRGConfFile(self, nelec, fciRestart, self.approx_maxIter)
        if self.verbose >= logger.DEBUG1:
            inFile = os.path.join(self.runtimeDir, self.configFile)
            logger.debug1(self, 'Block Input conf')
            logger.debug1(self, open(inFile, 'r').read())
        executeBLOCK(self)
        if self.verbose >= logger.DEBUG1:
            outFile = os.path.join(self.runtimeDir, self.outputFile)
            logger.debug1(self, open(outFile).read())
        calc_e = readEnergy(self)

        if self.nroots==1:
            roots = 0
        else:
            roots = range(self.nroots)
        return calc_e, roots

    def restart_scheduler_(self):
        def callback(envs):
            if (envs['norm_gorb'] < self.dmrg_switch_tol or
                ('norm_ddm' in envs and envs['norm_ddm'] < self.dmrg_switch_tol*10)):
                self._restart = True
            else :
                self._restart = False
        return callback

# Block code also allows non-spin-adapted calculation. S^2 is not available in
# this type of calculation
    if 'spin_adapted' in settings.BLOCKEXE:
        def spin_square(self, civec, norb, nelec):
            if isinstance(nelec, (int, numpy.integer)):
                nelecb = nelec//2
                neleca = nelec - nelecb
            else :
                neleca, nelecb = nelec
            s = (neleca - nelecb) * .5
            ss = s * (s+1)
            return ss, s*2+1

def make_schedule(sweeps, Ms, tols, noises, twodot_to_onedot):
    if len(sweeps) == len(Ms) == len(tols) == len(noises):
        schedule = ['schedule']
        for i, s in enumerate(sweeps):
            schedule.append('%d %6d  %8.4e  %8.4e' % (s, Ms[i], tols[i], noises[i]))
        schedule.append('end')
        if (twodot_to_onedot != 0):
            schedule.append('twodot_to_onedot %i'%twodot_to_onedot)
        return '\n'.join(schedule)
    else:

        return 'schedule default\nmaxM %s'%Ms[-1]

def writeDMRGConfFile(DMRGCI, nelec, Restart,
                      maxIter=None, with_2pdm=True, extraline=[]):
    confFile = os.path.join(DMRGCI.runtimeDir, DMRGCI.configFile)

    f = open(confFile, 'w')

    if isinstance(nelec, (int, numpy.integer)):
        nelecb = (nelec-DMRGCI.spin) // 2
        neleca = nelec - nelecb
    else :
        neleca, nelecb = nelec
    f.write('nelec %i\n'%(neleca+nelecb))
    f.write('spin %i\n' %(neleca-nelecb))
    if DMRGCI.groupname is not None:
        if isinstance(DMRGCI.wfnsym, str):
            wfnsym = dmrg_sym.irrep_name2id(DMRGCI.groupname, DMRGCI.wfnsym)
        else:
            gpname = dmrg_sym.d2h_subgroup(DMRGCI.groupname)
            assert(DMRGCI.wfnsym in dmrg_sym.IRREP_MAP[gpname])
            wfnsym = DMRGCI.wfnsym
        f.write('irrep %i\n' % wfnsym)

    if (not Restart):
        #f.write('schedule\n')
        #f.write('0 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, 1e-5, 10.0))
        #f.write('1 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, 1e-5, 1e-4))
        #f.write('10 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, 1e-6, 1e-5))
        #f.write('16 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, DMRGCI.tol/10.0, 0e-6))
        #f.write('end\n')
        schedule = make_schedule(DMRGCI.scheduleSweeps,
                                 DMRGCI.scheduleMaxMs,
                                 DMRGCI.scheduleTols,
                                 DMRGCI.scheduleNoises,
                                 DMRGCI.twodot_to_onedot)
        f.write('%s\n' % schedule)
    else :
        f.write('schedule\n')
        #if approx == True :
        #    f.write('0 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, DMRGCI.tol*10.0, 0e-6))
        #else :
        #    f.write('0 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, DMRGCI.tol, 0e-6))
        f.write('0 %6i  %8.4e  %8.4e \n' %(DMRGCI.maxM, DMRGCI.tol/10, 0e-6))
        f.write('end\n')
        f.write('fullrestart\n')
        f.write('onedot \n')

    if DMRGCI.groupname is not None:
        f.write('sym %s\n' % dmrg_sym.d2h_subgroup(DMRGCI.groupname).lower())
    f.write('orbitals %s\n' % DMRGCI.integralFile)
    if maxIter is None:
        maxIter = DMRGCI.maxIter
    f.write('maxiter %i\n'%maxIter)
    f.write('sweep_tol %8.4e\n'%DMRGCI.tol)

    f.write('outputlevel %s\n'%DMRGCI.outputlevel)
    f.write('hf_occ integral\n')
    if(with_2pdm and DMRGCI.twopdm):
        f.write('twopdm\n')
    if(DMRGCI.nonspinAdapted):
        f.write('nonspinAdapted\n')
    if(DMRGCI.scratchDirectory):
        f.write('prefix  %s\n'%DMRGCI.scratchDirectory)
    if (DMRGCI.nroots !=1):
        f.write('nroots %d\n'%DMRGCI.nroots)
        if (DMRGCI.weights==[]):
            DMRGCI.weights= [1.0/DMRGCI.nroots]* DMRGCI.nroots
        f.write('weights ')
        for weight in DMRGCI.weights:
            f.write('%f '%weight)
        f.write('\n')

    block_extra_keyword = DMRGCI.extraline + DMRGCI.block_extra_keyword + extraline
    if block_version(DMRGCI.executable).startswith('1.1'):
        for line in block_extra_keyword:
            if not ('num_thrds' in line or 'memory' in line):
                f.write('%s\n'%line)
    else:
        if DMRGCI.memory is not None:
            f.write('memory, %i, g\n'%(DMRGCI.memory))
        if DMRGCI.num_thrds > 1:
            f.write('num_thrds %d\n'%DMRGCI.num_thrds)
        for line in block_extra_keyword:
            f.write('%s\n'%line)
    f.close()
    #no reorder
    #f.write('noreorder\n')
    return confFile

def writeIntegralFile(DMRGCI, h1eff, eri_cas, ncas, nelec, ecore=0):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelec//2 + nelec%2
        nelecb = nelec - neleca
    else :
        neleca, nelecb = nelec
    integralFile = os.path.join(DMRGCI.runtimeDir, DMRGCI.integralFile)
    if DMRGCI.groupname is not None and DMRGCI.orbsym is not []:
# First removing the symmetry forbidden integrals. This has been done using
# the pyscf internal irrep-IDs (stored in DMRGCI.orbsym)
        orbsym = numpy.asarray(DMRGCI.orbsym) % 10
        pair_irrep = (orbsym.reshape(-1,1) ^ orbsym)[numpy.tril_indices(ncas)]
        sym_forbid = pair_irrep.reshape(-1,1) != pair_irrep.ravel()
        eri_cas = pyscf.ao2mo.restore(4, eri_cas, ncas)
        eri_cas[sym_forbid] = 0
        eri_cas = pyscf.ao2mo.restore(8, eri_cas, ncas)
# Then convert the pyscf internal irrep-ID to molpro irrep-ID
        orbsym = numpy.asarray(dmrg_sym.convert_orbsym(DMRGCI.groupname, orbsym))
    else:
        orbsym = []
        eri_cas = pyscf.ao2mo.restore(8, eri_cas, ncas)
    if not os.path.exists(DMRGCI.scratchDirectory):
        os.makedirs(DMRGCI.scratchDirectory)
    if not os.path.exists(DMRGCI.runtimeDir):
        os.makedirs(DMRGCI.runtimeDir)

    pyscf.tools.fcidump.from_integrals(integralFile, h1eff, eri_cas, ncas,
                                       neleca+nelecb, ecore, ms=abs(neleca-nelecb),
                                       orbsym=orbsym)
    return integralFile

def executeBLOCK(DMRGCI):

    inFile  = DMRGCI.configFile
    outFile = DMRGCI.outputFile
    try:
        cmd = ' '.join((DMRGCI.mpiprefix, DMRGCI.executable, inFile))
        cmd = "%s > %s 2>&1" % (cmd, outFile)
        check_call(cmd, cwd=DMRGCI.runtimeDir, shell=True)
    except CalledProcessError as err:
        logger.error(DMRGCI, cmd)
        outFile = os.path.join(DMRGCI.runtimeDir, outFile)
        DMRGCI.stdout.write(check_output(['tail', '-100', outFile]))
        raise err

def readEnergy(DMRGCI):
    file1 = open(os.path.join(DMRGCI.scratchDirectory, "node0", "dmrg.e"), "rb")
    format = ['d']*DMRGCI.nroots
    format = ''.join(format)
    calc_e = struct.unpack(format, file1.read())
    file1.close()
    if DMRGCI.nroots == 1:
        return calc_e[0]
    else:
        return numpy.asarray(calc_e)

def DMRGSCF(mf, norb, nelec, maxM=1000, tol=1.e-8, *args, **kwargs):
    '''Shortcut function to setup CASSCF using the DMRG solver.  The DMRG
    solver is properly initialized in this function so that the 1-step
    algorithm can applied with DMRG-CASSCF.

    Examples:

    >>> mol = gto.M(atom='C 0 0 0; C 0 0 1')
    >>> mf = scf.RHF(mol).run()
    >>> mc = DMRGSCF(mf, 4, 4)
    >>> mc.kernel()
    -74.414908818611522
    '''
    if (hasattr(mf,'with_df')):
      mc = mcscf.DFCASSCF(mf, norb, nelec, *args, **kwargs)
    else:
      mc = mcscf.CASSCF(mf, norb, nelec, *args, **kwargs)
    mc.fcisolver = DMRGCI(mf.mol, maxM, tol=tol)
    mc.callback = mc.fcisolver.restart_scheduler_()
    if mc.chkfile == mc._scf._chkfile.name:
        # Do not delete chkfile after mcscf
        mc.chkfile = tempfile.mktemp(dir=settings.BLOCKSCRATCHDIR)
        if not os.path.exists(settings.BLOCKSCRATCHDIR):
            os.makedirs(settings.BLOCKSCRATCHDIR)
    return mc

def dryrun(mc, mo_coeff=None):
    '''Generate FCIDUMP and dmrg config file'''
    if mo_coeff is None:
        mo_coeff = mc.mo_coeff
    mc.fcisolver.onlywriteIntegral, bak = True, mc.fcisolver.onlywriteIntegral
    mc.casci(mo_coeff)
    mc.fcisolver.onlywriteIntegral = bak

def block_version(blockexe):
    version = getattr(settings, 'BLOCKVERSION', None)
    if isinstance(version, str):
        return version

    try:
        msg = check_output([blockexe, '-v'], stderr=STDOUT)
        version = '1.1.0'
        for line in msg.split('\n'):
            if line.startswith('Block '):
                version = line.split()[1]
                break
        return version
    except CalledProcessError:
        f1 = tempfile.NamedTemporaryFile()
        f1.write('memory 1 m\n')
        f1.flush()
        try:
            msg = check_output([blockexe, f1.name], stderr=STDOUT)
        except CalledProcessError as err:
            if 'Unrecognized option :: memory' in err.output:
                version = '1.1.1'
            elif 'need to specify hf_occ' in err.output:
                version = '1.5'
            else:
                sys.stderr.write(err.output)
                raise err
        f1.close()
        return version

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import mcscf
    settings.MPIPREFIX =''
    b = 1.4
    mol = gto.Mole()
    mol.build(
        verbose = 7,
        output = 'out-dmrgci',
        atom = [['H', (0.,0.,i-3.5)] for i in range(8)],
        basis = {'H': 'sto-3g'},
        symmetry = True
    )
    m = scf.RHF(mol)
    m.scf()

    mc = DMRGSCF(m, 4, 4)
    mc.fcisolver.tol = 1e-9
    emc_1 = mc.mc2step()[0]

    mc = mcscf.CASCI(m, 4, 4)
    mc.fcisolver = DMRGCI(mol)
    mc.fcisolver.scheduleSweeps = []
    emc_0 = mc.casci()[0]

    b = 1.4
    mol = gto.Mole()
    mol.build(
        verbose = 7,
        output = 'out-casscf',
        atom = [['H', (0.,0.,i-3.5)] for i in range(8)],
        basis = {'H': 'sto-3g'},
        symmetry = True
    )
    m = scf.RHF(mol)
    m.scf()

    mc = mcscf.CASSCF(m, 4, 4)
    emc_1ref = mc.mc2step()[0]

    mc = mcscf.CASCI(m, 4, 4)
    emc_0ref = mc.casci()[0]

    print('DMRGCI  = %.15g CASCI  = %.15g' % (emc_0, emc_0ref))
    print('DMRGSCF = %.15g CASSCF = %.15g' % (emc_1, emc_1ref))


"""
Class for normalizing fission energy deposition
"""
import bisect
from collections import defaultdict
from copy import deepcopy
from itertools import product
from numbers import Real
import sys

from numpy import dot, zeros, newaxis, asarray

from openmc.mpi import comm
from openmc.checkvalue import check_type, check_greater_than
from openmc.data import JOULE_PER_EV, REACTION_MT
from openmc.lib import (
    Tally, MaterialFilter, EnergyFilter, EnergyFunctionFilter)
import openmc.lib
from .abc import (
    ReactionRateHelper, NormalizationHelper, FissionYieldHelper,
    TalliedFissionYieldHelper)

__all__ = (
    "DirectReactionRateHelper", "ChainFissionHelper", "EnergyScoreHelper"
    "SourceRateHelper", "ConstantFissionYieldHelper", "FissionYieldCutoffHelper",
    "AveragedFissionYieldHelper", "FluxCollapseHelper")

# -------------------------------------
# Helpers for generating reaction rates
# -------------------------------------

class DirectReactionRateHelper(ReactionRateHelper):
    """Class for generating one-group reaction rates with direct tallies

    This class generates reaction rate tallies for each nuclide and
    transmutation reaction relevant for a depletion calculation.

    Parameters
    ----------
    n_nucs : int
        Number of burnable nuclides tracked by :class:`openmc.deplete.Operator`
    n_react : int
        Number of reactions tracked by :class:`openmc.deplete.Operator`

    Attributes
    ----------
    nuclides : list of str
        All nuclides with desired reaction rates.
    """
    def __init__(self, n_nuc, n_react):
        super().__init__(n_nuc, n_react)
        self._rate_tally = None

        # Automatically pre-calculate reaction rates for depletion
        openmc.lib.settings.need_depletion_rx = True

    @ReactionRateHelper.nuclides.setter
    def nuclides(self, nuclides):
        ReactionRateHelper.nuclides.fset(self, nuclides)
        self._rate_tally.nuclides = nuclides

    def generate_tallies(self, materials, scores):
        """Produce one-group reaction rate tally

        Uses the :mod:`openmc.lib` to generate a tally
        of relevant reactions across all burnable materials.

        Parameters
        ----------
        materials : iterable of :class:`openmc.Material`
            Burnable materials in the problem. Used to
            construct a :class:`openmc.MaterialFilter`
        scores : iterable of str
            Reaction identifiers, e.g. ``"(n, fission)"``,
            ``"(n, gamma)"``, needed for the reaction rate tally.
        """
        self._rate_tally = Tally()
        self._rate_tally.writable = False
        self._rate_tally.scores = scores
        self._rate_tally.filters = [MaterialFilter(materials)]

    def get_material_rates(self, mat_id, nuc_index, react_index):
        """Return an array of reaction rates for a material

        Parameters
        ----------
        mat_id : int
            Unique ID for the requested material
        nuc_index : iterable of int
            Index for each nuclide in :attr:`nuclides` in the
            desired reaction rate matrix
        react_index : iterable of int
            Index for each reaction scored in the tally

        Returns
        -------
        rates : numpy.ndarray
            Array with shape ``(n_nuclides, n_rxns)`` with the
            reaction rates in this material
        """
        self._results_cache.fill(0.0)
        full_tally_res = self._rate_tally.mean[mat_id]
        for i_tally, (i_nuc, i_react) in enumerate(
                product(nuc_index, react_index)):
            self._results_cache[i_nuc, i_react] = full_tally_res[i_tally]

        return self._results_cache

class FluxCollapseHelper(ReactionRateHelper):
    """Class that generates one-group reaction rates using multigroup flux

    This class generates a multigroup flux tally that is used afterward to
    calculate a one-group reaction rate by collapsing it with continuous-energy
    cross section data. Additionally, select nuclides/reactions can be treated
    with a direct reaction rate tally when using a multigroup flux spectrum
    would not be sufficiently accurate. This is often the case for (n,gamma) and
    fission reactions.

    .. versionadded:: 0.12.1

    Parameters
    ----------
    n_nucs : int
        Number of burnable nuclides tracked by :class:`openmc.deplete.Operator`
    n_react : int
        Number of reactions tracked by :class:`openmc.deplete.Operator`
    energies : iterable of float
        Energy group boundaries for flux spectrum in [eV]
    reactions : iterable of str
        Reactions for which rates should be directly tallied
    nuclides : iterable of str
        Nuclides for which some reaction rates should be directly tallied. If
        None, then ``reactions`` will be used for all nuclides.

    Attributes
    ----------
    nuclides : list of str
        All nuclides with desired reaction rates.

    """
    def __init__(self, n_nucs, n_reacts, energies, reactions=None, nuclides=None):
        super().__init__(n_nucs, n_reacts)
        self._energies = asarray(energies)
        self._reactions_direct = list(reactions) if reactions is not None else []
        self._nuclides_direct = list(nuclides) if nuclides is not None else None

    @ReactionRateHelper.nuclides.setter
    def nuclides(self, nuclides):
        ReactionRateHelper.nuclides.fset(self, nuclides)
        if self._reactions_direct and self._nuclides_direct is None:
            self._rate_tally.nuclides = nuclides

    def generate_tallies(self, materials, scores):
        """Produce multigroup flux spectrum tally

        Uses the :mod:`openmc.lib` module to generate a multigroup flux tally
        for each burnable material.

        Parameters
        ----------
        materials : iterable of :class:`openmc.Material`
            Burnable materials in the problem. Used to construct a
            :class:`openmc.MaterialFilter`
        scores : iterable of str
            Reaction identifiers, e.g. ``"(n, fission)"``, ``"(n, gamma)"``,
            needed for the reaction rate tally.
        """
        self._materials = materials

        # adds an entry for fisson to the dictionary of reactions
        self._mts = [REACTION_MT[x] for x in scores]
        self._scores = scores

        # Create flux tally with material and energy filters
        self._flux_tally = Tally()
        self._flux_tally.writable = False
        self._flux_tally.filters = [
            MaterialFilter(materials),
            EnergyFilter(self._energies)
        ]
        self._flux_tally.scores = ['flux']

        # Create reaction rate tally
        if self._reactions_direct:
            self._rate_tally = Tally()
            self._rate_tally.writable = False
            self._rate_tally.scores = self._reactions_direct
            self._rate_tally.filters = [MaterialFilter(materials)]
            if self._nuclides_direct is not None:
                self._rate_tally.nuclides = self._nuclides_direct

    def get_material_rates(self, mat_index, nuc_index, react_index):
        """Return an array of reaction rates for a material

        Parameters
        ----------
        mat_index : int
            Index for material
        nuc_index : iterable of int
            Index for each nuclide in :attr:`nuclides` in the
            desired reaction rate matrix
        react_index : iterable of int
            Index for each reaction scored in the tally

        Returns
        -------
        rates : numpy.ndarray
            Array with shape ``(n_nuclides, n_rxns)`` with the reaction rates in
            this material

        """
        self._results_cache.fill(0.0)

        # Get flux for specified material
        shape = (len(self._materials), len(self._energies) - 1)
        mean_value = self._flux_tally.mean.reshape(shape)
        flux = mean_value[mat_index]

        # Get direct reaction rates
        if self._reactions_direct:
            nuclides_direct = self._rate_tally.nuclides
            shape = (len(nuclides_direct), len(self._reactions_direct))
            rx_rates = self._rate_tally.mean[mat_index].reshape(shape)

        mat = self._materials[mat_index]

        # Build nucname: density mapping to enable O(1) lookup in loop below
        densities = dict(zip(mat.nuclides, mat.densities))

        for name, i_nuc in zip(self.nuclides, nuc_index):
            # Determine density of nuclide
            density = densities[name]

            for mt, score, i_rx in zip(self._mts, self._scores, react_index):
                if score in self._reactions_direct and name in nuclides_direct:
                    # Determine index in rx_rates
                    i_rx_direct = self._reactions_direct.index(score)
                    i_nuc_direct = nuclides_direct.index(name)

                    # Get reaction rate from tally
                    self._results_cache[i_nuc, i_rx] = rx_rates[i_nuc_direct, i_rx_direct]
                else:
                    # Use flux to collapse reaction rate (per N)
                    nuc = openmc.lib.nuclides[name]
                    rate_per_nuc = nuc.collapse_rate(
                        mt, mat.temperature, self._energies, flux)

                    # Multiply by density to get absolute reaction rate
                    self._results_cache[i_nuc, i_rx] = rate_per_nuc * density

        return self._results_cache

# ------------------------------------------
# Helpers for obtaining normalization factor
# ------------------------------------------

class EnergyNormalizationHelper(NormalizationHelper):
    """Compute energy-based normalization."""

    def reset(self):
        """Reset energy produced prior to unpacking tallies"""
        self._energy = 0.0

    def factor(self, source_rate):
        # Reduce energy produced from all processes
        # J / source neutron
        energy = comm.allreduce(self._energy) * JOULE_PER_EV

        # Guard against divide by zero
        if energy == 0:
            if comm.rank == 0:
                sys.stderr.flush()
                print("No energy reported from OpenMC tallies. Do your HDF5 "
                      "files have heating data?\n", file=sys.stderr, flush=True)
            comm.barrier()
            comm.Abort(1)

        # Return normalization factor for scaling reaction rates. In this case,
        # the source rate is the power in [W], so [W] / [J/src] = [src/s]
        return source_rate / energy

class ChainFissionHelper(EnergyNormalizationHelper):
    """Computes normalization using fission Q values from depletion chain

    Attributes
    ----------
    nuclides : list of str
        All nuclides with desired reaction rates. Ordered to be
        consistent with :class:`openmc.deplete.Operator`
    energy : float
        Total energy [J/s/source neutron] produced in a transport simulation.
        Updated in the material iteration with :meth:`update`.
    """

    def __init__(self):
        super().__init__()
        self._fission_q_vector = None

    def prepare(self, chain_nucs, rate_index):
        """Populate the fission Q value vector from a chain.

        Parameters
        ----------
        chain_nucs : iterable of :class:`openmc.deplete.Nuclide`
            Nuclides used in this depletion chain. Do not need
            to be ordered
        rate_index : dict of str to int
            Dictionary mapping names of nuclides, e.g. ``"U235"``,
            to a corresponding index in the desired fission Q
            vector.
        """
        if (self._fission_q_vector is not None
                and self._fission_q_vector.shape == (len(rate_index),)):
            return

        fission_qs = zeros(len(rate_index))

        for nuclide in chain_nucs:
            if nuclide.name in rate_index:
                for rx in nuclide.reactions:
                    if rx.type == "fission":
                        fission_qs[rate_index[nuclide.name]] = rx.Q
                        break

        self._fission_q_vector = fission_qs

    def update(self, fission_rates):
        """Update energy produced with fission rates in a material

        Parameters
        ----------
        fission_rates : numpy.ndarray
            fission reaction rate for each isotope in the specified
            material. Should be ordered corresponding to initial
            ``rate_index`` used in :meth:`prepare`
        """
        self._energy += dot(fission_rates, self._fission_q_vector)

class EnergyScoreHelper(EnergyNormalizationHelper):
    """Class responsible for obtaining system energy via a tally score

    Parameters
    ----------
    score : string
        Valid score to use when obtaining system energy from OpenMC.
        Defaults to "heating-local"

    Attributes
    ----------
    nuclides : list of str
        List of nuclides with reaction rates. Not needed, but provided
        for a consistent API across other :class:`NormalizationHelper`
    energy : float
        System energy [eV] computed from the tally. Will be zero for
        all MPI processes that are not the "master" process to avoid
        artificially increasing the tallied energy.
    score : str
        Score used to obtain system energy

    """

    def __init__(self, score="heating-local"):
        super().__init__()
        self.score = score
        self._tally = None

    def prepare(self, *args, **kwargs):
        """Create a tally for system energy production

        Input arguments are not used, as the only information needed
        is :attr:`score`

        """
        self._tally = Tally()
        self._tally.writable = False
        self._tally.scores = [self.score]

    def reset(self):
        """Obtain system energy from tally

        Only the master process, ``comm.rank == 0`` will
        have a non-zero :attr:`energy` taken from the tally.
        This avoids accidentally scaling the system power by
        the number of MPI processes
        """
        super().reset()
        if comm.rank == 0:
            self._energy = self._tally.mean[0, 0]

class SourceRateHelper(NormalizationHelper):
    def prepare(self, *args, **kwargs):
        pass

    def factor(self, source_rate):
        return source_rate

# ------------------------------------
# Helper for collapsing fission yields
# ------------------------------------

class ConstantFissionYieldHelper(FissionYieldHelper):
    """Class that uses a single set of fission yields on each isotope

    Parameters
    ----------
    chain_nuclides : iterable of openmc.deplete.Nuclide
        Nuclides tracked in the depletion chain. All nuclides are
        not required to have fission yield data.
    energy : float, optional
        Key in :attr:`openmc.deplete.Nuclide.yield_data` corresponding
        to the desired set of fission yield data. Typically one of
        ``{0.0253, 500000, 14000000}`` corresponding to 0.0253 eV,
        500 keV, and 14 MeV yield libraries. If the specific key is not
        found, will fall back to closest energy present.
        Default: 0.0253 eV for thermal yields

    Attributes
    ----------
    constant_yields : collections.defaultdict
        Fission yields for all nuclides that only have one set of
        fission yield data. Dictionary of form ``{str: {str: float}}``
        representing yields for ``{parent: {product: yield}}``. Default
        return object is an empty dictionary
    energy : float
        Energy of fission yield libraries.
    """

    def __init__(self, chain_nuclides, energy=0.0253):
        check_type("energy", energy, Real)
        check_greater_than("energy", energy, 0.0, equality=True)
        self._energy = energy
        super().__init__(chain_nuclides)
        # Iterate over all nuclides with > 1 set of yields
        for name, nuc in self._chain_nuclides.items():
            yield_data = nuc.yield_data.get(energy)
            if yield_data is not None:
                self._constant_yields[name] = yield_data
                continue
            # Specific energy not found, use closest energy
            distances = [abs(energy - ene) for ene in nuc.yield_energies]
            min_E = min(nuc.yield_energies, key=lambda e: abs(e - energy))
            self._constant_yields[name] = nuc.yield_data[min_E]

    @classmethod
    def from_operator(cls, operator, **kwargs):
        """Return a new ConstantFissionYieldHelper using operator data

        All keyword arguments should be identical to their counterpart
        in the main ``__init__`` method

        Parameters
        ----------
        operator : openmc.deplete.TransportOperator
            operator with a depletion chain
        kwargs:
            Additional keyword arguments to be used in construction

        Returns
        -------
        ConstantFissionYieldHelper
        """
        return cls(operator.chain.nuclides, **kwargs)

    @property
    def energy(self):
        return self._energy

    def weighted_yields(self, _local_mat_index=None):
        """Return fission yields for all nuclides requested

        Parameters
        ----------
        _local_mat_index : int, optional
            Current material index. Not used since all yields are
            constant

        Returns
        -------
        library : collections.defaultdict
            Dictionary of ``{parent: {product: fyield}}``
        """
        return self.constant_yields

class FissionYieldCutoffHelper(TalliedFissionYieldHelper):
    """Helper that computes fission yields based on a cutoff energy

    Tally fission rates above and below the cutoff energy.
    Assume that all fissions below cutoff energy have use thermal fission
    product yield distributions, while all fissions above use a faster
    set of yield distributions.

    Uses a limit of 20 MeV for tallying fission.

    Parameters
    ----------
    chain_nuclides : iterable of openmc.deplete.Nuclide
        Nuclides tracked in the depletion chain. All nuclides are
        not required to have fission yield data.
    n_bmats : int, optional
        Number of burnable materials tracked in the problem
    cutoff : float, optional
        Cutoff energy in [eV] below which all fissions will be
        use thermal yields. All other fissions will use a
        faster set of yields. Default: 112 [eV]
    thermal_energy : float, optional
        Energy of yield data corresponding to thermal yields.
        Default: 0.0253 [eV]
    fast_energy : float, optional
        Energy of yield data corresponding to fast yields.
        Default: 500 [kev]

    Attributes
    ----------
    n_bmats : int
        Number of burnable materials tracked in the problem.
        Must be set prior to generating tallies
    thermal_yields : dict
        Dictionary of the form ``{parent: {product: yield}}``
        with thermal yields
    fast_yields : dict
        Dictionary of the form ``{parent: {product: yield}}``
        with fast yields
    constant_yields : collections.defaultdict
        Fission yields for all nuclides that only have one set of
        fission yield data. Dictionary of form ``{str: {str: float}}``
        representing yields for ``{parent: {product: yield}}``. Default
        return object is an empty dictionary
    results : numpy.ndarray
        Array of fission rate fractions with shape
        ``(n_mats, 2, n_nucs)``. ``results[:, 0]``
        corresponds to the fraction of all fissions
        that occured below ``cutoff``. The number
        of materials in the first axis corresponds
        to the number of materials burned by the
        :class:`openmc.deplete.Operator`
    """

    def __init__(self, chain_nuclides, n_bmats, cutoff=112.0,
                 thermal_energy=0.0253, fast_energy=500.0e3):
        check_type("cutoff", cutoff, Real)
        check_type("thermal_energy", thermal_energy, Real)
        check_type("fast_energy", fast_energy, Real)
        check_greater_than("thermal_energy", thermal_energy, 0.0, equality=True)
        check_greater_than("cutoff", cutoff, thermal_energy, equality=False)
        check_greater_than("fast_energy", fast_energy, cutoff, equality=False)
        self.n_bmats = n_bmats
        super().__init__(chain_nuclides)
        self._cutoff = cutoff
        self._thermal_yields = {}
        self._fast_yields = {}
        convert_to_constant = set()
        for name, nuc in self._chain_nuclides.items():
            yields = nuc.yield_data
            energies = nuc.yield_energies
            thermal = yields.get(thermal_energy)
            fast = yields.get(fast_energy)
            if thermal is None or fast is None:
                if cutoff <= energies[0]:
                    # use lowest energy yields as constant
                    self._constant_yields[name] = yields[energies[0]]
                    convert_to_constant.add(name)
                    continue
                if cutoff >= energies[-1]:
                    # use highest energy yields as constant
                    self._constant_yields[name] = yields[energies[-1]]
                    convert_to_constant.add(name)
                    continue
                cutoff_ix = bisect.bisect_left(energies, cutoff)
                # find closest energy to requested thermal, fast energies
                if thermal is None:
                    min_E = min(energies[:cutoff_ix],
                                key=lambda e: abs(e - thermal_energy))
                    thermal = yields[min_E]
                if fast is None:
                    min_E = min(energies[cutoff_ix:],
                                key=lambda e: abs(e - fast_energy))
                    fast = yields[min_E]
            self._thermal_yields[name] = thermal
            self._fast_yields[name] = fast
        for name in convert_to_constant:
            self._chain_nuclides.pop(name)

    @classmethod
    def from_operator(cls, operator, **kwargs):
        """Construct a helper from an operator

        All keyword arguments should be identical to their counterpart
        in the main ``__init__`` method

        Parameters
        ----------
        operator : openmc.deplete.Operator
            Operator with a chain and burnable materials
        kwargs:
            Additional keyword arguments to be used in construction

        Returns
        -------
        FissionYieldCutoffHelper

        """
        return cls(operator.chain.nuclides, len(operator.burnable_mats),
                   **kwargs)

    def generate_tallies(self, materials, mat_indexes):
        """Use C API to produce a fission rate tally in burnable materials

        Include a :class:`openmc.lib.EnergyFilter` to tally fission rates
        above and below cutoff energy.

        Parameters
        ----------
        materials : iterable of :class:`openmc.lib.Material`
            Materials to be used in :class:`openmc.lib.MaterialFilter`
        mat_indexes : iterable of int
            Indices of tallied materials that will have their fission
            yields computed by this helper. Necessary as the
            :class:`openmc.deplete.Operator` that uses this helper
            may only burn a subset of all materials when running
            in parallel mode.
        """
        super().generate_tallies(materials, mat_indexes)
        energy_filter = EnergyFilter([0.0, self._cutoff, self._upper_energy])
        self._fission_rate_tally.filters = (
            self._fission_rate_tally.filters + [energy_filter])

    def unpack(self):
        """Obtain fast and thermal fission fractions from tally"""
        if not self._tally_nucs or self._local_indexes.size == 0:
            self.results = None
            return
        fission_rates = self._fission_rate_tally.mean.reshape(
            self.n_bmats, 2, len(self._tally_nucs))
        self.results = fission_rates[self._local_indexes]
        total_fission = self.results.sum(axis=1)
        nz_mat, nz_nuc = total_fission.nonzero()
        self.results[nz_mat, :, nz_nuc] /= total_fission[nz_mat, newaxis, nz_nuc]

    def weighted_yields(self, local_mat_index):
        """Return fission yields for a specific material

        For nuclides with both yield data above and below
        the cutoff energy, the effective yield for nuclide ``A``
        will be a weighted sum of fast and thermal yields. The
        weights will be the fraction of ``A`` fission events
        in the above and below the cutoff energy.

        If ``A`` has fission product distribution ``F``
        for fast fissions and ``T`` for thermal fissions, and
        70% of ``A`` fissions are considered thermal, then
        the effective fission product yield distributions
        for ``A`` is ``0.7 * T + 0.3 * F``

        Parameters
        ----------
        local_mat_index : int
            Index for specific burnable material. Effective
            yields will be produced using
            ``self.results[local_mat_index]``

        Returns
        -------
        library : collections.defaultdict
            Dictionary of ``{parent: {product: fyield}}``
        """
        yields = self.constant_yields
        if not self._tally_nucs:
            return yields
        rates = self.results[local_mat_index]
        # iterate over thermal then fast yields, prefer __mul__ to __rmul__
        for therm_frac, fast_frac, nuc in zip(rates[0], rates[1], self._tally_nucs):
            yields[nuc.name] = (self._thermal_yields[nuc.name] * therm_frac
                                + self._fast_yields[nuc.name] * fast_frac)
        return yields

    @property
    def thermal_yields(self):
        return deepcopy(self._thermal_yields)

    @property
    def fast_yields(self):
        return deepcopy(self._fast_yields)

class AveragedFissionYieldHelper(TalliedFissionYieldHelper):
    r"""Class that computes fission yields based on average fission energy

    Computes average energy at which fission events occured with

    .. math::

        \bar{E} = \frac{
            \int_0^\infty E\sigma_f(E)\phi(E)dE
        }{
            \int_0^\infty\sigma_f(E)\phi(E)dE
        }

    If the average energy for a nuclide is below the lowest energy
    with yield data, that set of fission yields is taken.
    Conversely, if the average energy is above the highest energy
    with yield data, that set of fission yields is used.
    For the case where the average energy is between two sets
    of yields, the effective fission yield computed by
    linearly interpolating between yields provided at the
    nearest energies above and below the average.

    Parameters
    ----------
    chain_nuclides : iterable of openmc.deplete.Nuclide
        Nuclides tracked in the depletion chain. All nuclides are
        not required to have fission yield data.

    Attributes
    ----------
    constant_yields : collections.defaultdict
        Fission yields for all nuclides that only have one set of
        fission yield data. Dictionary of form ``{str: {str: float}}``
        representing yields for ``{parent: {product: yield}}``. Default
        return object is an empty dictionary
    results : None or numpy.ndarray
        If tallies have been generated and unpacked, then the array will
        have shape ``(n_mats, n_tnucs)``, where ``n_mats`` is the number
        of materials where fission reactions were tallied and ``n_tnucs``
        is the number of nuclides with multiple sets of fission yields.
        Data in the array are the average energy of fission events for
        tallied nuclides across burnable materials.
    """

    def __init__(self, chain_nuclides):
        super().__init__(chain_nuclides)
        self._weighted_tally = None

    def generate_tallies(self, materials, mat_indexes):
        """Construct tallies to determine average energy of fissions

        Parameters
        ----------
        materials : iterable of :class:`openmc.lib.Material`
            Materials to be used in :class:`openmc.lib.MaterialFilter`
        mat_indexes : iterable of int
            Indices of tallied materials that will have their fission
            yields computed by this helper. Necessary as the
            :class:`openmc.deplete.Operator` that uses this helper
            may only burn a subset of all materials when running
            in parallel mode.
        """
        super().generate_tallies(materials, mat_indexes)
        fission_tally = self._fission_rate_tally
        filters = fission_tally.filters

        ene_filter = EnergyFilter([0, self._upper_energy])
        fission_tally.filters = filters + [ene_filter]

        func_filter = EnergyFunctionFilter()
        func_filter.set_data((0, self._upper_energy), (0, self._upper_energy))
        weighted_tally = Tally()
        weighted_tally.writable = False
        weighted_tally.scores = ['fission']
        weighted_tally.filters = filters + [func_filter]
        self._weighted_tally = weighted_tally

    def update_tally_nuclides(self, nuclides):
        """Tally nuclides with non-zero density and multiple yields

        Must be run after :meth:`generate_tallies`.

        Parameters
        ----------
        nuclides : iterable of str
            Potential nuclides to be tallied, such as those with
            non-zero density at this stage.

        Returns
        -------
        nuclides : tuple of str
            Union of input nuclides and those that have multiple sets
            of yield data.  Sorted by nuclide name

        Raises
        ------
        AttributeError
            If tallies not generated
        """
        tally_nucs = super().update_tally_nuclides(nuclides)
        self._weighted_tally.nuclides = tally_nucs
        return tally_nucs

    def unpack(self):
        """Unpack tallies and populate :attr:`results` with average energies"""
        if not self._tally_nucs or self._local_indexes.size == 0:
            self.results = None
            return
        fission_results = (
            self._fission_rate_tally.mean[self._local_indexes])
        self.results = (
            self._weighted_tally.mean[self._local_indexes]).copy()
        nz_mat, nz_nuc = fission_results.nonzero()
        self.results[nz_mat, nz_nuc] /= fission_results[nz_mat, nz_nuc]

    def weighted_yields(self, local_mat_index):
        """Return fission yields for a specific material

        Use the computed average energy of fission
        events to determine fission yields. If average
        energy is between two sets of yields, linearly
        interpolate bewteen the two.
        Otherwise take the closet set of yields.

        Parameters
        ----------
        local_mat_index : int
            Index for specific burnable material. Effective
            yields will be produced using
            ``self.results[local_mat_index]``

        Returns
        -------
        library : collections.defaultdict
            Dictionary of ``{parent: {product: fyield}}``. Default return
            value is an empty dictionary
        """
        if not self._tally_nucs:
            return self.constant_yields
        mat_yields = defaultdict(dict)
        average_energies = self.results[local_mat_index]
        for avg_e, nuc in zip(average_energies, self._tally_nucs):
            nuc_energies = nuc.yield_energies
            if avg_e <= nuc_energies[0]:
                mat_yields[nuc.name] = nuc.yield_data[nuc_energies[0]]
                continue
            if avg_e >= nuc_energies[-1]:
                mat_yields[nuc.name] = nuc.yield_data[nuc_energies[-1]]
                continue
            # in-between two energies
            # linear search since there are usually ~3 energies
            for ix, ene in enumerate(nuc_energies[:-1]):
                if nuc_energies[ix + 1] > avg_e:
                    break
            lower, upper = nuc_energies[ix:ix + 2]
            fast_frac = (avg_e - lower) / (upper - lower)
            mat_yields[nuc.name] = (
                nuc.yield_data[lower] * (1 - fast_frac)
                + nuc.yield_data[upper] * fast_frac)
        mat_yields.update(self.constant_yields)
        return mat_yields

    @classmethod
    def from_operator(cls, operator, **kwargs):
        """Return a new helper with data from an operator

        All keyword arguments should be identical to their counterpart
        in the main ``__init__`` method

        Parameters
        ----------
        operator : openmc.deplete.TransportOperator
            Operator with a depletion chain
        kwargs :
            Additional keyword arguments to be used in construction

        Returns
        -------
        AveragedFissionYieldHelper
        """
        return cls(operator.chain.nuclides)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import math
import numpy

class Wavefront():

    def __init__(self, path, triangulateQuads = True):

        # These are the arrays we're going to produce, and which might
        # make sense to manipulate from the outside:

        self.vertexCoords = None  # XYZ coordinated for each vertex
        self.vertexNormals = None # Normal (in XYZ form) for each vertex
        self.faces = None         # Faces, specified by listing indexes for participating vertices

        # These are internal work arrays

        self._rawVertices = []
        self._rawTexCo = []
        self._rawVertexTexCo = [];
        self._rawFaces = []
        self._vertexBelongsToFaces = None
        self._faceVertCache = None

        self.triangulateQuads = triangulateQuads

        if not os.path.exists(path):
            raise IOError(path + " does not exist")

        self.content = []

        with open(path,'r') as file:
            self.content = file.readlines()

        # self._mode can be:
        #   ONLYTRIS:    The incoming mesh only contains tris (so no need to do anything)
        #   TRIANGULATE: The incoming mesh contains quads (and may contain tris). Triangulate to get only tris.
        #   ONLYQUADS:   The incoming mesh contains only quads. Keep these rather than triangulating

        self._mode = None

        # Check if mesh contains quads and/or tris
        self._scanForMode()

        # Make a sweep for vertices as faces need that information
        self._extractVertices()

        # Make a sweep for texture coordinates, as faces need that too
        self._extractTextureCoordinates()

        # Make a sweep for faces
        self._extractFaces()

        # TODO: Find texture coordinates for vertices

        # create numpy arrays to contain vertices, faces and normals
        self._createVerticesNumpyArray()
        self._createFacesNumpyArray()

        # These two operations need to be redone if vertex coordinates
        # are changed
        self.recalculateFaceNormals()
        self.recalculateVertexNormals()

    def _scanForMode(self):

        containsTris = False
        containsQuads = False

        for line in self.content:
            strippedLine = line.strip()
            if not strippedLine is None and not strippedLine == "" and not strippedLine[0] == "#":
                parts = strippedLine.split(' ')
                if len(parts) > 4:
                    containsQuads = True
                if len(parts) == 4:
                    containsTris = True
                if len(parts) > 5:
                    raise ValueError("Found a face with more than four vertices. N-gons are not supported.")
                # TODO: Check for n-gons?

        if containsQuads:
            if self.triangulateQuads:
                self._mode = "TRIANGULATE"
            else:
                if containsTris:
                    raise ValueError("Since the mesh contains both tris and quads, requesting the mesh to not be triangulated is illegal")
                else:
                    self._mode = "ONLYQUADS"
        else:
            if containsTris:
                self._mode = "ONLYTRIS"
            else:
                raise ValueError("The mesh didn't contain tris nor quads!?")

        if self._mode == "ONLYQUADS":
            raise ValueError("The ONLYQUADS mode is not implemented yet")

        print("Tris " + str(containsTris))
        print("Quads " + str(containsQuads))
        print(self._mode)

    def _extractVertices(self):
        for line in self.content:
            strippedLine = line.strip()
            if not strippedLine is None and not strippedLine == "" and not strippedLine[0] == "#":
                parts = strippedLine.split(' ')
                if len(parts) > 1:
                    command = parts[0]
                    if command == "v":
                        x = float(parts[1])
                        y = float(parts[2])
                        z = float(parts[3])
                        vertex = [x, y, z]
                        self._rawVertices.append(vertex)
                        self._rawVertexTexCo.append([0,0])

    def _extractTextureCoordinates(self):

        self.hasTexCo = False

        for line in self.content:
            strippedLine = line.strip()
            if not strippedLine is None and not strippedLine == "" and not strippedLine[0] == "#":
                parts = strippedLine.split(' ')
                if len(parts) > 1:
                    command = parts[0]
                    if command == "vt":
                        x = float(parts[1])
                        y = float(parts[2])
                        texco = [x, y]
                        self._rawTexCo.append(texco)

    def _distanceBetweenVerticesByIdx(self, idx1, idx2):

        vert1 = numpy.array(self._rawVertices[idx1])
        vert2 = numpy.array(self._rawVertices[idx2])

        difference = vert2 - vert1

        x = difference[0]
        y = difference[1]
        z = difference[2]

        distance = math.sqrt( x*x + y*y + z*z )
        return distance

    def _extractFaces(self):

        # Note that wavefront lists starts at 1, not 0

        for line in self.content:
            strippedLine = line.strip()
            if not strippedLine is None and not strippedLine == "" and not strippedLine[0] == "#":
                parts = strippedLine.split(' ')
                if len(parts) > 1:
                    command = parts[0]
                    if command == "f":

                        # Face info is vertIdx / texCoIdx / faceNormalIdx  OR  vertIdx / texCoIdx  OR  vertIdx

                        vInfo1 = parts[1].split('/')
                        vInfo2 = parts[2].split('/')
                        vInfo3 = parts[3].split('/')

                        # Find indexes of vertices making up the face. Note "-1" since wavefront indexes start
                        # at 1 rather than 0
                        vidx1 = int(vInfo1[0]) - 1
                        vidx2 = int(vInfo2[0]) - 1
                        vidx3 = int(vInfo3[0]) - 1

                        if len(parts) == 4:
                            if self._mode == "ONLYQUADS":
                                raise ValueError("Found tri although mode was ONLYQUADS")
                            face = [vidx1, vidx2, vidx3]
                            self._rawFaces.append(face)
                        else:
                            vInfo4 = parts[4].split('/')
                            vidx4 = int(vInfo4[0]) - 1
                            if self._mode == "ONLYTRIS":
                                raise ValueError("Found quad although mode was ONLYTRIS")
                            if self._mode == "ONLYQUADS":
                                raise ValueError("ONLYQUADS mode not implemented yet")

                            # Perform triangulation by splitting quad into two tris, using the shortest diagonal
                            distance13 = self._distanceBetweenVerticesByIdx(vidx1, vidx3)
                            distance24 = self._distanceBetweenVerticesByIdx(vidx2, vidx4)

                            if distance13 > distance24:
                                face = [vidx1, vidx2, vidx4]
                                self._rawFaces.append(face)
                                face = [vidx3, vidx4, vidx2]
                                self._rawFaces.append(face)
                            else:
                                face = [vidx1, vidx3, vidx4]
                                self._rawFaces.append(face)
                                face = [vidx2, vidx3, vidx1]
                                self._rawFaces.append(face)

                        i = 1
                        while i < len(parts):

                            f = parts[i].split('/')

                            if len(f) > 1:
                                vidx = int(f[0]) - 1 # Vertex index

                                ti = f[1] # May be empty if no UV unwrap
                                if ti != "":
                                    tidx = int(ti) - 1 # Texture coordinate index
                                    texco = self._rawTexCo[tidx] # Actual texture coordinats, x/y
                                    self._rawVertexTexCo[vidx] = texco
                                    self.hasTexCo = True

                            i = i + 1

    def _createFacesNumpyArray(self, assumeQuads = False):

        numberOfFaces = len(self._rawFaces)

        vertsPerFace = 3
        if assumeQuads:
            vertsPerFace = 4

        # Create a two-dimensional int array with shape (numFace/vertsPerFace) and
        # fill it values from the wavefront obj. This will contain vert indices.

        self.faces = numpy.array( self._rawFaces, dtype=int ) # Values will be copied from self._rawFaces

        # Create a two-dimensional float array with shape (numFace/ 3 ) and
        # fill it with zeros. This will contain faces normals, but needs to
        # be recalculated.

        self.faceNormals = numpy.zeros( (numberOfFaces, 3), dtype=float )

    def _createVerticesNumpyArray(self):

        numberOfVertices = len(self._rawVertices)

        if numberOfVertices != len(self._rawVertexTexCo):
            raise ValueError("Not same number of elements in texco array")

        # Convert raw coords from wavefront into a 2d numpy array
        self.vertexCoords = numpy.array( self._rawVertices, dtype=float )

        # Create a two-dimensional float array with shape (numVerts/3) and
        # fill it with zeros. This will contain vertex normals.
        self.vertexNormals = numpy.zeros( (numberOfVertices, 3), dtype=float )

        # Create a two-dimensional float array with shape (numVerts/2) and
        # fill it with texture coordinates.
        self.vertexTexCo = numpy.array( self._rawVertexTexCo, dtype=float )

    def recalculateVertexNormals(self, assumeQuads = False):

        # Build a cache where we, per vertex, list which faces are relevant
        # for it. We need this in order to calculate the vertex normal later,
        # as an average of the face normals surrounding it
        if self._vertexBelongsToFaces is None:

            self._vertexBelongsToFaces = []

            numberOfFaces = len(self.faces)
            numberOfVertices = len(self.vertexCoords)

            vertsPerFace = 3
            if assumeQuads:
                vertsPerFace = 4

            currentVert = 0
            while currentVert < numberOfVertices:
                self._vertexBelongsToFaces.append([])
                currentVert = currentVert + 1

            currentFace = 0
            while currentFace < numberOfFaces:
                fv = self.faces[currentFace]
                currentVert = 0
                while currentVert < vertsPerFace:
                    vertexIndex = fv[currentVert]
                    self._vertexBelongsToFaces[vertexIndex].append(currentFace)
                    currentVert = currentVert + 1
                currentFace = currentFace + 1

        # Calculate vertex normals as an average of the surrounding face
        # normals.
        currentVert = 0

        zeroNormal = numpy.array([0.0, 0.0, 0.0], dtype=float)

        while currentVert < numberOfVertices:
            faces = self._vertexBelongsToFaces[currentVert]
            numberOfFaces = len(faces)
            currentNormal = numpy.array([0,0,0], dtype=float)
            currentFace = 0
            firstNormal = None
            while currentFace < numberOfFaces:
                fidx = faces[currentFace]
                fnormal = self.faceNormals[fidx]
                if firstNormal is None:
                    firstNormal = fnormal
                currentNormal = currentNormal + fnormal
                currentFace = currentFace + 1
            if numberOfFaces < 1:
                raise ValueError("Found a vertex (" + str(currentVert) + ") which did not belong to any face")

            averageNormal = currentNormal / numberOfFaces

            if numpy.array_equal(averageNormal, zeroNormal):
                print("WARNING: found zero vertex normal for vertex " + str(currentVert))
                averageNormal = firstNormal

            self.vertexNormals[currentVert] = self._unitVector(averageNormal)

            currentVert = currentVert + 1

    def recalculateFaceNormals(self):

        self._copyVertCoordsToCache()

        # Calculate the face normal from the first three vertices. This might produce
        # strange results if using quads. In that case we should probably triangulate and
        # and weight together the face normals of the resulting two tris.
        numberOfFaces = len(self.faces)

        currentFace = 0
        while currentFace < numberOfFaces:

            U = self._faceVertCache[currentFace][1] - self._faceVertCache[currentFace][0]
            V = self._faceVertCache[currentFace][2] - self._faceVertCache[currentFace][0]

            cross = numpy.cross(U,V)
            N = self._unitVector(cross)

            self.faceNormals[currentFace] = N
            currentFace = currentFace + 1

    def _unitVector(self, normal):

        # There is probably a numpy version of doing this calculation

        x = normal[0] * normal[0]
        y = normal[1] * normal[1]
        z = normal[2] * normal[2]

        magnitude = math.sqrt(x + y + z)

        if magnitude == 0.0:
            print("\nINVALID MAGNITUDE:\n")
            print("Normal was: " + str(normal))
            raise ValueError("Invalid magnitude")
            sys.exit(1)

        return normal / magnitude

    def _copyVertCoordsToCache(self, assumeQuads = False):

        # Create a two-dimensional float array with shape ( numFace / vertsPerFace / 3 ) and
        # fill it with the coordinates for each vertex participating  in each face
        #
        # This cache is used when recalculating face normals

        numberOfFaces = len(self.faces)

        vertsPerFace = 3
        if assumeQuads:
            vertsPerFace = 4

        if self._faceVertCache is None:
            self._faceVertCache = numpy.zeros( (numberOfFaces, vertsPerFace, 3), dtype=float )

        currentFace = 0
        while currentFace < numberOfFaces:
            currentVertex = 0
            while currentVertex < vertsPerFace:
                vertexIndex = self.faces[currentFace][currentVertex]
                self._faceVertCache[currentFace][currentVertex] = self.vertexCoords[vertexIndex]
                currentVertex = currentVertex + 1
            currentFace = currentFace + 1

    def getVertexArray(self):
        return self.vertexCoords

    def getVertexNormals(self):
        return self.vertexNormals

    def getVertexAndNormalArray(self):
        # This should possibly be cached
        return numpy.hstack( (self.vertexCoords, self.vertexNormals) )

    def getVertexAndNormalAndTexCoArray(self):
        # This should possibly be cached
        return numpy.hstack( (self.vertexCoords, self.vertexNormals, self.vertexTexCo) )

    def getFaceArray(self):
        return self.faces

    def debugVertices(self):

        vertices = self.getVertexAndNormalArray()

        for vertex in vertices:
            out = "["
            for i in vertex:
                out = out + " " + str(round(i,4))
            out = out + " ]"
            print(out)