import os
import sys
import shutil
import warnings
import subprocess

import numpy as np
import matplotlib
%matplotlib inline

from pyne.utils import QAWarning
warnings.simplefilter('ignore', QAWarning)
from pyne import nucname
from pyne import data
from pyne import origen22
from pyne.material import Material, MaterialLibrary

H5NAME = 'origen-benchmark-py{0}k.h5'.format(sys.version_info[0])
DAYS_PER_SEC = 1.0 / (24.0 * 3600.0)

t9 = origen22.parse_tape9('TAPE9.INP')
nucs = set()
for i in [1, 2, 3]:
    nucs |= set(map(nucname.zzaaam_to_id, map(int, t9[i]['_cards']['f0'])))
nucs.discard(162500000)
nlb = tuple([n for n in t9.keys() if n not in [1, 2, 3]])
del t9
nucs = sorted(nucs)
print("# of nucs: ", len(nucs))

def decay_origen(nuc, t, fresh=None, nlb=(219, 220, 221)):
    """Decays 1 kg of a nuclide for time t using ORIGEN v2.2."""
    if fresh is None:
        fresh = Material({nuc: 1.0})
    if os.path.isdir('o22tmp'):
        shutil.rmtree('o22tmp')
    os.mkdir('o22tmp')
    shutil.copy('TAPE9.INP', 'o22tmp/TAPE9.INP')
    origen22.write_tape4(fresh, outfile='o22tmp/TAPE4.INP')
    origen22.write_tape5_decay(t*DAYS_PER_SEC, outfile='o22tmp/TAPE5.INP', xsfpy_nlb=nlb, cut_off=1e-300)
    subprocess.check_call('o2_therm_linux.exe', cwd='o22tmp')
    try:
        t6 = origen22.parse_tape6('o22tmp/TAPE6.OUT')
    except RuntimeError:
        return None
    mat = t6['materials'][-1]
    shutil.rmtree('o22tmp')
    return mat

%%time
fkey = 'fresh_{0}_0'
okey = 'origen_{0}_{1}'
dkey = 'decay_{0}_{1}'
mats = MaterialLibrary()
decades = np.logspace(-2, 2, 5)
for n, nuc in enumerate(nucs):
    hl = data.half_life(nuc)
    if np.isnan(hl):
        continue
    times = decades if np.isinf(hl) else hl * decades 
    fresh = Material({nuc: 1.0})
    fresh.metadata['decay_time'] = 0.0
    mats[fkey.format(nuc)] = fresh
    for i, t in enumerate(times, 1):
        print("Running:", nuc, n, t)
        t = float(t)
        o2mat = decay_origen(nuc, t, fresh, nlb=nlb)
        if o2mat is None:
            continue
        o2mat.metadata['decay_time'] = t
        mats[okey.format(nuc, i)] = o2mat
if os.path.isfile(H5NAME):
    os.remove(H5NAME)
mats.write_hdf5(H5NAME)