#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import time
import ctypes
import tempfile
from functools import reduce
import numpy
import h5py
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.cc import ccsd
from pyscf.cc import _ccsd

# t2,l2 as ijab

def kernel(mycc, eris, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    cput0 = (time.clock(), time.time())
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mycc.stdout, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if l1 is None: l1 = t1
    if l2 is None: l2 = t2

    nocc, nvir = t1.shape
    saved = make_intermediates(mycc, t1, t2, eris)

    if mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file)
        adiis.space = mycc.diis_space
    else:
        adiis = lambda t1,t2,*args: (t1, t2)
    cput0 = log.timer('CCSD lambda initialization', *cput0)

    conv = False
    for istep in range(max_cycle):
        l1new, l2new = update_amps(mycc, t1, t2, l1, l2, eris, saved)
        normt = numpy.linalg.norm(l1new-l1) + numpy.linalg.norm(l2new-l2)
        l1, l2 = l1new, l2new
        l1new = l2new = None
        if mycc.diis:
            l1, l2 = mycc.diis(l1, l2, istep, normt, 0, adiis)
        log.info('cycle = %d  norm(lambda1,lambda2) = %.6g', istep+1, normt)
        cput0 = log.timer('CCSD iter', *cput0)
        if normt < tol:
            conv = True
            break
    return conv, l1, l2

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[nocc:,nocc:]

    class _Saved:
        pass
    saved = _Saved()
    saved.ftmp = lib.H5TmpFile()
    saved.woooo = saved.ftmp.create_dataset('woooo', (nocc,nocc,nocc,nocc), 'f8')
    saved.wooov = saved.ftmp.create_dataset('wooov', (nocc,nocc,nocc,nvir), 'f8')
    saved.wOVov = saved.ftmp.create_dataset('wOVov', (nocc,nvir,nocc,nvir), 'f8')
    saved.wOvOv = saved.ftmp.create_dataset('wOvOv', (nocc,nvir,nocc,nvir), 'f8')
    saved.wovvv = saved.ftmp.create_dataset('wovvv', (nocc,nvir,nvir,nvir), 'f8')

# As we don't have l2 in memory, hold tau temporarily in memory
    w1 = fvv - numpy.einsum('ja,jb->ba', fov, t1)
    w2 = foo + numpy.einsum('ib,jb->ij', fov, t1)
    w3 = numpy.einsum('kc,jkbc->bj', fov, t2) * 2 + fov.T
    w3 -= numpy.einsum('kc,kjbc->bj', fov, t2)
    w3 += reduce(numpy.dot, (t1.T, fov, t1.T))
    w4 = fov.copy()

    _tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fswap = h5py.File(_tmpfile.name)

    time1 = time.clock(), time.time()
    unit = max(nocc*nvir**2*4 + nvir**3*2,
               nvir**3*3 + nocc*nvir**2,
               nocc*nvir**2*6 + nocc**2*nvir + nocc**3 + nocc**2*nvir)
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = max(ccsd.BLKMIN, int(max_memory*.95e6/8/unit))
    log.debug1('ccsd lambda make_intermediates: block size = %d, nocc = %d in %d blocks',
               blksize, nocc, int((nocc+blksize-1)//blksize))
    for istep, (p0, p1) in enumerate(prange(0, nocc, blksize)):
        eris_ovvv = _cp(eris.ovvv[p0:p1])
        eris_ovvv = lib.unpack_tril(eris_ovvv.reshape((p1-p0)*nvir,-1))
        eris_ovvv = eris_ovvv.reshape(p1-p0,nvir,nvir,nvir)
        w1 += numpy.einsum('jcba,jc->ba', eris_ovvv, t1[p0:p1]*2)
        w1 -= numpy.einsum('jabc,jc->ba', eris_ovvv, t1[p0:p1])
        #:w3 += numpy.einsum('kdcb,kjdc->bj', eris_ovvv, theta)
        for i in range(p1-p0):
            theta = t2[p0+i] * 2
            theta -= t2[p0+i].transpose(0,2,1)
            w3 += lib.dot(eris_ovvv[i].reshape(-1,nvir).T,
                          _cp(theta.reshape(nocc,-1)).T)
        theta = None
        #:wOVov = numpy.einsum('jbcd,kd->jbkc', eris_ovvv, t1)
        #:wOvOv = numpy.einsum('jdcb,kd->jbkc', eris_ovvv, -t1)
        wOVov = lib.dot(eris_ovvv.reshape(-1,nvir),
                        t1.T).reshape(-1,nvir,nvir,nocc).transpose(0,1,3,2).copy()
        g2ovvv = _cp(eris_ovvv.transpose(0,2,3,1))
        wOvOv = lib.dot(g2ovvv.reshape(-1,nvir),
                        -t1.T).reshape(-1,nvir,nvir,nocc).transpose(0,1,3,2).copy()
        for i in range(p1-p0):
            g2ovvv[i] *= 2
            g2ovvv[i] -= eris_ovvv[i].transpose(1,0,2)
        wooov = numpy.empty((p1-p0,nocc,nocc,nvir))
        woooo = numpy.empty((p1-p0,nocc,nocc,nocc))
        eris_ovov = _cp(_cp(eris.ovov[p0:p1]).transpose(0,2,1,3))
        for j0, j1 in prange(0, nocc, blksize):
            tau = _ccsd.make_tau(t2[j0:j1], t1[j0:j1], t1)
            #:wooov[:,j0:j1] = numpy.einsum('icbd,jkbd->ijkc', g2ovvv, tau)
            #:woooo[:,:,j0:j1] = numpy.einsum('icjd,klcd->ijkl', eris_ovov, tau)
            tmp = lib.dot(g2ovvv.reshape(-1,nvir**2), tau.reshape(-1,nvir**2).T)
            wooov[:,j0:j1] = tmp.reshape(-1,nvir,j1-j0,nocc).transpose(0,2,3,1)
            woooo[:,:,j0:j1] = lib.dot(eris_ovov.reshape(-1,nvir**2),
                                       tau.reshape(-1,nvir**2).T).reshape(-1,nocc,j1-j0,nocc)
        eris_ovov = eris_ovvv = g2ovvv = tau = tmp = None
#==== mem usage nocc*nvir**2*2 + nocc**2*nvir + nocc**3 + nvir**3*2 + nocc*nvir**2*2

        eris_ooov = _cp(eris.ooov[p0:p1])
        w2[p0:p1] += numpy.einsum('ijkb,kb->ij', eris_ooov, t1) * 2
        w2 -= numpy.einsum('kjib,kb->ij', eris_ooov, t1[p0:p1])
        #:w3 -= numpy.einsum('kjlc,klbc->bj', eris_ooov, theta)
        for i in range(p1-p0):
            theta = t2[p0+i].transpose(0,2,1) * 2
            theta -= t2[p0+i]
            w3 -= lib.dot(theta.reshape(-1,nvir).T, eris_ooov[i].reshape(nocc,-1).T)
        theta = None
        #:woooo += numpy.einsum('ikjc,lc->ijkl', eris_ooov, t1)
        #:wOvOv += numpy.einsum('jklb,lc->jbkc', eris_ooov, t1)
        woooo += lib.dot(eris_ooov.reshape(-1,nvir),
                         t1.T).reshape((-1,nocc,nocc,nocc)).transpose(0,2,1,3)
        for i in range(p1-p0):
            lib.dot(_cp(eris_ooov[i].transpose(2,0,1)).reshape(-1,nocc),
                    t1, 1, wOvOv[i].reshape(-1,nvir), 1)
            wooov[i] += eris_ooov[i].transpose(1,0,2)*2
            wooov[i] -= eris_ooov[i]

        eris_ovoo = _cp(eris.ovoo[p0:p1])
        #:woooo += numpy.einsum('icjl,kc->ijkl', eris_ovoo, t1)
        #:wOVov += numpy.einsum('jbkl,lc->jbkc', eris_ovoo, -t1)
        for i in range(p1-p0):
            woooo[i] += lib.dot(t1, eris_ovoo[i].reshape(nvir,-1)).reshape((nocc,)*3).transpose(1,0,2)
        lib.dot(eris_ovoo.reshape(-1,nocc), t1, -1, wOVov.reshape(-1,nvir), 1)
        #:wooov -= numpy.einsum('iblj,klbc->ijkc', eris_ovoo*1.5, t2)
        tmp_ovoo = _cp(-eris_ovoo.transpose(0,2,3,1)).reshape(-1,nov)
        for j in range(nocc):
            wooov[:,:,j] += lib.dot(tmp_ovoo, t2[j].reshape(-1,nvir),
                                    1.5).reshape(-1,nocc,nvir)
        #:g2ooov = eris_ooov * 2 - eris_ovoo.transpose(0,3,2,1)
        g2ooov, tmp_ovoo = tmp_ovoo.reshape(p1-p0,nocc,nocc,nvir), None
        g2ooov += eris_ooov * 2
        thetabuf = numpy.empty((blksize,nvir,nocc,nvir))
        vikjc = numpy.empty((p1-p0,nocc,blksize,nvir))
        for j0, j1 in prange(0, nocc, blksize):
            theta = thetabuf[:j1-j0]
            for i in range(j1-j0):
                theta[i] = t2[j0+i].transpose(1,0,2)*2 - t2[j0+i].transpose(2,0,1)
            #:vikjc = numpy.einsum('iklb,jlcb->ikjc', g2ooov, theta)
            if j1 == j0 + blksize:
                lib.dot(g2ooov.reshape(-1,nov), _cp(theta.reshape(-1,nov)).T,
                        1, vikjc.reshape((p1-p0)*nocc,-1), 0)
            else:
                vikjc = lib.dot(g2ooov.reshape(-1,nov),
                                _cp(theta.reshape(-1,nov)).T)
                vikjc = vikjc.reshape(p1-p0,nocc,j1-j0,nvir)
            wooov[:,j0:j1,:] += vikjc.transpose(0,2,1,3)
            wooov[:,:,j0:j1] -= vikjc*.5
        eris_ooov = eris_ovoo = g2ooov = vikjc = theta = thetabuf = None
#==== mem usage nocc*nvir**2*3 + nocc**2*nvir + nocc**3 + nocc*nvir**2 + nocc**2*nvir*3

        eris_ovov = _cp(eris.ovov[p0:p1])
        g2ovov = eris_ovov*2
        g2ovov -= eris_ovov.transpose(0,3,2,1)
        tmpw4 = numpy.einsum('kcld,ld->kc', g2ovov, t1)
        #:w1 -= numpy.einsum('kcja,kjcb->ba', g2ovov, t2[p0:p1])
        w1 -= lib.dot(t2[p0:p1].reshape(-1,nvir).T,
                      _cp(g2ovov.transpose(0,2,1,3).reshape(-1,nvir)))
        w1 -= numpy.einsum('ja,jb->ba', tmpw4, t1[p0:p1])
        #:w2[p0:p1] += numpy.einsum('ibkc,jkbc->ij', g2ovov, t2)
        w2[p0:p1] += lib.dot(_cp(g2ovov.transpose(0,2,1,3)).reshape(p1-p0,-1),
                             t2.reshape(nocc,-1).T)
        w2[p0:p1] += numpy.einsum('ib,jb->ij', tmpw4, t1)
        w3 += reduce(numpy.dot, (t1[p0:p1].T, tmpw4, t1.T))
        w4[p0:p1] += tmpw4
        vOVov = numpy.empty((nocc,nvir,p1-p0,nvir))
        #:vOVov += numpy.einsum('jbld,klcd->kcjb', g2ovov, t2)
        #:vOVov -= numpy.einsum('jbld,kldc->kcjb', eris_ovov, t2)
        for j in range(nocc):
            lib.dot(_cp(t2[j].transpose(1,0,2)).reshape(-1,nov),
                    g2ovov.reshape(-1,nov).T, 1, vOVov[j].reshape(nvir,-1))
            lib.dot(t2[j].reshape(nov,-1).T, eris_ovov.reshape(-1,nov).T,
                    -1, vOVov[j].reshape(nvir,-1), 1)
        vOVov = lib.transpose(vOVov.reshape(nov,-1)).reshape(p1-p0,nvir,nocc,nvir)
        vOVov += eris_ovov
        g2ovov = tmp = tmpw4 = None
#==== mem usage nocc*nvir**2*4 + nocc**2*nvir + nocc**3 + nocc*nvir**2

        #:tmp = numpy.einsum('jbld,kd->jlbk', eris_ovov, t1)
        #:wOVov -= numpy.einsum('jlbk,lc->jbkc', tmp, t1)
        #:tmp = numpy.einsum('jdlb,kd->jlbk', eris_ovov, t1)
        #:wOvOv += numpy.einsum('jlbk,lc->jbkc', tmp, t1)
        tmp = numpy.empty((nocc,nvir,nocc))
        for j in range(p1-p0):
            lib.dot(_cp(eris_ovov[j].transpose(1,0,2)).reshape(-1,nvir),
                    t1.T, 1, tmp.reshape(-1,nocc))
            lib.dot(tmp.reshape(nocc,-1).T, t1, -1, wOVov[j].reshape(-1,nvir), 1)
            lib.dot(eris_ovov[j].reshape(nvir,-1).T, t1.T, 1,
                    tmp.reshape(-1,nocc))
            lib.dot(tmp.reshape(nocc,-1).T, t1, 1, wOvOv[j].reshape(-1,nvir), 1)
        tmp = None

        #:vOvOv = numpy.einsum('jdlb,kldc->kcjb', eris_ovov, t2)
        ovovtmp = _cp(eris_ovov.transpose(0,3,2,1)).reshape(-1,nov)
        vOvOv = numpy.empty((nocc,nvir,p1-p0,nvir))
        for j in range(nocc):
            lib.dot(t2[j].reshape(-1,nvir).T, ovovtmp.T, 1,
                    vOvOv[j].reshape(nvir,-1))
        ovovtmp = eris_ovov = None
        vOvOv = lib.transpose(vOvOv.reshape(nov,-1)).reshape(p1-p0,nvir,nocc,nvir)
        vOvOv -= _cp(eris.oovv[p0:p1]).transpose(0,3,1,2)
        wOVov += vOVov
        wOvOv += vOvOv
        saved.wOVov[p0:p1] = wOVov
        saved.wOvOv[p0:p1] = wOvOv
        wOVov = wOvOv = None
#==== mem usage nocc*nvir**2*6 + nocc**2*nvir + nocc**3 + nocc**2*nvir

        ov1 = vOvOv*2 + vOVov
        #:wooov -= numpy.einsum('ibkc,jb->ijkc', ov1, t1)
        for i in range(p1-p0):
            lib.dot(t1, ov1[i].reshape(nvir,-1), -1, wooov[i].reshape(nocc,-1), 1)
        ov1 = lib.transpose(ov1.reshape(-1,nov))
        fswap['2vOvOv/%d'%istep] = ov1.reshape(nocc,nvir,-1,nvir)
        ov1 = None
        ov2 = vOVov*2 + vOvOv
        w3 += numpy.einsum('kcjb,kc->bj', ov2, t1[p0:p1])
        #:wooov += numpy.einsum('ibjc,kb->ijkc', ov2, t1)
        for i in range(p1-p0):
            wooov[i] += lib.dot(t1, ov2[i].reshape(nvir,-1)).reshape(nocc,nocc,nvir).transpose(1,0,2)
        ov2 = lib.transpose(ov2.reshape(-1,nov))
        fswap['2vovOV/%d'%istep] = ov2.reshape(nocc,nvir,-1,nvir)
        vOVov = vOvOv = None
        ov2 = None
#==== mem usage nocc*nvir**2*5 + nocc**2*nvir + nocc**3

        woooo += _cp(eris.oooo[p0:p1]).transpose(0,2,1,3)
        saved.woooo[p0:p1] = woooo
        saved.wooov[p0:p1] = wooov
        woooo = wooov = None
        time1 = log.timer_debug1('pass1 [%d:%d]'%(p0, p1), *time1)

    w3 += numpy.einsum('bc,jc->bj', w1, t1)
    w3 -= numpy.einsum('kj,kb->bj', w2, t1)

    for p0, p1 in prange(0, nocc, blksize):
        eris_ooov = _cp(eris.ooov[p0:p1])
        g2ooov = eris_ooov * 2
        g2ooov -= eris_ooov.transpose(0,2,1,3)
        #:tmp = numpy.einsum('kjla,jb->kabl', g2ooov, t1)
        #:wovvv = numpy.einsum('kabl,lc->kabc', tmp, t1)
        #:wovvv += numpy.einsum('kjla,jlbc->kabc', g2ooov, t2)
        tmp = lib.dot(_cp(g2ooov.transpose(1,0,2,3).reshape(nocc,-1)).T,
                      t1).reshape(-1,nocc,nvir,nvir).transpose(0,2,3,1)
        wovvv = lib.dot(_cp(tmp.reshape(-1,nocc)), t1).reshape(-1,nvir,nvir,nvir)
        wovvv += lib.dot(_cp(g2ooov.transpose(0,3,1,2).reshape(-1,nocc**2)),
                         t2.reshape(nocc**2,-1)).reshape(-1,nvir,nvir,nvir)
        tmp = g2ooov = None
        ov1 = numpy.empty((p1-p0,nvir,nocc,nvir))
        ov2 = numpy.empty((p1-p0,nvir,nocc,nvir))
        for istep, (j0, j1) in enumerate(prange(0, nocc, blksize)):
            ov1[:,:,j0:j1] = fswap['2vOvOv/%d'%istep][p0:p1]
            ov2[:,:,j0:j1] = fswap['2vovOV/%d'%istep][p0:p1]
        #:wovvv += numpy.einsum('kcja,jb->kabc', ov1, t1)
        #:wovvv -= numpy.einsum('kbja,jc->kabc', ov2, t1)
        wovvv += lib.dot(_cp(ov1.transpose(0,1,3,2).reshape(-1,nocc)),
                         t1).reshape(-1,nvir,nvir,nvir).transpose(0,2,3,1)
        wovvv -= lib.dot(_cp(ov2.transpose(0,1,3,2).reshape(-1,nocc)),
                         t1).reshape(-1,nvir,nvir,nvir).transpose(0,2,1,3)
#==== mem usage nvir**3 + nocc*nvir**2*2
        eris_ooov = ov1 = ov2 = None

        for j0, j1 in prange(0, nocc, blksize):
            eris_ovvv = _cp(eris.ovvv[j0:j1])
            eris_ovvv = lib.unpack_tril(eris_ovvv.reshape((j1-j0)*nvir,-1))
            eris_ovvv = eris_ovvv.reshape(j1-j0,nvir,nvir,nvir)
            #:wovvv += numpy.einsum('jabd,kjdc->kabc', eris_ovvv, t2[p0:p1,j0:j1]) * -1.5
            tmp_ovvv = numpy.empty((j1-j0,nvir,nvir,nvir))
            for i in range(j1-j0):
                tmp_ovvv[i] = eris_ovvv[i].transpose(1,0,2)*2
            tmp = lib.dot(_cp(t2[p0:p1,j0:j1].transpose(0,3,1,2).reshape((p1-p0)*nvir,-1)),
                          tmp_ovvv.reshape(-1,nvir**2), -1.5/2).reshape(-1,nvir,nvir,nvir)
            wovvv += tmp.transpose(0,2,3,1)
            if p0 == j0:
                for i in range(p1-p0):
                    tmp_ovvv[i] -= eris_ovvv[i].transpose(1,2,0)
                    wovvv[i] += tmp_ovvv[i]
            tmp = tmp_ovvv = None
            g2ovvv = numpy.empty((j1-j0,nvir,nvir,nvir))
            for i in range(j1-j0):
                g2ovvv[i] = eris_ovvv[i] * 2
                g2ovvv[i] -= eris_ovvv[i].transpose(1,2,0)
#==== mem usage nvir**3*3
            eris_ovvv = None
            theta = _cp(t2[p0:p1,j0:j1].transpose(0,2,1,3)*2)
            for i in range(p1-p0):
                theta[i] -= t2[p0+i,j0:j1].transpose(2,0,1)
            #:vkbca = numpy.einsum('jdca,kbjd->kbca', g2ovvv, theta)
            vkbca = lib.dot(theta.reshape((p1-p0)*nvir,-1),
                            g2ovvv.reshape(-1,nvir*nvir)).reshape(-1,nvir,nvir,nvir)
            wovvv += vkbca.transpose(0,3,1,2)
            wovvv -= vkbca.transpose(0,3,2,1) * .5
#==== mem usage nvir**3*3 + nocc*nvir**2
            g2ovvv = theta = vkabc = None
        saved.wovvv[p0:p1] = wovvv
        time1 = log.timer_debug1('pass2 [%d:%d]'%(p0, p1), *time1)

    fswap.close()

    saved.w1 = w1
    saved.w2 = w2
    saved.w3 = w3
    saved.w4 = w4
    saved.ftmp.flush()
    return saved

# update L1, L2
def update_amps(mycc, t1, t2, l1, l2, eris=None, saved=None):
    if saved is None:
        saved = make_intermediates(mycc, t1, t2, eris)
    time1 = time0 = time.clock(), time.time()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    fov = eris.fock[:nocc,nocc:]

    #:mba = numpy.einsum('klca,klcb->ba', l2, t2*2-t2.transpose(0,1,3,2))
    #:mij = numpy.einsum('ikcd,jkcd->ij', l2, t2*2-t2.transpose(0,1,3,2))
    #:theta = t2*2 - t2.transpose(0,1,3,2)
    theta = _ccsd.make_0132(t2, t2, 2, -1)
    mba = lib.dot(theta.reshape(-1,nvir).T, l2.reshape(-1,nvir))
    mij = lib.dot(l2.reshape(nocc,-1), theta.reshape(nocc,-1).T)
    theta = None
    mba1 = numpy.einsum('jc,jb->bc', l1, t1) + mba
    mij1 = numpy.einsum('kb,jb->kj', l1, t1) + mij
    mia1 =(t1 + numpy.einsum('kc,jkbc->jb', l1, t2) * 2
         - numpy.einsum('kc,jkcb->jb', l1, t2)
         - reduce(numpy.dot, (t1, l1.T, t1))
         - numpy.einsum('bd,jd->jb', mba, t1)
         - numpy.einsum('lj,lb->jb', mij, t1))

    tmp = mycc.add_wvvVV(numpy.zeros_like(l1), l2, eris)
    l2new = numpy.empty((nocc,nocc,nvir,nvir))
    ij = 0
    for i in range(nocc):
        for j in range(i):
            tmp1 = tmp[ij] * .5  # *.5 because of l2+l2.transpose(1,0,3,2) later
            l2new[i,j] = tmp1
            l2new[j,i] = tmp1.T
            ij += 1
        l2new[i,i] = tmp[ij] * .5
        ij += 1
    l1new =(numpy.einsum('ijab,jb->ia', l2new, t1) * 4
          - numpy.einsum('jiab,jb->ia', l2new, t1) * 2)
    tmp = tmp1 = None

    l1new += fov
    l1new += numpy.einsum('ib,ba->ia', l1, saved.w1)
    l1new -= numpy.einsum('ja,ij->ia', l1, saved.w2)
    l1new -= numpy.einsum('ik,ka->ia', mij, saved.w4)
    l1new -= numpy.einsum('ca,ic->ia', mba, saved.w4)
    l1new += numpy.einsum('ijab,bj->ia', l2, saved.w3) * 2
    l1new -= numpy.einsum('ijba,bj->ia', l2, saved.w3)

    l2new += numpy.einsum('ia,jb->ijab', l1, saved.w4)
    #:l2new += numpy.einsum('jibc,ca->jiba', l2, saved.w1)
    #:l2new -= numpy.einsum('kiba,jk->jiba', l2, saved.w2)
    lib.dot(l2.reshape(-1,nvir), saved.w1, 1, l2new.reshape(-1,nvir), 1)
    lib.dot(saved.w2, l2.reshape(nocc,-1),-1, l2new.reshape(nocc,-1), 1)

    eris_ooov = _cp(eris.ooov)
    l1new -= numpy.einsum('jkia,kj->ia', eris_ooov, mij1) * 2
    l1new += numpy.einsum('ikja,kj->ia', eris_ooov, mij1)
    #:l2new -= numpy.einsum('ka,kijb->jiba', l1, eris_ooov)
    lib.dot(_cp(eris_ooov.transpose(0,2,1,3).reshape(nocc,-1)).T,
            l1, -1, l2new.reshape(-1,nvir), 1)
    eris_ooov = None

    tau = _ccsd.make_tau(t2, t1, t1)
    #:l2tau = numpy.einsum('ijcd,klcd->ijkl', l2, tau)
    l2tau = lib.dot(l2.reshape(nocc**2,-1),
                    tau.reshape(nocc**2,-1).T).reshape((nocc,)*4)
    tau = None
    #:l2t1 = numpy.einsum('jidc,kc->ijkd', l2, t1)
    l2t1 = lib.dot(l2.reshape(-1,nvir), t1.T).reshape(nocc,nocc,nvir,nocc)
    l2t1 = _cp(l2t1.transpose(1,0,3,2))

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = max(nvir**3*2+nocc*nvir**2, nocc*nvir**2*5)
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*.95e6/8/unit)))
    log.debug1('block size = %d, nocc = %d is divided into %d blocks',
               blksize, nocc, int((nocc+blksize-1)/blksize))
    for p0, p1 in prange(0, nocc, blksize):
        eris_ovvv = _cp(eris.ovvv[p0:p1])
        eris_ovvv = lib.unpack_tril(eris_ovvv.reshape((p1-p0)*nvir,-1))
        eris_ovvv = eris_ovvv.reshape(p1-p0,nvir,nvir,nvir)

        l1new[p0:p1] += numpy.einsum('iabc,bc->ia', eris_ovvv, mba1) * 2
        l1new[p0:p1] -= numpy.einsum('ibca,bc->ia', eris_ovvv, mba1)
        #:l2new[p0:p1] += numpy.einsum('ic,jbac->jiba', l1, eris_ovvv)
        tmp = lib.dot(l1, eris_ovvv.reshape(-1,nvir).T)
        l2new[p0:p1] += tmp.reshape(nocc,-1,nvir,nvir).transpose(1,0,2,3)
        tmp = None
        m4buf = numpy.empty((blksize,nocc,nvir,nvir))
        eris_ovvv = _cp(eris_ovvv.transpose(0,2,1,3).reshape(-1,nvir**2))
        for j0, j1 in prange(0, nocc, blksize):
            #:m4 = numpy.einsum('ijkd,kadb->ijab', l2t1[j0:j1,:,p0:p1], eris_ovvv)
            m4 = m4buf[:j1-j0]
            lib.dot(_cp(l2t1[j0:j1,:,p0:p1].reshape((j1-j0)*nocc,-1)),
                    eris_ovvv, 1, m4.reshape(-1,nvir**2))
            l2new[j0:j1] -= m4
            l1new[j0:j1] -= numpy.einsum('ijab,jb->ia', m4, t1) * 2
            l1new -= numpy.einsum('ijab,ia->jb', m4, t1[j0:j1]) * 2
            l1new += numpy.einsum('jiab,jb->ia', m4, t1[j0:j1])
            l1new[j0:j1] += numpy.einsum('jiab,ia->jb', m4, t1)
        eris_ovvv = m4buf = m4 = None
#==== mem usage nvir**3*2 + nocc*nvir**2

        eris_ovov = _cp(eris.ovov[p0:p1])
        l1new[p0:p1] += numpy.einsum('jb,iajb->ia', l1, eris_ovov) * 2
        for i in range(p1-p0):
            l2new[p0+i] += eris_ovov[i].transpose(1,0,2) * .5
        #:l2new[p0:p1] -= numpy.einsum('icjb,ca->ijab', eris_ovov, mba1)
        #:l2new[p0:p1] -= numpy.einsum('jbka,ik->jiba', eris_ovov, mij1)
        tmp = numpy.empty((nocc,nvir,nvir))
        for j in range(p0,p1):
            lib.dot(eris_ovov[j-p0].reshape(nvir,-1).T, mba1, 1,
                    tmp.reshape(-1,nvir))
            l2new[j] -= tmp.transpose(0,2,1)
            lib.dot(mij1, _cp(eris_ovov[j-p0].transpose(1,0,2).reshape(nocc,-1)),
                    -1, l2new[j].reshape(nocc,-1), 1)
        tmp = None
        l1new[p0:p1] += numpy.einsum('iajb,jb->ia', eris_ovov, mia1) * 2
        l1new[p0:p1] -= numpy.einsum('ibja,jb->ia', eris_ovov, mia1)
        m4buf = numpy.empty((blksize,nocc,nvir,nvir))
        for j0, j1 in prange(0, nocc, blksize):
            #:m4 = numpy.einsum('kalb,ijkl->ijab', eris_ovov, l2tau[j0:j1,:,p0:p1])
            m4 = m4buf[:j1-j0]
            lib.dot(l2tau[j0:j1,:,p0:p1].reshape((j1-j0)*nocc,-1).copy(),
                    _cp(eris_ovov.transpose(0,2,1,3).reshape(-1,nvir**2)),
                    .5, m4.reshape(-1,nvir**2))
            l2new[j0:j1] += m4
            l1new[j0:j1] += numpy.einsum('ijab,jb->ia', m4, t1) * 4
            l1new[j0:j1] -= numpy.einsum('ijba,jb->ia', m4, t1) * 2
        eris_ovov = m4buf = m4 = None
#==== mem usage nocc*nvir**2 * 3

        eris_oovv = _cp(eris.oovv[p0:p1])
        l1new[p0:p1] -= numpy.einsum('jb,ijba->ia', l1, eris_oovv)
        eris_oovv = None

        saved_wooov = _cp(saved.wooov[p0:p1])
        #:l1new[p0:p1] -= numpy.einsum('jkca,ijkc->ia', l2, saved_wooov)
        l1new[p0:p1] -= lib.dot(saved_wooov.reshape(p1-p0,-1),
                                l2.reshape(-1,nvir))
        saved_wovvv = _cp(saved.wovvv[p0:p1])
        #:l1new += numpy.einsum('kibc,kabc->ia', l2[p0:p1], saved_wovvv)
        for j in range(p1-p0):
            lib.dot(l2[p0+j].reshape(nocc,-1),
                    saved_wovvv[j].reshape(nvir,-1).T, 1, l1new, 1)
        saved_wooov = saved_wovvv = None
#==== mem usage nvir**3 + nocc**2*nvir

        saved_wOvOv = _cp(saved.wOvOv[p0:p1])
        tmp_ovov = _cp(saved.wOVov[p0:p1]) * 2
        tmp_ovov += saved_wOvOv
        tmp_ovov = lib.transpose(tmp_ovov.reshape(-1,nov)).reshape(nocc,nvir,-1,nvir)
        tmp1 = numpy.empty((p1-p0,nvir,nocc,nvir))
        tmp = numpy.empty((blksize,nvir,nocc,nvir))
        for j0, j1 in prange(0, nocc, blksize):
            #:tmp = l2[j0:j1].transpose(0,2,1,3) - l2[j0:j1].transpose(0,3,1,2)*.5
            #:l2new[p0:p1] += numpy.einsum('kcia,kcjb->jiba', tmp, tmp_ovov[j0:j1])
            for i in range(j1-j0):
                tmp[i] = -.5 * l2[j0+i].transpose(2,0,1)
                tmp[i] += l2[j0+i].transpose(1,0,2)
            lib.dot(tmp_ovov[j0:j1].reshape((j1-j0)*nvir,-1).T,
                    tmp[:j1-j0].reshape((j1-j0)*nvir,-1), 1, tmp1.reshape(-1,nov))
            l2new[p0:p1] += tmp1.transpose(0,2,1,3)
        tmp = tmp1 = tmp_ovov = None
#==== mem usage nocc*nvir**2 * 5

        #:tmp = numpy.einsum('jkca,ibkc->ijab', l2, saved_wOvOv)
        tmp = numpy.empty((p1-p0,nvir,nvir))
        for j in range(nocc):
            lib.dot(saved_wOvOv.reshape(-1,nov), l2[j].reshape(nov,-1), 1,
                    tmp.reshape(-1,nvir))
            l2new[p0:p1,j] += tmp.transpose(0,2,1)
            l2new[p0:p1,j] += tmp * .5
        saved_wOvOv = tmp = None

        saved_woooo = _cp(saved.woooo[p0:p1])
        #:m3 = numpy.einsum('klab,ijkl->ijab', l2, saved_woooo)
        m3 = lib.dot(saved_woooo.reshape(-1,nocc**2),
                     l2.reshape(nocc**2,-1), .5).reshape(-1,nocc,nvir,nvir)
        l2new[p0:p1] += m3
        l1new[p0:p1] += numpy.einsum('ijab,jb->ia', m3, t1) * 4
        l1new[p0:p1] -= numpy.einsum('ijba,jb->ia', m3, t1) * 2
        saved_woooo = m3 = None
        time1 = log.timer_debug1('lambda pass [%d:%d]'%(p0, p1), *time1)

    mo_e = eris.fock.diagonal()
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    l1new /= eia
    l1new += l1

#    l2new = l2new + l2new.transpose(1,0,3,2)
#    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)
#    l2new += l2
    ij = 0
    for i in range(nocc):
        for j in range(i):
            dab = lib.direct_sum('a+b->ab', eia[i], eia[j])
            tmp = (l2new[i,j]+l2new[j,i].T) / dab + l2[i,j]
            l2new[i,j] = tmp
            l2new[j,i] = tmp.T
            ij += 1
        dab = lib.direct_sum('a+b->ab', eia[i], eia[i])
        l2new[i,i] = (l2new[i,i]+l2new[i,i].T)/dab + l2[i,i]
        ij += 1

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new

def prange(start, end, step):
    for i in range(start, end, step):
        yield i, min(i+step, end)

def _cp(a):
    return numpy.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import ccsd

    mol = gto.M()
    mf = scf.RHF(mol)

    mcc = ccsd.CCSD(mf)

    numpy.random.seed(12)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)

    eris = lambda:None
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    idx = numpy.tril_indices(nvir)
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:][:,:,idx[0],idx[1]].copy()
    eris.vvvv = ao2mo.restore(4,eri0[nocc:,nocc:,nocc:,nocc:],nvir)
    eris.fock = fock0

    saved = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_amps(mcc, t1, t2, l1, l2, eris, saved)
    print(abs(l1new).sum()-38172.7896467303)
    print(numpy.dot(l1new.flatten(), numpy.arange(35)) - 739312.005491083)
    print(numpy.dot(l1new.flatten(), numpy.sin(numpy.arange(35)))-7019.50937051188)
    print(numpy.dot(numpy.sin(l1new.flatten()), numpy.arange(35))-69.6652346635955)

    print(abs(l2new).sum()-72035.4931071527)
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())
    print(numpy.dot(l2new.flatten(), numpy.arange(35**2)) - 48427109.5409886)
    print(numpy.dot(l2new.flatten(), numpy.sin(numpy.arange(35**2)))-137.758016736487)
    print(numpy.dot(numpy.sin(l2new.flatten()), numpy.arange(35**2))-507.656936701192)

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-16
    rhf.scf()

    mcc = ccsd.CCSD(rhf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()

    nmo = rhf.mo_energy.size
    fock0 = numpy.diag(rhf.mo_energy)
    nocc = mol.nelectron // 2
    nvir = nmo - nocc

    eris = mcc.ao2mo()
    conv, l1, l2 = kernel(mcc, eris, t1, t2, tol=1e-8)
    print(numpy.linalg.norm(l1)-0.0132626841292)
    print(numpy.linalg.norm(l2)-0.212575609057)

    from pyscf.cc import ccsd_rdm
    dm1 = ccsd_rdm.make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = ccsd_rdm.make_rdm2(mcc, t1, t2, l1, l2)
    h1 = reduce(numpy.dot, (rhf.mo_coeff.T, rhf.get_hcore(), rhf.mo_coeff))
    eri = ao2mo.full(rhf._eri, rhf.mo_coeff)
    eri = ao2mo.restore(1, eri, nmo).reshape((nmo,)*4)
    e1 = numpy.einsum('pq,pq', h1, dm1)
    e2 = numpy.einsum('pqrs,pqrs', eri, dm2) * .5
    print(e1+e2+mol.energy_nuc() - rhf.e_tot - ecc)


import torch
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader

import argparse
import os

from tqdm import tqdm
import numpy as np
from scipy import linalg
from PIL import Image
from torchvision import datasets, transforms, utils
from torchvision.datasets.folder import find_classes
from torchvision.models import inception_v3, Inception3

parser = argparse.ArgumentParser(description='FID score calculator')
parser.add_argument('--img', required=True, help='path to image directory')
parser.add_argument('--batch', default=64, type=int, help='batch size')
parser.add_argument('--sample', default=5000, type=int,
                    help='number of samples generated for evaluation')
parser.add_argument('--code', default=128, type=int,
                    help='code size for generator')
parser.add_argument('--model', default='dcgan', choices=['dcgan', 'resnet'],
                    help='choice model class')
parser.add_argument('checkpoint', metavar='CHECKPOINT',
                    help='checkpoint of generator model')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def forward(self, x):
    x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=True)
    F.interpolate()

    x = self.Conv2d_1a_3x3(x)  # 299 x 299 x 3
    x = self.Conv2d_2a_3x3(x)  # 149 x 149 x 32
    x = self.Conv2d_2b_3x3(x)  # 147 x 147 x 32
    x = F.max_pool2d(x, kernel_size=3, stride=2)  # 147 x 147 x 64

    x = self.Conv2d_3b_1x1(x)  # 73 x 73 x 64
    x = self.Conv2d_4a_3x3(x)  # 73 x 73 x 80
    x = F.max_pool2d(x, kernel_size=3, stride=2)  # 71 x 71 x 192

    x = self.Mixed_5b(x)  # 35 x 35 x 192
    x = self.Mixed_5c(x)  # 35 x 35 x 256
    x = self.Mixed_5d(x)  # 35 x 35 x 288

    x = self.Mixed_6a(x)  # 35 x 35 x 288
    x = self.Mixed_6b(x)  # 17 x 17 x 768
    x = self.Mixed_6c(x)  # 17 x 17 x 768
    x = self.Mixed_6d(x)  # 17 x 17 x 768
    x = self.Mixed_6e(x)  # 17 x 17 x 768

    x = self.Mixed_7a(x)  # 17 x 17 x 768
    x = self.Mixed_7b(x)  # 8 x 8 x 1280
    x = self.Mixed_7c(x)  # 8 x 8 x 2048

    x = F.avg_pool2d(x, kernel_size=8)  # 8 x 8 x 2048

    return x.squeeze()  # 1 x 1 x 2048

def load_patched_inception_v3():
    inception = inception_v3(pretrained=True)
    inception.eval()
    inception.forward = forward.__get__(inception, Inception3)

    return inception.to(device)

transform = transforms.Compose([
    transforms.Resize(128),
    transforms.CenterCrop(128),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

class ImageData(Dataset):
    def __init__(self, root, transform=transform):
        self.root = root
        self.transform = transform
        self.images = os.listdir(root)

    def __getitem__(self, index):
        img = os.path.join(self.root, self.images[index])
        img = Image.open(img).convert('RGB')

        if self.transform is not None:
            img = transform(img)

        return img

    def __len__(self):
        return len(self.images)

if __name__ == '__main__':
    args = parser.parse_args()
    print(args)

    eps = 1e-6

    inception = load_patched_inception_v3()
    _, class2id = find_classes(args.img)
    total_class = len(class2id)

    if args.model == 'dcgan':
        from model import Generator

    elif args.model == 'resnet':
        from model_resnet import Generator

    generator = Generator(args.code, total_class).to(device)
    generator.load_state_dict(torch.load(args.checkpoint))
    generator.eval()

    fids = []

    for class_name, id in class2id.items():
        dataloader = DataLoader(ImageData(os.path.join(args.img, class_name)),
                                batch_size=args.batch, num_workers=4)
        real_feature = []
        with torch.no_grad():
            for image in tqdm(dataloader,
                              desc='Extract features of real images'):
                image = image.to(device)
                features = inception(image).detach().cpu().numpy()
                real_feature.append(features)

        real_feature = np.concatenate(real_feature)
        real_mean = np.mean(real_feature, 0)
        real_cov = np.cov(real_feature, rowvar=False)
        del real_feature

        sample_feature = []
        with torch.no_grad():
            n_iter = args.sample // args.batch
            resid = args.sample - n_iter * args.batch
            n_samples = [args.batch] * n_iter
            if resid != 0:
                n_samples += [resid]
            for i in tqdm(n_samples,
                          desc='Extract features of fake images'):
                input_class = torch.full([i], id, dtype=torch.long) \
                    .to(device)
                code = torch.randn(i, args.code).to(device)
                sample = generator(code, input_class)
                features = inception(sample).detach().cpu().numpy()
                sample_feature.append(features)

        sample_feature = np.concatenate(sample_feature)
        sample_mean = np.mean(sample_feature, 0)
        sample_cov = np.cov(sample_feature, rowvar=False)
        del sample_feature

        # Came from https://github.com/bioinf-jku/TTUR

        cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)
        if not np.isfinite(cov_sqrt).all():
            print('Product of cov matrices is singular')
            offset = np.eye(sample_cov.shape[0]) * eps
            cov_sqrt = linalg.sqrtm((sample_cov + offset)
                                    @ (real_cov + offset))

        if np.iscomplexobj(cov_sqrt):
            if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
                m = np.max(np.abs(cov_sqrt.imag))
                raise ValueError(f'Imaginary component {m}')

            cov_sqrt = cov_sqrt.real

        mean_diff = sample_mean - real_mean
        mean_norm = mean_diff @ mean_diff

        trace = np.trace(sample_cov) + np.trace(real_cov) \
            - 2 * np.trace(cov_sqrt)

        fid = mean_norm + trace

        print(f'FID score of class {id} ({class_name}):', fid)
        fids.append(fid)

    print(f'Mean FID is: {sum(fids) / len(fids)}')


# -*- coding: utf-8 -*-
"""
A submodule with classes used for supernova model files stored on disk. It
assumes models are available in a format usable by the AstroPy unified table
reader; see https://docs.astropy.org/en/stable/index.html for details.

Based on the ASTERIA (https://github.com/IceCubeOpenSource/ASTERIA) models
developed by Navya Uberoi and Spencer Griswold.

Updated summer 2020 by Jim Kneller & Arkin Worlikar.
"""

from abc import abstractmethod, ABC
from enum import IntEnum

import astropy
from astropy.io import ascii
from astropy.table import Table, join
from astropy.units.quantity import Quantity

import matplotlib as mpl
import matplotlib.pyplot as plt

import numpy as np
from scipy.interpolate import interp1d
from scipy.special import loggamma

import os
import logging
import re

import tarfile
import h5py

from .neutrino import Flavor
from .flavor_transformation import *

def get_value(x):
    """If quantity x has is an astropy Quantity with units, return just the
    value.

    Parameters
    ----------
    x : Quantity, float, or ndarray
        Input quantity.

    Returns
    -------
    value : float or ndarray
    """
    if type(x) == Quantity:
        return x.value
    return x

def get_closest(arr, x):
    """Get index of closest element in an array to input value.

    Parameters
    ----------
    arr : list or ndarray
        Array of values.
    x : float or int or str
        Value to search.

    Returns
    -------
    idx : int
        Index of closest element in the array.
    """
    return np.abs(np.asarray(arr) - x).argmin()

class SupernovaModel(ABC):
    """Base class defining an interface to a supernova model."""

    def __init__(self):
        pass

    @abstractmethod
    def get_time(self):
        """Returns
        -------
            returns array of snapshot times from the simulation
        """
        pass

    @abstractmethod
    def get_initialspectra(self, t, E):
        """Get neutrino spectra at the source.

        Parameters
        ----------
        t : float
            Time to evaluate spectra.
        E : float or ndarray
            Energies to evaluate spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of neutrino spectra, keyed by neutrino flavor.
        """
        pass

    def get_oscillatedspectra(self, t, E, flavor_xform):
        """Get neutrino spectra after applying oscillation.

        Parameters
        ----------
        t : astropy.Quantity
            Time to evaluate initial and oscillated spectra.
        E : astropy.Quantity or ndarray
            Energies to evaluate the initial and oscillated spectra.
        flavor_xform : FlavorTransformation
            An instance from the flavor_transformation module.

        Returns
        -------
        oscillatedspectra : dict
            Dictionary of oscillated spectra, keyed by neutrino flavor.
        """
        initialspectra = self.get_initialspectra(t, E)
        oscillatedspectra = {}

        oscillatedspectra[Flavor.NU_E] = \
            flavor_xform.prob_ee(t, E) * initialspectra[Flavor.NU_E] + \
            flavor_xform.prob_ex(t, E) * initialspectra[Flavor.NU_X]

        oscillatedspectra[Flavor.NU_X] = \
            flavor_xform.prob_xe(t, E) * initialspectra[Flavor.NU_E] + \
            flavor_xform.prob_xx(t, E) * initialspectra[Flavor.NU_X]

        oscillatedspectra[Flavor.NU_E_BAR] = \
            flavor_xform.prob_eebar(t, E) * initialspectra[Flavor.NU_E_BAR] + \
            flavor_xform.prob_exbar(t, E) * initialspectra[Flavor.NU_X_BAR]

        oscillatedspectra[Flavor.NU_X_BAR] = \
            flavor_xform.prob_xebar(t, E) * initialspectra[Flavor.NU_E_BAR] + \
            flavor_xform.prob_xxbar(t, E) * initialspectra[Flavor.NU_X_BAR]

        return oscillatedspectra

class Nakazato_2013(SupernovaModel):
    """Set up a model based on simulations from Nakazato et al., ApJ S 205:2,
    2013 and ApJ 804:75, 2015. See also http://asphwww.ph.noda.tus.ac.jp/snn/.
    """

    def __init__(self, filename):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to FITS file with model data.
        """
        # Store model metadata.
        if 't_rev' in filename:
            self.progenitor_mass = float(filename.split('-')[-1].strip('s%.fits')) * u.Msun
            self.revival_time = float(filename.split('-')[-2].strip('t_rev%ms')) * u.ms
            self.metallicity = float(filename.split('-')[-3].strip('z%'))
            self.EOS = filename.split('-')[-4].upper()
        # No revival time because the explosion "failed" (BH formation).
        else:
            self.progenitor_mass = float(filename.split('-')[-1].strip('s%.fits')) * u.Msun
            self.metallicity = float(filename.split('-')[-2].strip('z%'))
            self.revival_time = 0 * u.ms
            self.EOS = filename.split('-')[-4].upper()

        # Read FITS table using the astropy reader.
        simtab = Table.read(filename)
        self.filename = os.path.basename(filename)

        # Get grid of model times.
        self.time = simtab['TIME'].to('s')

        # Set up dictionary of luminosity, mean energy and shape parameter
        # alpha, keyed by neutrino flavor (NU_E, NU_X, NU_E_BAR, NU_X_BAR).
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}

        for flavor in Flavor:
            # Note: file only contains NU_E, NU_E_BAR, and NU_X, so double up
            # the use of NU_X for NU_X_BAR.
            _flav = Flavor.NU_X if flavor == Flavor.NU_X_BAR else flavor

            self.luminosity[flavor] = simtab['L_{}'.format(_flav.name)].to('erg/s')
            self.meanE[flavor] = simtab['E_{}'.format(_flav.name)].to('MeV')
            self.pinch[flavor] = simtab['ALPHA_{}'.format(_flav.name)]

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.time

    def get_initialspectra(self, t, E):
        """Get neutrino spectra/luminosity at the source.

        Parameters
        ----------
        t : float or astropy.Quantity
            Time to evaluate initial spectra.
        E : float or ndarray of astropy.Quantity
            Energies to evaluate the initial spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor.
        """
        initialspectra = {}

        # Avoid division by zero in energy PDF below.
        E[E==0] = np.finfo(float).eps * E.unit

        # Estimate L(t), <E_nu(t)> and alpha(t). Express all energies in erg.
        E = E.to('erg').value

        # Make sure input time uses the same units as the model time grid, or
        # the interpolation will not work correctly.
        t = t.to(self.time.unit)

        for flavor in Flavor:
            # Use np.interp rather than scipy.interpolate.interp1d because it
            # can handle dimensional units (astropy.Quantity).
            L  = get_value(np.interp(t, self.time, self.luminosity[flavor].to('erg/s')))
            Ea = get_value(np.interp(t, self.time, self.meanE[flavor].to('erg')))
            a  = np.interp(t, self.time, self.pinch[flavor])

            # For numerical stability, evaluate log PDF and then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

    def __repr__(self):
        """Default representation of the model.
        """
        mod = 'Nakazato_2013 Model: {}\n'.format(self.filename)
        s = ['Progenitor mass : {}'.format(self.progenitor_mass),
             'Metallicity     : {}'.format(self.metallicity),
             'Revival time    : {}'.format(self.revival_time),
             'Eq. of state    : {}'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

    def _repr_markdown_(self):
        """Markdown representation of the model, for Jupyter notebooks.
        """
        mod = '**Nakazato_2013 Model**: {}\n\n'.format(self.filename)
        s = ['|Parameter|Value|',
             '|:---------|:-----:|',
             '|Progenitor mass | ${0.value:g}$ {0.unit:latex}|'.format(self.progenitor_mass),
             '|Metallicity | ${:g}$|'.format(self.metallicity),
             '|Revival time | ${0.value:g}$ {0.unit:latex}|'.format(self.revival_time),
             '|EOS | {}|'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

class Sukhbold_2015(SupernovaModel):
    """Set up a model based on simulations from Sukhbold et al., ApJ 821:38,2016. Models were shared privately by email.
    """

    def __init__(self, filename):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to FITS file with model data.
        """
        # Store model metadata.
        self.progenitor_mass = float(filename.split('-')[-1].strip('z%.fits')) * u.Msun
        self.EOS = filename.split('-')[-2]

        # Read FITS table using the astropy unified Table reader.
        simtab = Table.read(filename)
        self.filename = os.path.basename(filename)

        # Get grid of model times.
        self.time = simtab['TIME'].to('s')

        # Set up dictionary of luminosity, mean energy, and shape parameter,
        # keyed by neutrino flavor (NU_E, NU_X, NU_E_BAR, NU_X_BAR).
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}

        for flavor in Flavor:
            self.luminosity[flavor] = simtab['L_{}'.format(flavor.name)].to('erg/s')
            self.meanE[flavor] = simtab['E_{}'.format(flavor.name)].to('MeV')
            self.pinch[flavor] = simtab['ALPHA_{}'.format(flavor.name)]

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.time

    def get_initialspectra(self,t,E):
        """Get neutrino spectra/luminosity curves after oscillation.

        Parameters
        ----------
        t : astropy.Quantity
            Time to evaluate initial spectra.
        E : astropy.Quantity or ndarray
            Energies to evaluate the initial spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor.
        """
        initialspectra = {}

        # Avoid division by zero in energy PDF below.
        E[E==0] = np.finfo(float).eps * E.unit

        # Estimate L(t), <E_nu(t)>, and alpha(t). Express all energies in erg.
        E = E.to('erg').value

        # Make sure input time uses the same units as the global time grid or
        # the interpolation will not work properly.
        t = t.to(self.time.unit)

        for flavor in Flavor:
            L  = get_value(np.interp(t, self.time, self.luminosity[flavor].to('erg/s')))
            Ea = get_value(np.interp(t, self.time, self.meanE[flavor].to('erg')))
            a  = np.interp(t, self.time, self.pinch[flavor])

            # For numerical stability, evaluate log PDF then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

    def __repr__(self):
        """Default representation of the model.
        """
        mod = 'Sukhbold_2015 Model: {}\n'.format(self.filename)
        s = ['Progenitor mass : {}'.format(self.progenitor_mass),
             'Eq. of state    : {}'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

    def _repr_markdown_(self):
        """Markdown representation of the model, for Jupyter notebooks.
        """
        mod = '**Sukhbold_2015 Model**: {}\n\n'.format(self.filename)
        s = ['|Parameter|Value|',
             '|:---------|:-----:|',
             '|Progenitor mass | ${0.value:g}$ {0.unit:latex}|'.format(self.progenitor_mass),
             '|EOS | {}|'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

class Bollig_2016(SupernovaModel):
    """Set up a model based on simulations from Bollig et al. (2016). Models were taken, with permission, from the Garching Supernova Archive.
    """
    def __init__(self, filename, eos='LS220'):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to file prefix, we add nue/nuebar/nux.
        eos : string
            Equation of state used in simulation.
        """
        self.time = {}
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}

        # Store model metadata.
        self.filename = os.path.basename(filename)
        self.EOS = eos
        self.progenitor_mass = float(self.filename.strip('s%c')) * u.Msun

        # Read through the several ASCII files for the chosen simulation and
        # merge the data into one giant table.
        mergtab = None
        for flavor in Flavor:
            _flav = Flavor.NU_X if flavor == Flavor.NU_X_BAR else flavor
            _sfx = _flav.name.replace('_', '').lower()
            _filename = '{}_{}_{}'.format(filename, eos, _sfx)
            _lname  = 'L_{}'.format(flavor.name)
            _ename  = 'E_{}'.format(flavor.name)
            _e2name = 'E2_{}'.format(flavor.name)
            _aname  = 'ALPHA_{}'.format(flavor.name)

            simtab = Table.read(_filename,
                                names=['TIME', _lname, _ename, _e2name],
                                format='ascii')
            simtab['TIME'].unit = 's'
            simtab[_lname].unit = '1e51 erg/s'
            simtab[_aname] = (2*simtab[_ename]**2 - simtab[_e2name]) / (simtab[_e2name] - simtab[_ename]**2)
            simtab[_ename].unit = 'MeV'
            del simtab[_e2name]

            if mergtab is None:
                mergtab = simtab
            else:
                mergtab = join(mergtab, simtab, keys='TIME', join_type='left')
                mergtab[_lname].fill_value = 0.
                mergtab[_ename].fill_value = 0.
                mergtab[_aname].fill_value = 0.
        simtab = mergtab.filled()

        self.time = simtab['TIME'].to('s')

        for flavor in Flavor:
            # Set the dictionary of luminosity, mean energy, and shape
            # parameter keyed by NU_E, NU_X, NU_E_BAR, NU_X_BAR.
            _lname  = 'L_{}'.format(flavor.name)
            self.luminosity[flavor] = simtab[_lname].to('erg/s')

            _ename  = 'E_{}'.format(flavor.name)
            self.meanE[flavor] = simtab[_ename].to('MeV')

            _aname  = 'ALPHA_{}'.format(flavor.name)
            self.pinch[flavor] = simtab[_aname]

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.time

    def get_initialspectra(self,t,E):
        """Get neutrino spectra/luminosity curves before oscillation.

        Parameters
        ----------
        t : float
            Time to evaluate initial spectra.
        E : float or ndarray
            Energies to evaluate the initial spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor.
        """
        initialspectra = {}

        # Avoid division by zero in energy PDF below.
        E[E==0] = np.finfo(float).eps * E.unit

        # Estimate L(t), <E_nu(t)> and alpha(t). Express all energies in erg.
        E = E.to('erg').value

        # Make sure input time uses the same units as the model time grid, or
        # the interpolation will not work correctly.
        t = t.to(self.time.unit)

        for flavor in Flavor:
            # Use np.interp rather than scipy.interpolate.interp1d because it
            # can handle dimensional units (astropy.Quantity).
            L  = get_value(np.interp(t, self.time, self.luminosity[flavor].to('erg/s')))
            Ea = get_value(np.interp(t, self.time, self.meanE[flavor].to('erg')))
            a  = np.interp(t, self.time, self.pinch[flavor])

            # For numerical stability, evaluate log PDF and then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

    def __repr__(self):
        """Default representation of the model.
        """
        mod = 'Bollig_2016 Model: {}\n'.format(self.filename)
        s = ['Progenitor mass : {}'.format(self.progenitor_mass),
             'Eq. of state    : {}'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

    def _repr_markdown_(self):
        """Markdown representation of the model, for Jupyter notebooks.
        """
        mod = '**Bollig_2016 Model**: {}\n\n'.format(self.filename)
        s = ['|Parameter|Value|',
             '|:---------|:-----:|',
             '|Progenitor mass | ${0.value:g}$ {0.unit:latex}|'.format(self.progenitor_mass),
             '|EOS | {}|'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

class OConnor_2015(SupernovaModel):
    """Set up a model based on the black hole formation simulation in O'Connor (2015).
    """
    def __init__(self, filename, eos='LS220'):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to file prefix, we add nue/nuebar/nux
        eos : string
            Equation of state used in simulation
        """
        simtab = Table.read(filename,
                     names= ['TIME','L_NU_E','L_NU_E_BAR','L_NU_X',
                                    'E_NU_E','E_NU_E_BAR','E_NU_X',
                                    'RMS_NU_E','RMS_NU_E_BAR','RMS_NU_X'],
                     format='ascii')

        simtab['ALPHA_NU_E'] = (2.0*simtab['E_NU_E']**2 - simtab['RMS_NU_E']**2)/(simtab['RMS_NU_E']**2 - simtab['E_NU_E']**2)
        simtab['ALPHA_NU_E_BAR'] = (2.0*simtab['E_NU_E_BAR']**2 - simtab['RMS_NU_E_BAR']**2)/(simtab['RMS_NU_E_BAR']**2 - simtab['E_NU_E_BAR']**2)
        simtab['ALPHA_NU_X'] = (2.0*simtab['E_NU_X']**2 - simtab['RMS_NU_X']**2)/(simtab['RMS_NU_X']**2 - simtab['E_NU_X']**2)

        # SYB: double-check on this factor of 4. Should be factor of 2?
        simtab['L_NU_X'] /= 4.0

        self.filename = 'OConnor2015_s40WH07_LS220'
        self.EOS = eos
        self.progenitor_mass = 40 * u.Msun

        # Get grid of model times.
        self.time = simtab['TIME'] * u.s

        # Set up dictionary of luminosity, mean energy and shape parameter
        # alpha, keyed by neutrino flavor (NU_E, NU_X, NU_E_BAR, NU_X_BAR).
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}

        for flavor in Flavor:
            # Note: file only contains NU_E, NU_E_BAR, and NU_X, so double up
            # the use of NU_X for NU_X_BAR.
            _flav = Flavor.NU_X if flavor == Flavor.NU_X_BAR else flavor

            self.luminosity[flavor] = simtab['L_{}'.format(_flav.name)] * u.erg/u.s
            self.meanE[flavor] = simtab['E_{}'.format(_flav.name)] * u.MeV
            self.pinch[flavor] = simtab['ALPHA_{}'.format(_flav.name)]

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.time

    def get_initialspectra(self,t,E):
        """Get neutrino spectra/luminosity curves before oscillation.

        Parameters
        ----------
        t : float
            Time to evaluate initial and oscillated spectra.
        E : float or ndarray
            Energies to evaluate the initial and oscillated spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor.
        """
        initialspectra = {}

        # Avoid division by zero in energy PDF below.
        E[E==0] = np.finfo(float).eps * E.unit

        # Estimate L(t), <E_nu(t)> and alpha(t). Express all energies in erg.
        E = E.to('erg').value

        # Make sure input time uses the same units as the model time grid, or
        # the interpolation will not work correctly.
        t = t.to(self.time.unit)

        for flavor in Flavor:
            # Use np.interp rather than scipy.interpolate.interp1d because it
            # can handle dimensional units (astropy.Quantity).
            L  = get_value(np.interp(t, self.time, self.luminosity[flavor].to('erg/s')))
            Ea = get_value(np.interp(t, self.time, self.meanE[flavor].to('erg')))
            a  = np.interp(t, self.time, self.pinch[flavor])

            # For numerical stability, evaluate log PDF and then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

    def __repr__(self):
        """Default representation of the model.
        """
        mod = 'OConnor_2015 Model: {}\n'.format(self.filename)
        s = ['Progenitor mass : {}'.format(self.progenitor_mass),
             'Eq. of state    : {}'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

    def _repr_markdown_(self):
        """Markdown representation of the model, for Jupyter notebooks.
        """
        mod = '**OConnor_2015 Model**: {}\n\n'.format(self.filename)
        s = ['|Parameter|Value|',
             '|:---------|:-----:|',
             '|Progenitor mass | ${0.value:g}$ {0.unit:latex}|'.format(self.progenitor_mass),
             '|EOS | {}|'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

class Warren_2020(SupernovaModel):
    """Set up a model based on simulations from Warren et al. (2020)."""

    def __init__(self, filename, eos='LS220'):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to file prefix, we add nue/nuebar/nux
        eos : string
            Equation of state used in simulation
        """
        # Read data from HDF5 files, then store.
        f = h5py.File(filename, 'r')
        simtab = Table()

        for i in range(len(f['nue_data']['lum'])):
            if f['sim_data']['shock_radius'][i][1] > 0.00001:
                bounce = f['sim_data']['shock_radius'][i][0]
                break

        simtab['TIME'] = f['nue_data']['lum'][:, 0] - bounce
        simtab['L_NU_E'] = f['nue_data']['lum'][:, 1] * 1e51
        simtab['L_NU_E_BAR'] = f['nuae_data']['lum'][:, 1] * 1e51
        simtab['L_NU_X'] = f['nux_data']['lum'][:, 1] * 1e51 / 4.0
        simtab['E_NU_E'] = f['nue_data']['avg_energy'][:, 1]
        simtab['E_NU_E_BAR'] = f['nuae_data']['avg_energy'][:, 1]
        simtab['E_NU_X'] = f['nux_data']['avg_energy'][:, 1]
        simtab['RMS_NU_E'] = f['nue_data']['rms_energy'][:, 1]
        simtab['RMS_NU_E_BAR'] = f['nuae_data']['rms_energy'][:, 1]
        simtab['RMS_NU_X'] = f['nux_data']['rms_energy'][:, 1]

        simtab['ALPHA_NU_E'] = (2.0 * simtab['E_NU_E'] ** 2 - simtab['RMS_NU_E'] ** 2) / (simtab['RMS_NU_E'] ** 2 - simtab['E_NU_E'] ** 2)
        simtab['ALPHA_NU_E_BAR'] = (2.0 * simtab['E_NU_E_BAR'] ** 2 - simtab['RMS_NU_E_BAR'] ** 2) / (simtab['RMS_NU_E_BAR'] ** 2 - simtab['E_NU_E_BAR'] ** 2)
        simtab['ALPHA_NU_X'] = (2.0 * simtab['E_NU_X'] ** 2 - simtab['RMS_NU_X'] ** 2) / (simtab['RMS_NU_X'] ** 2 - simtab['E_NU_X'] ** 2)

        # Set model metadata.
        self.filename = os.path.basename(filename)
        self.EOS = eos
        self.progenitor_mass = float(filename.split('_')[-1].strip('m%.h5')) * u.Msun
        self.turbmixing_param = float(filename.split('_')[-2].strip('a%'))

        # Get grid of model times.
        self.time = simtab['TIME'] * u.s

        # Set up dictionary of luminosity, mean energy and shape parameter
        # alpha, keyed by neutrino flavor (NU_E, NU_X, NU_E_BAR, NU_X_BAR).
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}

        for flavor in Flavor:
            # Note: file only contains NU_E, NU_E_BAR, and NU_X, so double up
            # the use of NU_X for NU_X_BAR.
            _flav = Flavor.NU_X if flavor == Flavor.NU_X_BAR else flavor

            self.luminosity[flavor] = simtab['L_{}'.format(_flav.name)] * u.erg/u.s
            self.meanE[flavor] = simtab['E_{}'.format(_flav.name)] * u.MeV
            self.pinch[flavor] = simtab['ALPHA_{}'.format(_flav.name)]

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.time

    def get_initialspectra(self, t, E):
        """Get neutrino spectra/luminosity curves before oscillation.

        Parameters
        ----------
        t : float
            Time to evaluate initial and oscillated spectra.
        E : float or ndarray
            Energies to evaluate the initial and oscillated spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor.
        """
        initialspectra = {}

        # Avoid division by zero in energy PDF below.
        E[E==0] = np.finfo(float).eps * E.unit

        # Estimate L(t), <E_nu(t)> and alpha(t). Express all energies in erg.
        E = E.to('erg').value

        # Make sure input time uses the same units as the model time grid, or
        # the interpolation will not work correctly.
        t = t.to(self.time.unit)

        for flavor in Flavor:
            # Use np.interp rather than scipy.interpolate.interp1d because it
            # can handle dimensional units (astropy.Quantity).
            L  = get_value(np.interp(t, self.time, self.luminosity[flavor].to('erg/s')))
            Ea = get_value(np.interp(t, self.time, self.meanE[flavor].to('erg')))
            a  = np.interp(t, self.time, self.pinch[flavor])

            # For numerical stability, evaluate log PDF and then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

    def __repr__(self):
        """Default representation of the model.
        """
        mod = 'Warren_2020 Model: {}\n'.format(self.filename)
        s = ['Progenitor mass : {}'.format(self.progenitor_mass),
             'Turb. mix param : {}'.format(self.turbmixing_param),
             'Eq. of state    : {}'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

    def _repr_markdown_(self):
        """Markdown representation of the model, for Jupyter notebooks.
        """
        mod = '**Warren_2020 Model**: {}\n\n'.format(self.filename)
        s = ['|Parameter|Value|',
             '|:---------|:-----:|',
             '|Progenitor mass | ${0.value:g}$ {0.unit:latex}|'.format(self.progenitor_mass),
             '|Turb. mixing param. | {}|'.format(self.turbmixing_param),
             '|EOS | {}|'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

class Kuroda_2020(SupernovaModel):
    """Set up a model based on simulations from Kuroda et al. (2020)."""

    def __init__(self, filename, eos='LS220'):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to file prefix, we add nue/nuebar/nux
        eos : string
            Equation of state used in simulation
        """
        # Load up model metadata.
        self.filename = filename
        self.EOS = eos

        # Read ASCII data.
        simtab = Table.read(filename, format='ascii')

        # Get grid of model times.
        self.time = (simtab['Tpb[ms]'] * u.ms).to('s')

        # Set up dictionary of luminosity, mean energy and shape parameter
        # alpha, keyed by neutrino flavor (NU_E, NU_X, NU_E_BAR, NU_X_BAR).
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}

        for flavor in Flavor:
            # Note: file only contains NU_E, NU_E_BAR, and NU_X, so double up
            # the use of NU_X for NU_X_BAR.
            _flav = Flavor.NU_X if flavor == Flavor.NU_X_BAR else flavor
            if _flav.is_neutrino:
                _fkey = _flav.name.lower()
            else:
                _fkey = _flav.name.strip('%_BAR').lower().replace('_', '_a')

            self.luminosity[flavor] = simtab['<L{}>'.format(_fkey)] * 1e51 * u.erg/u.s
            self.meanE[flavor] = simtab['<E{}>'.format(_fkey)] * u.MeV

            # There is no pinch parameter so use alpha=2.0.
            self.pinch[flavor] = np.full_like(self.meanE[flavor].value, 2.)

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.time

    def get_initialspectra(self, t, E):
        """Get neutrino spectra/luminosity curves before oscillation.

        Parameters
        ----------
        t : float
            Time to evaluate initial and oscillated spectra.
        E : float or ndarray
            Energies to evaluate the initial and oscillated spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor.
        """
        initialspectra = {}

        # Avoid division by zero in energy PDF below.
        E[E==0] = np.finfo(float).eps * E.unit

        # Estimate L(t), <E_nu(t)> and alpha(t). Express all energies in erg.
        E = E.to('erg').value

        # Make sure input time uses the same units as the model time grid, or
        # the interpolation will not work correctly.
        t = t.to(self.time.unit)

        for flavor in Flavor:
            # Use np.interp rather than scipy.interpolate.interp1d because it
            # can handle dimensional units (astropy.Quantity).
            L  = get_value(np.interp(t, self.time, self.luminosity[flavor].to('erg/s')))
            Ea = get_value(np.interp(t, self.time, self.meanE[flavor].to('erg')))
            a  = np.interp(t, self.time, self.pinch[flavor])

            # For numerical stability, evaluate log PDF and then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

    def __repr__(self):
        """Default representation of the model.
        """
        mod = 'Kuroda_2020 Model: {}\n'.format(self.filename)
        s = ['Eq. of state    : {}'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

    def _repr_markdown_(self):
        """Markdown representation of the model, for Jupyter notebooks.
        """
        mod = '**Kuroda_2020 Model**: {}\n\n'.format(self.filename)
        s = ['|Parameter|Value|',
             '|:---------|:-----:|',
             '|EOS | {}|'.format(self.EOS)
             ]
        return mod + '\n'.join(s)

class Janka(SupernovaModel):
    """Set up a model based on simulations from Janka, I'll have to update this descriptioin later because I dont know where this is from
    """

    def __init__(self, filename):
        """Initialize model.

        Parameters
        ----------
        filename : str
            Absolute or relative path to FITS file with model data.
        """
        self.file = Table.read(filename)
        self.filename = filename
        self.luminosity = {}
        self.meanE = {}
        self.pinch = {}
        for flavor in Flavor:
            self.luminosity[flavor] = interp1d(self.get_time(), self.get_luminosity(flavor))
            self.meanE[flavor] = interp1d(self.get_time(), self.get_mean_energy(flavor))
            self.pinch[flavor] = interp1d(self.get_time(), self.get_pinch_param(flavor))

    def get_time(self):
        """Get grid of model times.

        Returns
        -------
        time : ndarray
            Grid of times used in the model.
        """
        return self.file['TIME']

    def get_luminosity(self, flavor):
        """Get model luminosity L_nu.

        Parameters
        ----------
        flavor : Flavor
            Neutrino flavor type.

        Returns
        -------
        luminosity : ndarray
            Grid of luminosity values (erg/s) for this flavor.
        """
        if flavor == Flavor.NU_X_BAR:
            flavor = Flavor.NU_X
        return self.file['L_{}'.format(flavor.name.upper())]

    def get_mean_energy(self, flavor):
        """Get model mean energy <E_nu>.

        Parameters
        ----------
        flavor : Flavor
            Neutrino flavor type.

        Returns
        -------
        energy : ndarray
            Grid of mean energy versus time.
        """
        if flavor == Flavor.NU_X_BAR:
            flavor = Flavor.NU_X
        return self.file['E_{}'.format(flavor.name.upper())]

    def get_pinch_param(self, flavor):
        """Get spectral pinch parameter alpha(t).

        Parameters
        ----------
        flavor : Flavor
            Neutrino flavor type.

        Returns
        -------
        alpha : ndarray
            Grid of alpha versus time.
        """
        if (flavor == Flavor.NU_X_BAR):
            flavor = Flavor.NU_X
        return self.file['ALPHA_{}'.format(flavor.name.upper())]

    def get_EOS(self):
        """Model equation of state.

        Returns
        -------
        eos : str
            Model equation of state.
        """
        return self.filename.split('-')[1]

    def get_progenitor_mass(self):
        """Progenitor mass.

        Returns
        -------
        mass : float
            Progenitor mass, in units of solar mass.
        """
        return float(self.split('-')[-1].split('.')[0].strip('s'))

    def get_initialspectra(self,t,E):
        """Get neutrino spectra/luminosity curves after oscillation.

        Parameters
        ----------
        t : float
            Time to evaluate initial and oscillated spectra.
        E : float or ndarray
            Energies to evaluate the initial and oscillated spectra.

        Returns
        -------
        initialspectra : dict
            Dictionary of model spectra, keyed by neutrino flavor."""

        initialspectra = {}
        for flavor in Flavor:
            L = self.luminosity[flavor](t)
            Ea = self.meanE[flavor](t)          # <E_nu(t)>
            Ea = Ea*1e6 * 1.60218e-12
            a = self.pinch[flavor](t)           # alpha_nu(t)
            E[E==0] = np.finfo(float).eps       # Avoid division by zero.

            # For numerical stability, evaluate log PDF then exponentiate.
            initialspectra[flavor] = \
                np.exp(np.log(L) - (2+a)*np.log(Ea) + (1+a)*np.log(1+a)
                       - loggamma(1+a) + a*np.log(E) - (1+a)*(E/Ea)) / (u.erg * u.s)

        return initialspectra

# class Fornax2019(SupernovaModel):
#
#     def __init__(self, filename):
#         self.file = Table.read(filename)
#         self.filename = filename
#
#     def get_time(self):
#         return self.file['TIME']
#
#     def get_luminosity(self, flavor):
#         if flavor == Flavor.NU_X_BAR:
#             flavor = Flavor.NU_X
#         return self.file['L_{}'.format(flavor.name.upper())]
#
#     def get_mean_energy(self, flavor):
#         if flavor == Flavor.NU_X_BAR:
#             flavor = Flavor.NU_X
#         return self.file['E_{}'.format(flavor.name.upper())]
#
#     def get_pinch_param(self, flavor):
#         if (flavor == Flavor.NU_X_BAR):
#             flavor = Flavor.NU_X
#         return self.file['ALPHA_{}'.format(flavor.name.upper())]
#
#     def get_EOS(self):
#         return self.filename.split('-')[1]
#
#     def get_progenitor_mass(self):
#         return float(self.split('-')[-1].split('.')[0].strip('s'))
#
# class OConnor_2013(SupernovaModel):
#
#     eos = 'LS220'
#     mass = 30
#     def __init__(self, filename, FlavorTransformation):
#         spectra_tuple = get_spectra(mass, eos)
#         self.eos = eos
#         self.luminosity = get_luminosity
#         self.meanE = spectra_list(2)
#         self.luminosity = get_luminosity(mass, eos)
#         self.
#
#     def get_luminosity(mass, eos):
#         # Open luminosity file.
#         tf = tarfile.open('{}_timeseries.tar.gz'.format(eos))
#
#         # Extract luminosity data.
#         dataname = 's{:d}_{}_timeseries.dat'.format(mass, eos)
#         datafile = tf.extractfile(dataname)
#         lumdata = ascii.read(datafile, names=['t', 'Le', 'Lae', 'Lx',
#                                               'Ee_avg', 'Eae_avg', 'Ex_avg',
#                                               'Ee_rms', 'Eae_rms', 'Ex_rms'])
#         return lumdata
#
#
#     def get_spectra(mass, eos):
#         # Open spectra file.
#         tf = tarfile.open('{}_timeseries_spectra.tar.gz'.format(eos))
#
#         # Extract luminosity data.
#         dataname = 's{:d}_{}_timeseries_spectra.dat'.format(mass, eos)
#         datafile = tf.extractfile(dataname)
#
#         t = []
#         E = []
#         spectra = []
#
#         for line in datafile:
#             tokens = [float(x) for x in line.strip().split()]
#             if len(tokens) == 1:
#                 if t:
#                     E = _E
#                     spectra.append(Table([_Fe, _Fae, _Fx],
#                                          names=['Fe', 'Fae', 'Fx'],
#                                          meta={'t' : t[-1]}))
#                 t.append(tokens[0])
#                 _E, _Fe, _Fae, _Fx = [[] for _ in range(4)]
#             elif len(tokens) == 4:
#                 _E.append(tokens[0])
#                 _Fe.append(tokens[1])
#                 _Fae.append(tokens[2])
#                 _Fx.append(tokens[3])
#
#         spectra.append(Table([_Fe, _Fae, _Fx],
#                              names=['Fe', 'Fae', 'Fx'],
#                              meta={'t' : t[-1]}))
#
#         return t, E, spectra

class SNOwGLoBES:
    """A model that does not inherit from SupernovaModel (yet) and imports a group of SNOwGLoBES files."""

    def __init__(self, tarfilename):
        """Initialize model from a tar archive.

        Parameters
        ----------
        tarfilename: str
            Absolute or relative path to tar archive with SNOwGLoBES files.
        """
        self.tfname = tarfilename
        tf = tarfile.open(self.tfname)

        # For now just pull out the "NoOsc" files.
        datafiles = sorted([f.name for f in tf if '.dat' in f.name])
        noosc = [df for df in datafiles if 'NoOsc' in df]
        noosc.sort(key=len)

        # Loop through the noosc files and pull out the number fluxes.
        self.time = []
        self.energy = None
        self.flux = {}
        self.fmin = 1e99
        self.fmax = -1e99

        for nooscfile in noosc:
            with tf.extractfile(nooscfile) as f:
                logging.debug('Reading {}'.format(nooscfile))
                meta = f.readline()
                metatext = meta.decode('utf-8')
                t = float(metatext.split('TBinMid=')[-1].split('sec')[0])
                dt = float(metatext.split('tBinWidth=')[-1].split('s')[0])
                dE = float(metatext.split('eBinWidth=')[-1].split('MeV')[0])

                data = Table.read(f, format='ascii.commented_header', header_start=-1)
                data.meta['t'] = t
                data.meta['dt'] = dt
                data.meta['dE'] = dE

                self.time.append(t)
                if self.energy is None:
                    self.energy = (data['E(GeV)'].data*1000).tolist()

            for flavor in ['NuE', 'NuMu', 'NuTau', 'aNuE', 'aNuMu', 'aNuTau']:
                if flavor in self.flux:
                    self.flux[flavor].append(data[flavor].data.tolist())
                else:
                    self.flux[flavor] = [data[flavor].data.tolist()]

        # We now have a table with rows=times and columns=energies. Transpose
        # so that rows=energy and cols=time.
        for k, v in self.flux.items():
            self.flux[k] = np.transpose(self.flux[k])
            self.fmin = np.minimum(self.fmin, np.min(self.flux[k]))
            self.fmax = np.maximum(self.fmax, np.max(self.flux[k]))

    def get_fluence(self, t):
        """Return the fluence at a given time t.

        Parameters
        ----------
        t : float
            Time in seconds.

        Returns
        -------
        fluence : dict
            A dictionary giving fluence at time t, keyed by flavor.
        """
        idx = get_closest(self.time, t)

        fluence = {}
        for k, fl in self.flux.items():
            fluence[k] = fl[:,idx]

        return fluence

import numpy,from scipy,import astropy,from astropy