import os
import sys
import shutil
import warnings
import subprocess
from functools import lru_cache

import numpy as np
import matplotlib
%matplotlib inline
import pandas as pd
import matplotlib.pyplot as plt

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

import sys
sys.path.insert(0, '/home/scopatz/pyne/src')
import decaygen

exp2 = np.exp2

H5NAME = 'origen-benchmark-py{0}k.h5'.format(sys.version_info[0])
mats = MaterialLibrary(H5NAME)

t9 = origen22.parse_tape9('TAPE9.INP')
for key in list(t9.keys()):
    if key not in [1, 2, 3]:
        del t9[key]
        
def t9_half_life(nuc):
    nuczz = nucname.zzaaam(nuc)
    hl = None
    for i in [1, 2, 3]:
        hl = t9[i]['half_life'].get(nuczz, None)
        if hl is not None:
            break
    return hl

@lru_cache()
def hl_relerr(nuc):
    """Half-life relative error."""
    dhl = data.half_life(nuc) or np.inf
    ohl = t9_half_life(nuc) or np.inf
    if np.isinf(ohl) and np.isinf(dhl):
        return 0.0
    hlre = np.abs(ohl - dhl) * 2 / (ohl + dhl)
    return np.inf if np.isnan(hlre) else hlre

import json
t9hls = {}
for i in [1, 2, 3]:
    t9hls.update({nucname.id(int(nuc)): v for nuc, v in t9[i]['half_life'].items()})
with open('o2hls.json', 'w') as f:
    json.dump(t9hls, f, sort_keys=True)

metastable_blacklist = {
    771940001,  # have 2+ metastables that ORIGEN lumps together, or
    340830001,  # ...
    350830000,  # decay to metastables without reason.
    340830000,  # ...
    481170000,  # ...
    461110000,  # ...
    441030000,  # ...
    471110001,  # missing branch in origen
    501130001,  # ...
}
nuc_keys = {}
for key in mats.keys():
    _, nuc, t = key.split('_')
    nuc = int(nuc)
    if nuc in metastable_blacklist:
        continue
    chains = decaygen.genchains([(nuc,)])
    maxrelerr = max([max(list(map(hl_relerr, c))) for c in chains])
    if maxrelerr > 0.1:
        continue
    t = int(t)
    if nuc not in nuc_keys:
        nuc_keys[nuc] = []
    nuc_keys[nuc].append(key)
sfunc = lambda x: int(x.rsplit('_', 1)[1])
for keys in nuc_keys.values():
    keys.sort(key=sfunc)

def mat_to_series(x):
    xmass = x.mult_by_mass()
    s = pd.Series(list(x.values()), list(map(int, x.keys())))
    return s

def matdiff(x, y):
    if not isinstance(x, pd.Series):
        x = mat_to_series(x)
    if not isinstance(y, pd.Series):
        y = mat_to_series(y)
    return (x - y).abs() * (2 / (x + y))

idx = []
err = []
for nuc, keys in list(nuc_keys.items()):
    fresh = mats[keys[0]]
    for matkey in keys[1:]:
        mat = mats[matkey]
        if (mat.mass < 0.99) or (mat.mass > 1.1):
            # Origen lost too much mass
            continue
        t = mat.metadata['decay_time']
        dmat = fresh.decay(t)
        oser = mat_to_series(mat)
        dser = mat_to_series(dmat)
        md = matdiff(oser, dser)
        md = md.fillna(value=0.0)
        mask = (dser > 1e-3) | (oser > 1e-3)
        mdmasked = md.where(mask)
        maxmd = mdmasked.max()
        child = mdmasked.argmax()
        if np.isnan(child):
            continue
        childmass = max(dser[child], oser[child])
        weightederr = maxmd * childmass
        idx.append(matkey)
        err.append((nuc, t, maxmd, child, childmass, weightederr))
df = pd.DataFrame(err, index=idx, columns=['nuc', 't', 'relerr', 'child', 'childmass', 'weightederr'])

def nucsummary(nuc, i=None):
    print('Nulcide:', nucname.name(nuc))
    print("origen half life:", t9_half_life(nuc))
    print("data.half_life():", data.half_life(nuc))
    print('relerr half life:', hl_relerr(nuc))
    print()
    subdf = df[df['nuc'] == nuc]
    plt.figure(figsize=(10, 6))
    subdf.plot(x='t', y='relerr', logx=True, label=nucname.name(int(nuc)))
    plt.legend(loc=0)
    plt.ylabel('t [sec]')
    plt.ylabel('Relative Error')
    if i is not None:
        okey = 'origen_{0}_{1}'.format(nuc, i)
        fkey = 'fresh_{0}_0'.format(nuc)
        omat = mats[okey]
        fresh = mats[fkey]
        t = omat.metadata['decay_time']
        dmat = fresh.decay(t)
        md = matdiff(omat, dmat)
        print(omat)
        print()
        print(dmat)
        print()
        print(md)
    print(subdf.to_latex())
    return subdf

df.sort(['relerr', 'childmass', 'nuc'], ascending=[False, False, True])

(df.relerr > 0.5).sum() / len(df)

(df.weightederr > 0.5).sum() / len(df)

((df.relerr > 0.5) & (df.weightederr > 0.5)).sum() / len(df)

len(df)

N = 20
plt.figure(figsize=(10, 6))
for i, (key, grp) in enumerate(df.groupby(['nuc'])):
    if i > N:
        break
    nn = nucname.name(int(key))
    p = grp.plot(x='t', y='relerr', logx=True, label=nn)
    plt.xlabel('t [sec]')
    plt.ylabel('Relative Error')
    p.legend(p.lines, (nn,), loc=0)
    plt.savefig('decay-relative-error-{0}.eps'.format(nn))

nucsummary(420990000, 1)

nucsummary(491140000, 3)

nucsummary(380910000, 1)

df.relerr.plot(kind='hist', bins=10000, loglog=True, figsize=(10, 6))
plt.xlabel('Relative Error, on range [0, 2]')
plt.ylabel('Number of Calculations, Total = {0}'.format(len(df)))
plt.savefig('decay-relative-error-histogram.eps')