#!/usr/bin/env python3

################################################################################
##                                                                            ##
##  This file is part of NCrystal (see https://mctools.github.io/ncrystal/)   ##
##                                                                            ##
##  Copyright 2015-2022 NCrystal developers                                   ##
##                                                                            ##
##  Licensed under the Apache License, Version 2.0 (the "License");           ##
##  you may not use this file except in compliance with the License.          ##
##  You may obtain a copy of the License at                                   ##
##                                                                            ##
##      http://www.apache.org/licenses/LICENSE-2.0                            ##
##                                                                            ##
##  Unless required by applicable law or agreed to in writing, software       ##
##  distributed under the License is distributed on an "AS IS" BASIS,         ##
##  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  ##
##  See the License for the specific language governing permissions and       ##
##  limitations under the License.                                            ##
##                                                                            ##
################################################################################

"""Script for extracting neutron scattering kernels from ENDF files and creating
corresponding .ncmat files. The created files might need a bit of manual editing
after creation, to correct the dummy material densities inserted, or (for
crystalline materials) to remove the densities and instead add sections with
information about crystal structure. For materials based on more than one
scattering kernel (e.g. heavy water, D2O, where S(alpha,beta) tables are needed
for both atoms), it might be necessary to manually combine two resulting files
into one.

Input files are read with the Python Nuclear Engineering Toolkit
(http://pyne.io/), which must be first installed.

Input files can for instance be downloaded from:

   https://www.nndc.bnl.gov/endf/b8.0/download.html

(download and open zip-file from the "Thermal Neutron Scattering Sublibrary")
"""

################################################################################################
####### Common code for all NCrystal cmdline scripts needing to import NCrystal modules ########
import sys
pyversion = sys.version_info[0:3]
_minpyversion=(3,6,0)
if pyversion < _minpyversion:
    raise SystemExit('Unsupported python version %i.%i.%i detected (needs %i.%i.%i or later).'%(pyversion+_minpyversion))
import os
import pathlib

def maybeThisIsConda():
    return ( os.environ.get('CONDA_PREFIX',None) or
             os.path.exists(os.path.join(sys.base_prefix, 'conda-meta')) )

def fixSysPathAndImportNCrystal( *, allowfail = False ):
    thisdir = pathlib.Path( __file__ ).parent
    def extract_cmake_pymodloc():
        p = thisdir / 'ncrystal-config'
        if not p.exists():
            return
        with p.open('rt') as fh:
            for i,l in enumerate(fh):
                if i==30:
                    break
                if l.startswith('#CMAKE_RELPATH_TO_PYMOD:'):
                    l = l[24:].strip()
                    return ( thisdir / l ) if l else None
    pml = extract_cmake_pymodloc()
    hack_syspath = pml and ( pml / 'NCrystal' / '__init__.py' ).exists()
    if hack_syspath:
        sys.path.insert(0,str(pml.absolute().resolve()))
    try:
        import NCrystal
    except ImportError:
        if allowfail:
            return
        msg = 'ERROR: Could not import the NCrystal Python module'
        if maybeThisIsConda():
            msg += ' (if using conda it might help to close your terminal and activate your environment again)'
        elif not hack_syspath:
            msg += ' (perhaps your PYTHONPATH is misconfigured)'
        raise SystemExit(msg)
    return NCrystal
################################################################################################
NC = fixSysPathAndImportNCrystal()

import warnings
import argparse

__pynecache=[None]
def import_pyne():
    global __pynecache
    if __pynecache[0] is not None:
        return __pynecache[0]

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")#silence annoying pyne.endf QAWarning's.
            import pyne.endf
    except ImportError:
        raise SystemExit('ERROR: Could not "import pyne.endf". You probably need to install PyNE (see http://pyne.io/)')

    #Monkey-patch pyne.endf reader, which was ignoring optional sections with teff
    #curves of non-principal atoms (see
    #https://github.com/pyne/pyne/issues/1209). The fix was accepted upstream, so
    #eventually we can remove this monkey-patching once we think all users will have
    #recent enough PyNE installations:
    _orig_rti = pyne.endf.Evaluation._read_thermal_inelastic
    def _fix_rti(_self):
        _orig_rti(_self)
        inel=_self.thermal_inelastic
        B=inel['B']
        for ii in range(len(B)//6-1):
            if B[6*(ii+1)]==0.0:
                if 'teff_%i'%(ii+1) in inel:
                    #If upstream already includes fix, do nothing.
                    continue
                _, teff = _self._get_tab1_record()
                inel['teff_%i'%(ii+1)] = teff#Store as teff_1, teff_2, etc. for now.
    pyne.endf.Evaluation._read_thermal_inelastic=_fix_rti
    __pynecache[0] = pyne
    return pyne

def elementZToName(z):
    d = NCrystal.atomDB(z,throwOnErrors=False)
    return d.elementName() if d else None

def guessElementFromMass(mass_amu):
    elem_data=[ (abs(d.averageMassAMU()-mass_amu),d) for d in [NCrystal.atomDB(z,throwOnErrors=False) for z in range(1,120)] if d ]
    elem_data.sort()
    print (f'Guessing mass={mass_amu:.8}u is element {elem_data[0][1].elementName()}')
    return elem_data[0][1].elementName()

def guessElementName(ZA,mass_amu):
    Z=ZA//1000
    if ZA==155:
        return 'Y'#YH2, whatever it is
    if ZA==140:
        return 'H'#benzene...  C6H6. Guessing H is dominant. But tsl-benzene.endf is actually a compound SAB...
    if ZA==1002:
        return 'D'
    if ZA in (133,134):
        return 'H'#solid or liquid methane (CH4 happens to have mass 12+1*4=16 close to oxygen, so can't guess from mass)
    if ZA in (112,113):
        return 'D'#ortho-D or para-D
    guess_by_mass = guessElementFromMass(mass_amu)
    guess_by_z = elementZToName(Z) if (Z>0 and Z<120) else None
    if guess_by_z:
        if guess_by_z!=guess_by_mass:
            raise SystemExit(f"ERROR: Mass and Z guess gives different element name (Z={Z}=>'{guess_by_z}', mass={mass_amu}amu=>'{guess_by_mass}')")
        return guess_by_z
    if Z==0:
        return guess_by_mass
    raise ValueError(f'Could not determine element name from ZA={ZA}, mass={mass_amu}amu')

#Custom bare-bones reader, so we can extract the data header as-is in the file:
def parse_endf6(src,name=None,selectfct=None):
    """Parse TSL data from ENDF-6 file according to ENDF manual. Select function can
    used to determine which mat,mf,mt lines to read. E.g. "selectfct=lambda
    mat,mf,mt:mf==7". Remember to select (mf,mt)=(1,451) for the description at
    the start of each material.

    """
    if isinstance(src,str) or isinstance(src,bytes) or hasattr(src,'__fspath__'):
        p=pathlib.Path(src).expanduser()
        with p.open('rt') as srcstream:
            return parse_endf6(srcstream,name=p.name)
    #according to the manual section 0.6, each line is broken down as follows:
    #Columns 1-66 content, 67-70 MAT number, 71-72 MF, 73-75 MT, 76-80 [optional] NS, 81+ undefined
    #Content can either be one 66 char long text field, or there are 6 fields of 11 chars.

    errprefix = '' if not name else f'in {name}'
    def require(check,errmsg):
        if not check:
            raise ValueError(f'{errprefix}: {errmsg}')

    def parse_endf6_line(l):
        require(len(l)>=75,"line too short")
        content=l[0:66]
        mat=int(l[66:70])
        mf=int(l[70:72])
        mt=int(l[72:75])
        return (content,mat,mf,mt)

    content,mat,mf,mt = parse_endf6_line(next(src))
    require((mat,mf,mt)==(1,0,0),'Not an ENDF-6 file?')
    global_header=content
    d={}#material -> material info
    for l in src:
        #print(l)
        content,mat,mf,mt = parse_endf6_line(l)
        if mat<1:
            continue#Ignore dummy lines (section dividers, etc.)
        if selectfct and not selectfct(mat,mf,mt):
            continue
        if not mat in d:
            d[mat]=dict()
        dd=d[mat]
        if not mf in dd:
            dd[mf]={}
        ddd=dd[mf]
        if not mt in ddd:
            ddd[mt]=[content]
        else:
            ddd[mt]+=[content]
    return dict(global_header=global_header,name=name),d

def extract_and_format_endf_file_header(filename):
    custom_parse = parse_endf6(filename,selectfct=lambda mat,mf,mt : (mf,mt)==(1,451))
    mats=custom_parse[1]
    assert len(mats)==1,("File contains more than one material number which is not supported by the"
                         +" present script (please contact NCrystal developers if you really need this).")
    material_number,material_info = next(iter(mats.items()))
    return dict( filename = custom_parse[0]['name'],
                 material_number = material_number,
                 global_header = custom_parse[0]['global_header'],
                 infosection=material_info[1][451] )

def parse_endf_file(filename):
    print(f"Attempting to load ENDF file {filename}...")
    result={}
    warnings_=[]
    pyne = import_pyne()
    data=pyne.endf.Evaluation(filename)
    data.read()
    result['input_file']=pathlib.Path(filename)
    result['pynedata']=data
    print('Performing a few sanity checks...')
    if not hasattr(data,'thermal_inelastic') or not 'B' in data.thermal_inelastic or not 'scattering_law' in data.thermal_inelastic:
        raise SystemExit('Error: File does not appear to contain any inelastic neutron scattering kernel.')

    ti=data.thermal_inelastic

    if ti['ln(S)']:
        raise SystemExit('The ln(S) flag is enabled. For the reasons given in the ENDF manual just before'
                         +'section 7.4.1, it is not safe to simple convert these to S-values. We need ncmat to support a logsab_scaled option first.'
                         +' Another (easier?) alternative would be to use mpmath module here and calculate unscaled S-values and produce an ncmat file with S=...')

    assert ti['temperature_used'] in ('actual','0.0253 eV')
    if ti['temperature_used'].strip()=='actual':
        alphabetagrid_T0=None
    else:
        alphabetagrid_T0 = float(ti['temperature_used'].split('eV',1)[0].strip())/NCrystal.constant_boltzmann#Must stretch betagrid values relative to this

    #print(data.target)
    element_name=guessElementName(data.target['ZA'],data.target['mass'])
    zsymam=data.target['zsymam'].strip()
    if not zsymam in ['s-ch4','l-ch4','ortho-d','para-d','BENZINE'] and not 'graphit' in zsymam.lower():
        assert element_name in zsymam,f"{element_name} not in {zsymam}"
    result['element_name_principal']=element_name

    lenB = len(ti['B'])
    B  = [None] + list(ti['B'])#B-array, with fortran indexing (for easier reference with ENDF manual):
    A0 = float(B[3]) #Needed to unscale
    if B[1]==0.0:
        raise SystemExit('Material has no principal scatterers - thus no S(alpha,beta)')
    suggested_emax = B[4]#In principle... but in practice not sure if this holds! Also, data.info['energy_max'] gives a different number (5.0)?!?
    count_principal = B[6]
    num_non_principal = ti['num_non_principal']

    non_principal_data = []
    for i in range(num_non_principal):
        offset = 6*(i+1)
        assert lenB >= 6*(i+2)
        a1=B[offset+1]#0.0: SCT (short collision time approx), 1.0: free-gas, 2.0: diffusive motion.
        if a1!=1.0:
            warnings_+=['WARNING: A non-free-gas model is proposed for the first non-principal atom']
            warnings_+=['         in the original file. This is currently not handled by NCrystal.']
            print(warnings_[-2])
            print(warnings_[-1])
        effective_mass = B[offset+3]#mass in units of neutron mass
        count = B[offset+6]#count per molecule or unit cell
        theElementName = guessElementFromMass(effective_mass*NCrystal.const_neutron_mass_amu)
        non_principal_data+=[[theElementName,count,effective_mass]]

    count_total = count_principal + sum(count_ for elemname,count_,effmass in non_principal_data)
    format_fraction = lambda c,ctot : f'{c:.14g}/{ctot:.14g}' if c!=ctot else '1'
    for e in non_principal_data:
        elementName,thecount,effective_mass=e
        e += [format_fraction(thecount,count_total)]

    header = extract_and_format_endf_file_header(filename)
    result['header']=header
    result['warnings']=warnings_
    result['non_principal_data']=non_principal_data
    result['count_total']=count_total
    result['count_principal']=count_principal
    result['fraction_str_principal']=format_fraction(count_principal,count_total)
    result['suggested_emax']=suggested_emax

    ndatablock = len(ti['teff'].x)
    result_datablocks=[]
    result['result_datablocks'] = result_datablocks
    for idatablock in range(ndatablock):
        result_datablock={}
        result_datablocks += [result_datablock]
        temperature=ti['teff'].x[idatablock]
        temperature_effective=ti['teff'].y[idatablock]
        result_datablock['T']=temperature
        result_datablock['Teff']=temperature_effective
        sabt=ti['scattering_law'][:,:,idatablock] #indexing t,a,b
        sab=sabt.T
        result_datablock['sab']=sab
        #We must .copy() alpha/beta grids due to '*=' operators below:
        alphagrid = ti['alpha'].copy()*A0#undo funny ENDF alpha definition which divides by A0
        betagrid = ti['beta'].copy()
        if alphabetagrid_T0 is not None:
            #Undo funny ENDF scaling of grids at each temperature
            betagrid *= (alphabetagrid_T0/temperature)
            alphagrid *= (alphabetagrid_T0/temperature)
        result_datablock['alphagrid']=alphagrid.copy()
        result_datablock['betagrid']=betagrid.copy()

        assert ti['symmetric'] == bool(betagrid[0]>=0.0),"inconsistencies detected"

        warningsfmt=''
        if warnings_:
            warningsfmt='{}\n#'.format('\n#  ----> '.join(['']+warnings_))
    return result

def format_endf_block_as_ncmatdyninfo_for_principal_element(parsed_endf_data,temperature,fraction_str=None):
    elem_name = parsed_endf_data["element_name_principal"]

    #Find block by temperature:
    temperature_exact,block_idx =  list(sorted((abs(temperature-_["T"]),_["T"],idx) for idx,_ in enumerate(parsed_endf_data['result_datablocks'])))[0][1:]
    assert abs(temperature_exact-temperature)<1e-6
    temperature=temperature_exact

    res =   '@DYNINFO\n'
    res += f"  # Scattering kernel for {elem_name} @ {temperature}K extracted from {parsed_endf_data['input_file'].name}\n"
    for w in parsed_endf_data['warnings']:
        res += f'\n  #  ----> {w}'
    res += '  #\n'
    res += '  # For reference the description embedded in input ENDF file (section MF1,MT451) is repeated here:\n'
    res += '  #\n'
    header=parsed_endf_data['header']
    for l in [header['global_header']]+header['infosection']:
        res+=f'  #     --->{(" "+l).rstrip()}\n'

    datablock = parsed_endf_data['result_datablocks'][block_idx]
    res+=f'  element  {elem_name}\n'
    if fraction_str is None:
        fraction_str = parsed_endf_data["fraction_str_principal"]
    else:
        assert parsed_endf_data["fraction_str_principal"]=='1'
    res+=f'  fraction {fraction_str}\n'
    res +='  type     scatknl\n'
    res+=f'  temperature {temperature:.10} #NB: ENDF file specified "effective temperature" as {datablock["Teff"]:.10}K\n'
    #NB: Not suggesting an emax value, since it seems to be unreliable! Same for B[4] from above...
    #res+=f"  egrid {parsed_endf_data['pynedata'].info['energy_max']}#Value (in eV) as specified in source ENDF file.\n"#argh... e.g. D2O@300K shows this can't be trusted!!!
    res += NCrystal.formatVectorForNCMAT('alphagrid',datablock['alphagrid'])
    res += NCrystal.formatVectorForNCMAT('betagrid',datablock['betagrid'])
    res += NCrystal.formatVectorForNCMAT('sab_scaled',datablock['sab'])
    return res

def _parseArgs():
    descr="""
Converts neutron scattering kernels ("thermal scattering laws, tsl") found in
ENDF files to NCMAT format (specifically @DYNINFO sections of NCMAT files).

Note that this script exists for the benefit of experts only. Most users are
recommended to simply download and use one of the pre-converted files shipped
with NCrystal or found at:

    https://github.com/mctools/ncrystal-extra/tree/master/data

Only information for NCMAT @DYNINFO sections are extracted from ENDF files, so
it is recommended to provide the initial parts of the target NCMAT file in a
separate file specified via the --filehdr parameter. That file should include
both initial comments for the file as well as either @DENSITY or
@CELL/@SPACEGROUP/@ATOMPOSITIONS/@DEBYETEMPERATUE sections, depending on whether
or not the material is crystalline. Note that comments describing the origin of
the @DYNINFO sections extracted from the ENDF files will be inserted in the
relevant @DYNINFO sections. The --filehdr should include the string
'<<STDNOTICE>>' on a separate line, which will be expanded to a notice about
availability of files for other temperatures.

If the --filehdr argument is not provided, a dummy header will be
automatically generated, with a fake @DENSITY section, using the density of
water (1g/cm3). This allows the file to be technically valid, but will of
course result in incorrect physics in case the file is used for simulations.

As each ENDF file contains a list of temperatures for which scattering kernels
are available, one NCMAT file will be generated for each such temperature. The
resulting filename will be <basename>_T<temperature>K.ncmat, where <basename>
will default to the input filename unless --outbn is provided.

The script internally relies on the Python Nuclear Engineering Toolkit
(http://pyne.io/) for ENDF parsing, so this must be installed on the system.

Input ENDF files can for instance be downloaded from:

   https://www.nndc.bnl.gov/endf/b8.0/download.html

(download and open zip-file from the "Thermal Neutron Scattering Sublibrary").

"""

    parser = argparse.ArgumentParser(description=descr,
                                     formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('ENDFFILE',help="Primary ENDF file with tsl data.",type=str)
    parser.add_argument('--secondary',metavar='ENDFFILE2:fraction',
                        type=str,
                        help="ENDF file for secondary element, along with the fraction of that element")
    parser.add_argument("--filehdr",help="File containing initial part of generated files.",type=str)
    parser.add_argument("--outbn",type=str,help="Basename of generated ncmat files.")
    parser.add_argument("--ignoretemp",type=float,nargs='+',metavar='T',
                        help="If some temperature blocks in input should be ignored, provide the temperature values here (kelvin).")

    def to_path(parser,fn):
        _ = pathlib.Path(fn)
        if not _.exists():
            parser.error(f'File not found: {_}')
        return _


    args=parser.parse_args()
    args.ENDFFILE = to_path(parser,args.ENDFFILE)
    args.fraction1='1'
    args.fraction2=None
    if not args.secondary:
        args.ENDFFILE2=None
    else:
        if not args.secondary.count(':')==1:
            parser.error('Argument to --secondary must contain exactly one semicolon (:)')
        sn,fr2 = args.secondary.split(':',1)
        args.secondary = True
        if fr2.count('/')==1:
            a,b = [int(e) for e in fr2.split('/')]
            if b<=0 or not ( 0 < a < b ):
                parser.error('Invalid fraction specified in --secondary')
            args.fraction1=f'{b-a}/{b}'
            args.fraction2=f'{a}/{b}'
        else:
            args.fraction2=fr2
            if not ( 0.0 < float(fr2) < 1.0):
                parser.error('Invalid fraction specified in --secondary')
            args.fraction1=f'{1.0-float(fr2):10}'
        args.ENDFFILE2 = to_path(parser,sn)
    if args.filehdr:
        _filehdrtxt=to_path(parser,args.filehdr).read_text()
    else:
        _filehdrtxt=f"""
NCMAT v2
#
# Material with scattering kernel extracted from ENDF file.
#
<<STDNOTICE>>
#
@DENSITY
  1 g_per_cm3 #FIX{""}ME. Dummy number!!! (update or add unit cell sections and then remove @DENSITY section)
"""
    args.filehdr = [l.rstrip() for l in _filehdrtxt.strip().splitlines()]

    if not any('<<STDNOTICE>>' in l for l in args.filehdr):
        parser.error('ERROR: File specified with --filehdr should contain the string "<<STDNOTICE>>" on a separate line (just before the first data section)')

    if not args.outbn:
        if args.ENDFFILE2:
            parser.error('--outbn is required when providing two input files')
        else:
            args.outbn=args.ENDFFILE.name
            if args.outbn.endswith('.endf'):
                args.outbn=args.outbn[0:-5]
        assert args.outbn

    return args

def genstdheaders():
    def do_write(fn,suggestedcmdopts,content):
        pathlib.Path(fn).write_text(content)
        print(f'Wrote {fn}')
        progname=os.path.basename(sys.argv[0])
        print( f'  -> Suggested conversion cmd: {progname} {suggestedcmdopts}')

    fixstr='NO'+'COMM'+'IT/FI'+'XME'
    do_write('hdrH2O.txt',
             'tsl-HinH2O.endf --outbn LiquidWaterH2O --filehdr=hdrH2O.txt --ignoretemp 650 800',
             f"""NCMAT v2
#
# Water (H2O) based on ENDF/B-VIII.0 scattering kernels by J.I. Marquez Damian,
# et. al. (see below for references). Please check the @DENSITY carefully below, as
# the density of liquid water depends on the material pressure.
<<STDNOTICE>>
#
# In case your local installation does not provide all these files, they can be
# downloaded from: https://github.com/mctools/ncrystal-extra/tree/master/data
#
@DENSITY
#{fixstr}: Edit to leave just the relevant line (and the relevant comment for the reference):
#NB: Density depends on pressure (value here provided by https://www.thermexcel.com/english/tables/eau_atm.htm):
#Uncomment for T=283.6K: 0.99973 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=293.6K: 0.99820 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=300.0K: 0.99663 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=323.6K: 0.98781 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=350.0K: 0.97355 g_per_cm3 #(at 1atm pressure)
#NB: Density depends on pressure (value here provided by engineeringtoolbox.com):
#Uncomment for T=373.6K: 0.95805 g_per_cm3 #(at saturation pressure of 1.0171atm)
#Uncomment for T=400.0K: 0.9376 g_per_cm3 #(at saturation pressure of 2.41atm)
#Uncomment for T=423.6K: 0.9162 g_per_cm3 #(at saturation pressure of 4.77atm)
#Uncomment for T=450.0K: 0.8903 g_per_cm3 #(at saturation pressure of 9.21atm)
#Uncomment for T=473.6K: 0.8645 g_per_cm3 #(at saturation pressure of 15.48atm)
#Uncomment for T=500.0K: 0.8315 g_per_cm3 #(at saturation pressure of 26.03atm)
#Uncomment for T=523.6K: 0.7979 g_per_cm3 #(at saturation pressure of 39.55atm)
#Uncomment for T=550.0K: 0.7554 g_per_cm3 #(at saturation pressure of 60.39atm)
#Uncomment for T=573.6K: 0.7116 g_per_cm3 #(at saturation pressure of 85.29atm)
#Uncomment for T=600.0K: 0.6499 g_per_cm3 #(at saturation pressure of 121.8atm)
#Uncomment for T=623.6K: 0.5709 g_per_cm3 #(at saturation pressure of 164.06atm)
1 g_per_cm3 #DUMMY {fixstr}
#{fixstr}Also pick appropriate egrid max and move down to @DYNINFO(H) section:
#Uncomment and place below for T=283.6K: egrid 3.92 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=293.6K: egrid 4.02 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=300.0K: egrid 4.04 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=323.6K: egrid 4.3 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=350.0K: egrid 4.65 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=373.6K: egrid 4.9 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=400.0K: egrid 5.4 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=423.6K: egrid 5.65 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=450.0K: egrid 5.95 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=473.6K: egrid 6.2 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=500.0K: egrid 6.55 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=523.6K: egrid 6.85# Value tuned manually by T. Kittelmann
#Uncomment and place below for T=550.0K: egrid 7.2 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=573.6K: egrid 7.6 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=600.0K: egrid 7.95 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=623.6K: egrid 8.35 # Value tuned manually by T. Kittelmann
""")

    do_write('hdrD2O.txt',
             'tsl-DinD2O.endf --secondary tsl-OinD2O.endf:1/3 --outbn LiquidHeavyWaterD2O --filehdr=hdrD2O.txt --ignoretemp=650',
             f"""NCMAT v2
#
# Heavy water (D2O) based on ENDF/B-VIII.0 scattering kernels by J.I. Marquez
# Damian, et. al. (see below for references). Please check the @DENSITY
# carefully below, as the density of liquid water depends on the material
# pressure.
<<STDNOTICE>>
#
# In case your local installation does not provide all these files, they can be
# downloaded from: https://github.com/mctools/ncrystal-extra/tree/master/data
#
@DENSITY
#{fixstr}: Edit to leave just the relevant line (and the relevant comment for the reference):
#NB: Density depends on pressure (value here provided by engineeringtoolbox.com):
#Uncomment for T=283.6K: 1.10586 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=293.6K: 1.10526 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=300.0K: 1.10404 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=323.6K: 1.09547 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=350.0K: 1.08037 g_per_cm3 #(at 1atm pressure)
#Uncomment for T=373.6K: 1.06312 g_per_cm3 #(at 1atm pressure)
#NB: Density depends on pressure (value here from P.G., MacMillan, et. al. (1981),
#    "Tables of thermodynamic properties of heavy water in SI units"):
#Uncomment for T=400.0K:  1.04019 g_per_cm3 #(at saturation pressure of 2.34atm)
#Uncomment for T=423.6K:  1.01678 g_per_cm3 #(at saturation pressure of 4.65atm)
#Uncomment for T=450.0K:  0.987245 g_per_cm3 #(at saturation pressure of 9.08atm)
#Uncomment for T=473.6K:  0.957643 g_per_cm3 #(at saturation pressure of 15.4atm)
#Uncomment for T=500.0K:  0.920446 g_per_cm3 #(at saturation pressure of 26.1atm)
#Uncomment for T=523.6K:  0.883010 g_per_cm3 #(at saturation pressure of 39.7atm)
#Uncomment for T=550.0K:  0.834808 g_per_cm3 #(at saturation pressure of 60.9atm)
#Uncomment for T=573.6K:  0.783788 g_per_cm3 #(at saturation pressure of 86.3atm)
#Uncomment for T=600.0K:  0.712194 g_per_cm3 #(at saturation pressure of 123.6atm)
#Uncomment for T=623.6K:  0.621322 g_per_cm3 #(at saturation pressure of 166.9atm)
1 g_per_cm3 #DUMMY {fixstr}
#{fixstr}Also pick appropriate egrid max and move down to @DYNINFO(H) section:
#Uncomment and place below for T=283.6K D: egrid 3.6 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=283.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=293.6K D: egrid 3.7 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=293.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=300.0K D: egrid 3.8 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=300.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=323.6K D: egrid 4.1 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=323.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=350.0K D: egrid 4.4 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=350.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=373.6K D: egrid 4.7 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=373.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=400.0K D: egrid 5.1 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=400.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=423.6K D: egrid 5.3 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=423.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=450.0K D: egrid 5.55 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=450.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=473.6K D: egrid 6.1 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=473.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=500.0K D: egrid 6.4 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=500.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=523.6K D: egrid 6.8 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=523.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=550.0K D: egrid 7.1 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=550.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=573.6K D: egrid 7.4 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=573.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=600.0K D: egrid 7.6 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=600.0K O: egrid 10.0 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=623.6K D: egrid 8.1 # Value tuned manually by T. Kittelmann
#Uncomment and place below for T=623.6K O: egrid 10.0 # Value tuned manually by T. Kittelmann
""")
    print(f"\nNOTICE: Some files produced by above commands contain FIX{'ME'}s (might need density values updated for given temperature)")
    print("\nNOTICE: It might also be worth investigating each file in order to manually provide kernel egrid max via egrid keyword.")

if __name__=='__main__':
    if '--genstdheaders' in sys.argv[1:]:
        #Hidden option to prepare some conversions as used for official NCrystal files.
        genstdheaders()
        sys.exit(0)
    args=_parseArgs()

    p1=parse_endf_file(args.ENDFFILE)
    p2=parse_endf_file(args.ENDFFILE2) if args.ENDFFILE2 else None

    temperatures1 = [ _["T"] for _ in p1['result_datablocks'] ]
    if p2:
        temperatures2 = [ _["T"] for _ in p2['result_datablocks'] ]
        temperatures_combined = [ t1 for t1 in temperatures1 if min(abs(t1-t2) for t2 in temperatures2)<1e-5 ]
        if len(temperatures_combined)<len(temperatures1):
            print(f"WARNING: {len(temperatures1)-len(temperatures_combined)} temperature blocks in {args.ENDFFILE} could not be used due to missing data in other file!")
            if len(temperatures_combined)<len(temperatures2):
                print(f"WARNING: {len(temperatures2)-len(temperatures_combined)} temperature blocks in {args.ENDFFILE} could not be used due to missing data in other file!")
    else:
        temperatures_combined = temperatures1

    temperatures_combined=sorted(list(set(temperatures_combined)))
    for t in (args.ignoretemp or []):
        removed=False
        for _ in temperatures_combined:
            if abs(t-_)<1e-6:
                temperatures_combined.remove(_)
                print (f"Ignoring T={_}K as requested")
                removed=True
                break
        if not removed:
            raise SystemExit(f'ERROR: Asked to ignore T={t}K, but no such temperature was found in the input')


    if p2 and (p1['non_principal_data'] or p2['non_principal_data']):
        raise SystemExit("ERROR: When providing two ENDF files as input both must files must be without non-principal"
                         +" elements! If you know how to handle this, You can try to convert separately and combine the"
                         +" resulting .ncmat files manually.")

    for t in temperatures_combined:
        fn=pathlib.Path(f'{args.outbn}_T{t}K.ncmat')
        with fn.open('wt') as fh:
            print(f'   -> Writing {fn}')
            fh.write('NCMAT v2\n')
            stdnotice = f'#\n# Notice: This NCMAT file is valid at T={t}K only.'
            if len(temperatures_combined)>1:
                stdnotice+=' Other files alternatively provide\n'
                stdnotice+='# the same material at temperatures:\n#\n'
                nperline=7
                t_to_write=list(_ for _ in temperatures_combined if _!=t)
                for i in range(0,len(t_to_write),nperline):
                    stdnotice += ('#       '+' '.join((f'{_}K' for _ in t_to_write[i:i+nperline]))+'\n')
            for l in args.filehdr:
                if l.startswith('NCMAT '):
                    continue
                if stdnotice and '<<STDNOTICE>>' in l:
                    l=l.replace('<<STDNOTICE>>',stdnotice)
                    stdnotice=''
                fh.write(l if l.endswith('\n') else f'{l}\n')
            if stdnotice:
                fh.write(stdnotice)
            for elementName,count_,effmass,fraction_str in p1['non_principal_data']:
                assert not p2
                fh.write('@DYNINFO\n')
                fh.write(f'  element  {elementName}\n')
                fh.write(f'  fraction {fraction_str}\n')
                fh.write('  type     freegas\n')
            fh.write(format_endf_block_as_ncmatdyninfo_for_principal_element(p1,t,
                                                                             args.fraction1 if p2 else None))
            if p2:
                fh.write(format_endf_block_as_ncmatdyninfo_for_principal_element(p2,t,args.fraction2))
        print('   -> Testing that NCrystal can load this file')
        NCrystal.createScatter(f'{fn};dcutoff=0.8')
    print("All done.")