#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#         Alexander Sokolov <alexander.y.sokolov@gmail.com>
#

'''
Selected CI using Heat-Bath CI algorithm
(JCTC 2016, 12, 3674-3680)

Simple usage::

'''

import numpy
import time
import ctypes
from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.fci import cistring
from pyscf.fci import direct_spin1

libhci = lib.load_library('libhci')

def contract_2e_ctypes(h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
    h1, eri = h1_h2
    strs = civec._strs
    ndet = len(strs)
    if hdiag is None:
        hdiag = make_hdiag(h1, eri, strs, norb, nelec)
    ci1 = numpy.zeros_like(civec)

    h1 = numpy.asarray(h1, order='C')
    eri = numpy.asarray(eri, order='C')
    strs = numpy.asarray(strs, order='C')
    civec = numpy.asarray(civec, order='C')
    hdiag = numpy.asarray(hdiag, order='C')
    ci1 = numpy.asarray(ci1, order='C')

    libhci.contract_h_c(h1.ctypes.data_as(ctypes.c_void_p),
                        eri.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_int(norb),
                        ctypes.c_int(nelec[0]),
                        ctypes.c_int(nelec[1]),
                        strs.ctypes.data_as(ctypes.c_void_p),
                        civec.ctypes.data_as(ctypes.c_void_p),
                        hdiag.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_ulonglong(ndet),
                        ci1.ctypes.data_as(ctypes.c_void_p))

    return ci1

def contract_2e(h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
    h1, eri = h1_h2
    strs = civec._strs
    ndet = len(strs)
    if hdiag is None:
        hdiag = make_hdiag(h1, eri, strs, norb, nelec)
    ci1 = numpy.zeros_like(civec)

    eri = eri.reshape([norb]*4)

    for ip in range(ndet):
        for jp in range(ip):
            stria, strib = strs[ip].reshape(2,-1)
            strja, strjb = strs[jp].reshape(2,-1)
            desa, crea = str_diff(stria, strja)
            if len(desa) > 2:
                continue
            desb, creb = str_diff(strib, strjb)
            if len(desb) + len(desa) > 2:
                continue
            if len(desa) + len(desb) == 1:
# alpha->alpha
                if len(desb) == 0:
                    i,a = desa[0], crea[0]
                    occsa = str2orblst(stria, norb)[0]
                    occsb = str2orblst(strib, norb)[0]
                    fai = h1[a,i]
                    for k in occsa:
                        fai += eri[k,k,a,i] - eri[k,i,a,k]
                    for k in occsb:
                        fai += eri[k,k,a,i]
                    sign = cre_des_sign(a, i, stria)
                    ci1[jp] += sign * fai * civec[ip]
                    ci1[ip] += sign * fai * civec[jp]
# beta ->beta
                elif len(desa) == 0:
                    i,a = desb[0], creb[0]
                    occsa = str2orblst(stria, norb)[0]
                    occsb = str2orblst(strib, norb)[0]
                    fai = h1[a,i]
                    for k in occsb:
                        fai += eri[k,k,a,i] - eri[k,i,a,k]
                    for k in occsa:
                        fai += eri[k,k,a,i]
                    sign = cre_des_sign(a, i, strib)
                    ci1[jp] += sign * fai * civec[ip]
                    ci1[ip] += sign * fai * civec[jp]

            else:
# alpha,alpha->alpha,alpha
                if len(desb) == 0:
                    i,j = desa
                    a,b = crea
# 6 conditions for i,j,a,b
# --++, ++--, -+-+, +-+-, -++-, +--+
                    if a > j or i > b:
# condition --++, ++--
                        v = eri[a,j,b,i]-eri[a,i,b,j]
                        sign = cre_des_sign(b, i, stria)
                        sign*= cre_des_sign(a, j, stria)
                    else:
# condition -+-+, +-+-, -++-, +--+
                        v = eri[a,i,b,j]-eri[a,j,b,i]
                        sign = cre_des_sign(b, j, stria)
                        sign*= cre_des_sign(a, i, stria)
                    ci1[jp] += sign * v * civec[ip]
                    ci1[ip] += sign * v * civec[jp]
# beta ,beta ->beta ,beta
                elif len(desa) == 0:
                    i,j = desb
                    a,b = creb
                    if a > j or i > b:
                        v = eri[a,j,b,i]-eri[a,i,b,j]
                        sign = cre_des_sign(b, i, strib)
                        sign*= cre_des_sign(a, j, strib)
                    else:
                        v = eri[a,i,b,j]-eri[a,j,b,i]
                        sign = cre_des_sign(b, j, strib)
                        sign*= cre_des_sign(a, i, strib)
                    ci1[jp] += sign * v * civec[ip]
                    ci1[ip] += sign * v * civec[jp]
# alpha,beta ->alpha,beta
                else:
                    i,a = desa[0], crea[0]
                    j,b = desb[0], creb[0]
                    v = eri[a,i,b,j]
                    sign = cre_des_sign(a, i, stria)
                    sign*= cre_des_sign(b, j, strib)
                    ci1[jp] += sign * v * civec[ip]
                    ci1[ip] += sign * v * civec[jp]
        ci1[ip] += hdiag[ip] * civec[ip]

    return ci1

def spin_square(civec, norb, nelec):
    ss = numpy.dot(civec.T, contract_ss(civec, norb, nelec))
    s = numpy.sqrt(ss+.25) - .5
    multip = s*2+1
    return ss, multip

def contract_ss(civec, norb, nelec):
    strs = civec._strs
    ndet = len(strs)
    ci1 = numpy.zeros_like(civec)

    strs = numpy.asarray(strs, order='C')
    civec = numpy.asarray(civec, order='C')
    ci1 = numpy.asarray(ci1, order='C')

    libhci.contract_ss_c(ctypes.c_int(norb),
                        ctypes.c_int(nelec[0]),
                        ctypes.c_int(nelec[1]),
                        strs.ctypes.data_as(ctypes.c_void_p),
                        civec.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_ulonglong(ndet),
                        ci1.ctypes.data_as(ctypes.c_void_p))

    return ci1

def make_hdiag(h1e, eri, strs, norb, nelec):
    eri = ao2mo.restore(1, eri, norb)
    diagj = numpy.einsum('iijj->ij',eri)
    diagk = numpy.einsum('ijji->ij',eri)

    ndet = len(strs)
    hdiag = numpy.zeros(ndet)
    for idet, (stra, strb) in enumerate(strs.reshape(ndet,2,-1)):
        aocc = str2orblst(stra, norb)[0]
        bocc = str2orblst(strb, norb)[0]
        e1 = h1e[aocc,aocc].sum() + h1e[bocc,bocc].sum()
        e2 = diagj[aocc][:,aocc].sum() + diagj[aocc][:,bocc].sum() \
           + diagj[bocc][:,aocc].sum() + diagj[bocc][:,bocc].sum() \
           - diagk[aocc][:,aocc].sum() - diagk[bocc][:,bocc].sum()
        hdiag[idet] = e1 + e2*.5
    return hdiag

def cre_des_sign(p, q, string):
    nset = len(string)
    pg, pb = p//64, p%64
    qg, qb = q//64, q%64

    if pg > qg:
        n1 = 0
        for i in range(nset-pg, nset-qg-1):
            n1 += bin(string[i]).count('1')
        n1 += bin(string[-1-pg] & numpy.uint64((1<<pb) - 1)).count('1')
        n1 += string[-1-qg] >> numpy.uint64(qb+1)
    elif pg < qg:
        n1 = 0
        for i in range(nset-qg, nset-pg-1):
            n1 += bin(string[i]).count('1')
        n1 += bin(string[-1-qg] & numpy.uint64((1<<qb) - 1)).count('1')
        n1 += string[-1-pg] >> numpy.uint64(pb+1)
    else:
        if p > q:
            mask = numpy.uint64((1 << pb) - (1 << (qb+1)))
        else:
            mask = numpy.uint64((1 << qb) - (1 << (pb+1)))
        n1 = bin(string[-1-pg]&mask).count('1')

    if n1 % 2:
        return -1
    else:
        return 1

def argunique(strs):
    def order(x, y):
        for i in range(y.size):
            if x[i] > y[i]:
                return 1
            elif y[i] > x[i]:
                return -1
        return 0
    def qsort_idx(idx):
        nstrs = len(idx)
        if nstrs <= 1:
            return idx
        else:
            ref = idx[-1]
            group_lt = []
            group_gt = []
            for i in idx[:-1]:
                c = order(strs[i], strs[ref])
                if c == -1:
                    group_lt.append(i)
                elif c == 1:
                    group_gt.append(i)
            return qsort_idx(group_lt) + [ref] + qsort_idx(group_gt)
    return qsort_idx(range(len(strs)))

def argunique_ctypes(strs):
    nstrs, nset = strs.shape

    sort_idx = numpy.empty(nstrs, dtype=numpy.uint64)

    strs = numpy.asarray(strs, order='C')
    sort_idx = numpy.asarray(sort_idx, order='C')

    nstrs_ = numpy.array([nstrs])

    libhci.argunique(strs.ctypes.data_as(ctypes.c_void_p),
                     sort_idx.ctypes.data_as(ctypes.c_void_p),
                     nstrs_.ctypes.data_as(ctypes.c_void_p),
                     ctypes.c_int(nset))

    sort_idx = sort_idx[:nstrs_[0]]

    return sort_idx.tolist()

def str_diff(string0, string1):
    des_string0 = []
    cre_string0 = []
    nset = len(string0)
    off = 0
    for i in reversed(range(nset)):
        df = string0[i] ^ string1[i]
        des_string0.extend([x+off for x in find1(df & string0[i])])
        cre_string0.extend([x+off for x in find1(df & string1[i])])
        off += 64
    return des_string0, cre_string0

def excitation_level(string, nelec=None):
    nset = len(string)
    if nelec is None:
        nelec = 0
        for i in range(nset):
            nelec += bin(string[i]).count('1')

    g, b = nelec//64, nelec%64
    tn = nelec - bin(string[-1-g])[-b:].count('1')
    for s in string[nset-g:]:
        tn -= bin(s).count('1')
    return tn

def find1(s):
    return [i for i,x in enumerate(bin(s)[2:][::-1]) if x is '1']

def toggle_bit(s, place):
    nset = len(s)
    g, b = place//64, place%64
    s[-1-g] ^= numpy.uint64(1<<b)
    return s

def select_strs_ctypes(myci, civec, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec):
    strs = civec._strs
    ndet = strs.shape[0]
    ndet, nset = strs.shape
    nset = nset // 2
    neleca, nelecb = nelec

    h1 = numpy.asarray(h1, order='C')
    eri = numpy.asarray(eri, order='C')
    jk = numpy.asarray(jk, order='C')
    civec = numpy.asarray(civec, order='C')
    strs = numpy.asarray(strs, order='C')
    eri_sorted = numpy.asarray(eri_sorted, order='C')
    jk_sorted = numpy.asarray(jk_sorted, order='C')

    str_add = numpy.empty((0,strs.shape[1]), dtype=numpy.uint64)

    batch_size = max(1, 8 * 4 * neleca * nelecb * (norb-neleca) * (norb-nelecb))
    ndet_batch = int(myci.max_memory * 1024**2) // batch_size
    nbatches = ndet // ndet_batch + 1

    for i in range(nbatches):
        ndet_start = ndet_batch * i
        ndet_finish = min(ndet_batch * (i + 1), ndet)
        ndet_select_max = 4 * neleca * nelecb * (norb-neleca) * (norb-nelecb) * ndet_batch

        str_add_batch = numpy.empty((ndet_select_max, strs.shape[1]), dtype=numpy.uint64)
        n_str_add_batch = numpy.array([str_add_batch.shape[0]])

        str_add_batch = numpy.asarray(str_add_batch, order='C')

        libhci.select_strs(h1.ctypes.data_as(ctypes.c_void_p),
                           eri.ctypes.data_as(ctypes.c_void_p),
                           jk.ctypes.data_as(ctypes.c_void_p),
                           eri_sorted.ctypes.data_as(ctypes.c_void_p),
                           jk_sorted.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(norb),
                           ctypes.c_int(neleca),
                           ctypes.c_int(nelecb),
                           strs.ctypes.data_as(ctypes.c_void_p),
                           civec.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_ulonglong(ndet_start),
                           ctypes.c_ulonglong(ndet_finish),
                           ctypes.c_double(myci.select_cutoff),
                           str_add_batch.ctypes.data_as(ctypes.c_void_p),
                           n_str_add_batch.ctypes.data_as(ctypes.c_void_p))

        n_str_add_batch = n_str_add_batch[0]
        str_add_batch = str_add_batch[:n_str_add_batch]
        str_add = numpy.vstack((str_add, str_add_batch))

    str_add = numpy.asarray(str_add)
    return str_add

def enlarge_space(myci, civec, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec):
    if not isinstance(civec, (tuple, list)):
        civec = [civec]

    strs = civec[0]._strs

    nroots = len(civec)

    cidx = abs(civec[0]) > myci.ci_coeff_cutoff
    for p in range(1,nroots):
        cidx += abs(civec[p]) > myci.ci_coeff_cutoff

    strs = strs[cidx]

    ci_coeff = [as_SCIvector(c[cidx], strs) for c in civec]

    strs_new = strs.copy()

    for p in range(nroots):
        str_add = select_strs_ctypes(myci, ci_coeff[p], h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)
        strs_new = numpy.vstack((strs, str_add))

    # Add strings together and remove duplicate strings
    tmp = numpy.ascontiguousarray(strs_new).view(numpy.dtype((numpy.void, strs_new.dtype.itemsize * strs_new.shape[1])))
    _, tmpidx = numpy.unique(tmp, return_index=True)

    new_ci = []
    for p in range(nroots):
        c = numpy.zeros(strs_new.shape[0])
        c[:ci_coeff[p].shape[0]] = ci_coeff[p]
        new_ci.append(c[tmpidx])

    strs_new = strs_new[tmpidx]

    return [as_SCIvector(ci, strs_new) for ci in new_ci]

def str2orblst(string, norb):
    occ = []
    vir = []
    nset = len(string)
    off = 0
    for k in reversed(range(nset)):
        s = string[k]
        occ.extend([x+off for x in find1(s)])
        for i in range(0, min(64, norb-off)):
            if not (s & numpy.uint64(1<<i)):
                vir.append(i+off)
        off += 64
    return occ, vir

def orblst2str(lst, norb):
    nset = (norb+63) // 64
    string = numpy.zeros(nset, dtype=numpy.uint64)
    for i in lst:
        toggle_bit(string, i)
    return string

def kernel_float_space(myci, h1e, eri, norb, nelec, ci0=None,
                       tol=None, lindep=None, max_cycle=None, max_space=None,
                       nroots=None, davidson_only=None, max_iter=None,
                       max_memory=None, verbose=None, ecore=0, return_integrals=False,
                       eri_sorted=None, jk=None, jk_sorted=None, **kwargs):
    if verbose is None:
        log = logger.Logger(myci.stdout, myci.verbose)
    elif isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(myci.stdout, verbose)
    if tol is None: tol = myci.conv_tol
    if lindep is None: lindep = myci.lindep
    if max_cycle is None: max_cycle = myci.max_cycle
    if max_space is None: max_space = myci.max_space
    if max_memory is None: max_memory = myci.max_memory
    if nroots is None: nroots = myci.nroots
    if max_iter is None: max_iter = myci.max_iter
    if myci.verbose >= logger.WARN:
        myci.check_sanity()

    log.info('\nStarting heat-bath CI algorithm...\n')
    log.info('Selection threshold:                  %8.5e',    myci.select_cutoff)
    log.info('CI coefficient cutoff:                %8.5e',    myci.ci_coeff_cutoff)
    log.info('Energy convergence tolerance:         %8.5e',    tol)
    log.info('Number of determinants tolerance:     %8.5e',    myci.conv_ndet_tol)
    log.info('Number of electrons:                  %s',       nelec)
    log.info('Number of orbitals:                   %3d',      norb)
    log.info('Number of roots:                      %3d',    nroots)

    nelec = direct_spin1._unpack_nelec(nelec, myci.spin)
    eri = ao2mo.restore(1, eri, norb)

    # Avoid resorting the integrals by storing them in memory
    eri = eri.ravel()

    if eri_sorted is None and jk is None and jk_sorted is None:
        log.debug("\nSorting two-electron integrals...")
        t_start = time.time()
        eri_sorted = abs(eri).argsort()[::-1]
        jk = eri.reshape([norb]*4)
        jk = jk - jk.transpose(2,1,0,3)
        jk = jk.ravel()
        jk_sorted = abs(jk).argsort()[::-1]
        t_current = time.time() - t_start
        log.debug('Timing for sorting the integrals: %10.3f', t_current)

    # Initial guess
    if ci0 is None:
        hf_str = numpy.hstack([orblst2str(range(nelec[0]), norb), orblst2str(range(nelec[1]), norb)]).reshape(1,-1)
        ci0 = [as_SCIvector(numpy.ones(1), hf_str)]
    else:
        assert(nroots == len(ci0))

    ci0 = myci.enlarge_space(ci0, h1e, eri, jk, eri_sorted, jk_sorted, norb, nelec)

    def hop(c):
        hc = myci.contract_2e((h1e, eri), as_SCIvector(c, ci_strs), norb, nelec, hdiag)
        return hc.ravel()
    precond = lambda x, e, *args: x/(hdiag-e+myci.level_shift)

    e_last = 0
    float_tol = 3e-4
    conv = False
    for icycle in range(max_iter):
        ci_strs = ci0[0]._strs
        float_tol = max(float_tol*.3, tol*1e2)
        log.info('\nMacroiteration %d', icycle)
        log.info('Number of CI configurations: %d', ci_strs.shape[0])
        hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
        t_start = time.time()
        e, ci0 = myci.eig(hop, ci0, precond, tol=float_tol, lindep=lindep,
                          max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                          max_memory=max_memory, verbose=log, **kwargs)
        if not isinstance(ci0, (tuple, list)):
            ci0 = [ci0]
            e = [e]
        t_current = time.time() - t_start
        log.debug('Timing for solving the eigenvalue problem: %10.3f', t_current)
        ci0 = [as_SCIvector(c, ci_strs) for c in ci0]
        de, e_last = min(e)-e_last, min(e)
        log.info('Cycle %d  E = %s  dE = %.8g', icycle, numpy.array(e)+ecore, de)

        if abs(de) < tol*1e3:
            conv = True
            break

        last_ci0_size = float(len(ci_strs))
        t_start = time.time()
        ci0 = myci.enlarge_space(ci0, h1e, eri, jk, eri_sorted, jk_sorted, norb, nelec)
        t_current = time.time() - t_start
        log.debug('Timing for selecting configurations: %10.3f', t_current)
        if (((1 - myci.conv_ndet_tol) < len(ci0[0]._strs)/last_ci0_size < (1 + myci.conv_ndet_tol))):
            conv = True
            break

    ci_strs = ci0[0]._strs
    log.info('\nExtra CI in the final selected space')
    log.info('Number of CI configurations: %d', ci_strs.shape[0])
    hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
    e, c = myci.eig(hop, ci0, precond, tol=tol, lindep=lindep,
                    max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                    max_memory=max_memory, verbose=log, **kwargs)
    if not isinstance(c, (tuple, list)):
        c = [c]
        e = [e]
    log.info('\nSelected CI  E = %s', numpy.array(e)+ecore)

    if (return_integrals):
        return (numpy.array(e)+ecore), [as_SCIvector(ci, ci_strs) for ci in c], eri_sorted, jk, jk_sorted
    else:
        return (numpy.array(e)+ecore), [as_SCIvector(ci, ci_strs) for ci in c]

def fix_spin(myci, shift=.2, ss=None, **kwargs):
    r'''If Selected CI solver cannot stick on spin eigenfunction, modify the solver by
    adding a shift on spin square operator

    .. math::

        (H + shift*S^2) |\Psi\rangle = E |\Psi\rangle

    Args:
        myci : An instance of :class:`SelectedCI`

    Kwargs:
        shift : float
            Level shift for states which have different spin
        ss : number
            S^2 expection value == s*(s+1)

    Returns
            A modified Selected CI object based on myci.
    '''
    if 'ss_value' in kwargs:
        sys.stderr.write('fix_spin_: kwarg "ss_value" will be removed in future release. '
                         'It was replaced by "ss"\n')
        ss_value = kwargs['ss_value']
    else:
        ss_value = ss

    def contract_2e(h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
        if isinstance(nelec, (int, numpy.number)):
            sz = (nelec % 2) * .5
        else:
            sz = abs(nelec[0]-nelec[1]) * .5
        if ss_value is None:
            ss = sz*(sz+1)
        else:
            ss = ss_value

        h1, eri = h1_h2
        strs = civec._strs
        ndet = len(strs)
        if hdiag is None:
            hdiag = make_hdiag(h1, eri, strs, norb, nelec)
        ci1 = numpy.zeros_like(civec)
        ci2 = numpy.zeros_like(civec)

        h1 = numpy.asarray(h1, order='C')
        eri = numpy.asarray(eri, order='C')
        strs = numpy.asarray(strs, order='C')
        civec = numpy.asarray(civec, order='C')
        hdiag = numpy.asarray(hdiag, order='C')
        ci1 = numpy.asarray(ci1, order='C')
        ci2 = numpy.asarray(ci2, order='C')

        libhci.contract_h_c_ss_c(h1.ctypes.data_as(ctypes.c_void_p),
                                 eri.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_int(norb),
                                 ctypes.c_int(nelec[0]),
                                 ctypes.c_int(nelec[1]),
                                 strs.ctypes.data_as(ctypes.c_void_p),
                                 civec.ctypes.data_as(ctypes.c_void_p),
                                 hdiag.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_ulonglong(ndet),
                                 ci1.ctypes.data_as(ctypes.c_void_p),
                                 ci2.ctypes.data_as(ctypes.c_void_p))

        if ss < sz*(sz+1)+.1:
# (S^2-ss)|Psi> to shift state other than the lowest state
            ci2 -= ss * civec
        else:
# (S^2-ss)^2|Psi> to shift states except the given spin.
# It still relies on the quality of initial guess
            tmp = ci2.copy()
            tmp -= ss * civec
            ci2 = -ss * tmp
            ci2 += myci.contract_ss(as_SCIvector_if_not(tmp, strs), norb, nelec)
            tmp = None
        ci2 *= shift
        ci1 += ci2

        return as_SCIvector_if_not(ci1, strs)

    myci.contract_2e = contract_2e
    return myci

def to_fci(civec, norb, nelec, root=0):
    assert(norb <= 64)
    neleca, nelecb = nelec
    strsa = cistring.gen_strings4orblist(range(norb), neleca)
    stradic = dict(zip(strsa,range(strsa.__len__())))
    strsb = cistring.gen_strings4orblist(range(norb), nelecb)
    strbdic = dict(zip(strsb,range(strsb.__len__())))
    na = len(stradic)
    nb = len(strbdic)
    ndet = len(civec[root])
    fcivec = numpy.zeros((na,nb))
    for idet, (stra, strb) in enumerate(civec[root]._strs.reshape(ndet,2,-1)):
        ka = stradic[stra[0]]
        kb = strbdic[strb[0]]
        fcivec[ka,kb] = civec[root][idet]
    return fcivec

def from_fci(fcivec, ci_strs, norb, nelec):
    neleca, nelecb = nelec
    strsa = cistring.gen_strings4orblist(range(norb), neleca)
    stradic = dict(zip(strsa,range(strsa.__len__())))
    strsb = cistring.gen_strings4orblist(range(norb), nelecb)
    strbdic = dict(zip(strsb,range(strsb.__len__())))
    na = len(stradic)
    nb = len(strbdic)
    fcivec = fcivec.reshape(na,nb)
    ta = [excitation_level(s, neleca) for s in strsa.reshape(-1,1)]
    tb = [excitation_level(s, nelecb) for s in strsb.reshape(-1,1)]
    ndet = len(ci_strs)
    civec = numpy.zeros(ndet)
    for idet, (stra, strb) in enumerate(ci_strs.reshape(ndet,2,-1)):
        ka = stradic[stra[0]]
        kb = strbdic[strb[0]]
        civec[idet] = fcivec[ka,kb]
    return as_SCIvector(civec, ci_strs)

def make_rdm12s(civec, norb, nelec):
    '''Spin orbital 1- and 2-particle reduced density matrices (aa, bb, aaaa, aabb, bbbb)
    '''
    strs = civec._strs
    ndet = len(strs)
    rdm1a = numpy.zeros(norb*norb)
    rdm1b = numpy.zeros(norb*norb)
    rdm2aa = numpy.zeros(norb*norb*norb*norb)
    rdm2ab = numpy.zeros(norb*norb*norb*norb)
    rdm2bb = numpy.zeros(norb*norb*norb*norb)

    civec = numpy.asarray(civec, order='C')
    strs = numpy.asarray(strs, order='C')
    rdm1a  = numpy.asarray(rdm1a, order='C')
    rdm1b  = numpy.asarray(rdm1b, order='C')
    rdm2aa = numpy.asarray(rdm2aa, order='C')
    rdm2ab = numpy.asarray(rdm2ab, order='C')
    rdm2bb = numpy.asarray(rdm2bb, order='C')

    # Compute 1- and 2-RDMs
    libhci.compute_rdm12s(ctypes.c_int(norb),
                          ctypes.c_int(nelec[0]),
                          ctypes.c_int(nelec[1]),
                          strs.ctypes.data_as(ctypes.c_void_p),
                          civec.ctypes.data_as(ctypes.c_void_p),
                          ctypes.c_ulonglong(ndet),
                          rdm1a.ctypes.data_as(ctypes.c_void_p),
                          rdm1b.ctypes.data_as(ctypes.c_void_p),
                          rdm2aa.ctypes.data_as(ctypes.c_void_p),
                          rdm2ab.ctypes.data_as(ctypes.c_void_p),
                          rdm2bb.ctypes.data_as(ctypes.c_void_p))

    rdm1a = rdm1a.reshape([norb]*2)
    rdm1b = rdm1b.reshape([norb]*2)
    rdm2aa = rdm2aa.reshape([norb]*4)
    rdm2ab = rdm2ab.reshape([norb]*4)
    rdm2bb = rdm2bb.reshape([norb]*4)

    # Sort 2-RDM into chemists' notation: <p_1 q_2|r_1 s_2> -> (p_1 r_1| q_2 s_2)
    rdm2aa = rdm2aa.transpose(0,2,1,3)
    rdm2ab = rdm2ab.transpose(0,2,1,3)
    rdm2bb = rdm2bb.transpose(0,2,1,3)

    return (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb)

class SelectedCI(direct_spin1.FCISolver):
    def __init__(self, mol=None):
        direct_spin1.FCISolver.__init__(self, mol)
        self.ci_coeff_cutoff = .5e-3
        self.select_cutoff = .5e-3
        self.conv_tol = 1e-9
        self.conv_ndet_tol = 0.001
        self.nroots = 1
        self.max_iter = 10
        # Maximum memory in MB for storing lists of selected strings
        self.max_memory = 1000

##################################################
# don't modify the following attributes, they are not input options
        #self.converged = False
        #self.ci = None
        self._strs = None
        self._keys = set(self.__dict__.keys())

    def dump_flags(self, verbose=None):
        direct_spin1.FCISolver.dump_flags(self, verbose)
        logger.info(self, 'ci_coeff_cutoff %g', self.ci_coeff_cutoff)
        logger.info(self, 'select_cutoff   %g', self.select_cutoff)

    # define absorb_h1e for compatibility to other FCI solver
    def absorb_h1e(h1, eri, *args, **kwargs):
        return (h1, eri)

    def contract_2e(self, h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
        if getattr(civec, '_strs', None) is not None:
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)
        return contract_2e_ctypes(h1_h2, civec, norb, nelec, hdiag, **kwargs)
#        return contract_2e(h1_h2, civec, norb, nelec, hdiag, **kwargs)

    def contract_ss(self, civec, norb, nelec):
        if getattr(civec, '_strs', None) is not None:
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)
        return contract_ss(civec, norb, nelec)

    def spin_square(self, civec, norb, nelec):
        if getattr(civec, '_strs', None) is not None:
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)
        return spin_square(civec, norb, nelec)

    def make_hdiag(self, h1e, eri, strs, norb, nelec):
        return make_hdiag(h1e, eri, strs, norb, nelec)

    def to_fci(self, civec, norb, nelec):

        if getattr(civec, '_strs', None) is not None:
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)

        return to_fci(civec, norb, nelec)

    def make_rdm12s(self, civec, norb, nelec):

        if getattr(civec, '_strs', None) is not None:
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)

        return make_rdm12s(civec, norb, nelec)

    enlarge_space = enlarge_space
    kernel = kernel_float_space

SCI = SelectedCI

class _SCIvector(numpy.ndarray):
    def __array_finalize__(self, obj):
        self._strs = getattr(obj, '_strs', None)

def as_SCIvector(civec, ci_strs):
    civec = civec.view(_SCIvector)
    civec._strs = ci_strs
    return civec

def as_SCIvector_if_not(civec, ci_strs):
    if getattr(civec, '_strs', None) is None:
        civec = as_SCIvector(civec, ci_strs)
    return civec

if __name__ == '__main__':
    numpy.random.seed(3)
    strs = (numpy.random.random((14,3)) * 4).astype(numpy.uint64)
    print(strs)
    print(argunique(strs))

    norb = 6
    nelec = 3,3
    hf_str = numpy.hstack([orblst2str(range(nelec[0]), norb),
                           orblst2str(range(nelec[1]), norb)]).reshape(1,-1)
    numpy.random.seed(3)
    h1 = numpy.random.random([norb]*2)**4 * 1e-2
    h1 = h1 + h1.T
    eri = numpy.random.random([norb]*4)**4 * 1e-2
    eri = eri + eri.transpose(0,1,3,2)
    eri = eri + eri.transpose(1,0,2,3)
    eri = eri + eri.transpose(2,3,0,1)
    eri_sorted = abs(eri).argsort()[::-1]
    jk = eri.reshape([norb]*4)
    jk = jk - jk.transpose(2,1,0,3)
    jk = jk.ravel()
    jk_sorted = abs(jk).argsort()[::-1]
    ci1 = [as_SCIvector(numpy.ones(1), hf_str)]

    myci = SelectedCI()
    myci.select_cutoff = .001
    myci.ci_coeff_cutoff = .001

    ci2 = enlarge_space(myci, ci1, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)
    print(len(ci2[0]))

    ci2 = enlarge_space(myci, ci1, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)
    numpy.random.seed(1)
    ci3 = numpy.random.random(ci2[0].size)
    ci3 *= 1./numpy.linalg.norm(ci3)
    ci3 = [ci3]
    ci3 = enlarge_space(myci, ci2, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)

    efci = direct_spin1.kernel(h1, eri, norb, nelec, verbose=5)[0]

    ci4 = contract_2e_ctypes((h1, eri), ci3[0], norb, nelec)

    fci3 = to_fci(ci3, norb, nelec)
    h2e = direct_spin1.absorb_h1e(h1, eri, norb, nelec, .5)
    fci4 = direct_spin1.contract_2e(h2e, fci3, norb, nelec)
    fci4 = from_fci(fci4, ci3[0]._strs, norb, nelec)
    print(abs(ci4-fci4).sum())

    e = myci.kernel(h1, eri, norb, nelec, verbose=5)[0]
    print(e, efci)


##########################
##########################
###                    ###
###  MUTE              ###
###  William Woodley   ###
###  19 December 2021  ###
###                    ###
##########################
##########################

# Import packages

import os

import numpy as np
from tqdm import tqdm

import mute.constants as constants

try:

    import proposal as pp

except ImportError:

    pass

# Create the propagator

def _create_propagator(force):

    """This function creates the propagator object in PROPOSAL for use in _propagation_loop()"""

    # Check values

    constants.check_constants(force=force)

    # Global variables
    # The propagator is used in every iteration of the doubly-nested propagation loop
    # Make it a global variable so it only has to be created once

    global propagator

    if constants.get_verbose() > 1:
        print("Creating propagator.")

    # Propagator arguments

    mu = pp.particle.MuMinusDef()
    cuts = pp.EnergyCutSettings(500, 0.05, True)

    if constants.get_medium() == "rock":

        medium = pp.medium.StandardRock()

    elif constants.get_medium() == "water":

        medium = pp.medium.Water()

    elif constants.get_medium() == "ice":

        medium = pp.medium.Ice()

    elif constants.get_medium() == "air":

        medium = pp.medium.Air()

    else:

        raise NotImplementedError(
            "Medium type {0} not implemented.".format(constants.get_medium())
        )

    args = {"particle_def": mu, "target": medium, "interpolate": True, "cuts": cuts}

    # Initialise standard cross-sections, then specify and set parametrisation models

    cross_sections = pp.crosssection.make_std_crosssection(**args)

    brems_param = pp.parametrization.bremsstrahlung.KelnerKokoulinPetrukhin(lpm=False)
    epair_param = pp.parametrization.pairproduction.KelnerKokoulinPetrukhin(lpm=False)
    ionis_param = pp.parametrization.ionization.BetheBlochRossi(energy_cuts=cuts)
    shado_param = pp.parametrization.photonuclear.ShadowButkevichMikheyev()
    photo_param = pp.parametrization.photonuclear.AbramowiczLevinLevyMaor97(
        shadow_effect=shado_param
    )

    cross_sections[0] = pp.crosssection.make_crosssection(brems_param, **args)
    cross_sections[1] = pp.crosssection.make_crosssection(epair_param, **args)
    cross_sections[2] = pp.crosssection.make_crosssection(ionis_param, **args)
    cross_sections[3] = pp.crosssection.make_crosssection(photo_param, **args)

    # Propagation utility

    collection = pp.PropagationUtilityCollection()

    collection.interaction = pp.make_interaction(cross_sections, True)
    collection.displacement = pp.make_displacement(cross_sections, True)
    collection.time = pp.make_time(cross_sections, mu, True)
    collection.decay = pp.make_decay(cross_sections, mu, True)

    pp.PropagationUtilityCollection.cont_rand = False

    utility = pp.PropagationUtility(collection=collection)

    # Other settings

    pp.do_exact_time = False

    # Set up geometry

    detector = pp.geometry.Sphere(
        position=pp.Cartesian3D(0, 0, 0), radius=10000000, inner_radius=0
    )
    density_distr = pp.density_distribution.density_homogeneous(
        mass_density=constants.get_density()
    )

    propagator = pp.Propagator(mu, [(detector, utility, density_distr)])

    if constants.get_verbose() > 1:
        print("Finished creating propagator.")

    return propagator

# Propagation function

def _propagation_loop(energy, slant_depth, force):

    # This function propagates n_muon muons, looping over the energies and slant depths, and returns the muons' underground energies

    # Check values

    n_muon = constants.get_n_muon()

    constants.check_constants(force=force)

    # Convert the slant depth from [km.w.e.] to [cm]

    convert_to_cm = 1e5 * 0.997 / constants.get_density()

    # Initialise the list of underground energies

    u_energies_ix = []

    # Define the initial state of the muon

    mu_initial = pp.particle.ParticleState()
    mu_initial.energy = energy + constants.MU_MASS
    mu_initial.position = pp.Cartesian3D(0, 0, 0)
    mu_initial.direction = pp.Cartesian3D(0, 0, -1)

    # Propagate n_muon muons

    for _ in range(n_muon):

        # Propagate the muons

        track = propagator.propagate(mu_initial, slant_depth * convert_to_cm)

        # Test whether or not the muon has energy left (has not lost all of its energy or has not decayed)
        # If it does, record its energy
        # If it does not, ignore this muon and proceed with the next loop iteration

        if (
            track.track_energies()[-1] != constants.MU_MASS
            and track.track_types()[-1] != pp.particle.Interaction_Type.decay
        ):

            # Store the final underground energy of the muon

            u_energies_ix.append(track.track_energies()[-1])

    # Return the underground energies for the muon

    return u_energies_ix

# Propagate the muons and return underground energies

def propagate_muons(seed=0, job_array_number=0, output=None, force=False):

    """
    Propagate muons for the default surface energy grid and slant depths.

    The default surface energy grid is given by constants.ENERGIES, and the default slant depths are given by constants.SLANT_DEPTHS.

    Parameters
    ----------
    seed : str, optional (default: 0)
        The random seed for use in the PROPOSAL propagator.

    job_array_number : int, optional (default: 0)
        The job array number from a high-statistics run on a computer cluster. This is set so the underground energy files from each job in the job array will be named differently.

    output : bool, optional (default: taken from constants.get_output())
        If True, an output file will be created to store the results.

    force : bool, optional (default: False)
        If True, this will force the creation of an underground_energies directory if one does not already exist.

    Returns
    -------
    u_energies : NumPy ndarray
        A two-dimensional array containing lists of underground energies for muons that survived the propagation.
    """

    # Check values

    assert type(job_array_number) == int, "job_array_number must be an integer."

    constants.check_constants(force=force)

    if output is None:
        output = constants.get_output()

    # Create the propagator once

    _create_propagator(force=force)

    # Set the random seed

    pp.RandomGenerator.get().set_seed(seed)

    # Initialise the matrix of underground energies

    u_energies = np.zeros(
        (len(constants.ENERGIES), len(constants.SLANT_DEPTHS)), dtype=np.ndarray
    )

    # Run the propagation function and print the underground energies

    if constants.get_verbose() >= 1:
        print(
            "Propagating "
            + str(
                constants.get_n_muon()
                * len(constants.ENERGIES)
                * len(constants.SLANT_DEPTHS)
            )
            + " muons."
        )

    for i in (
        tqdm(range(len(constants.ENERGIES)))
        if constants.get_verbose() >= 1
        else range(len(constants.ENERGIES))
    ):

        for x in range(len(constants.SLANT_DEPTHS)):

            u_energies[i, x] = _propagation_loop(
                constants.ENERGIES[i], constants.SLANT_DEPTHS[x], force=force
            )

    if constants.get_verbose() >= 1:
        print("Finished propagation.")

    if output:

        constants.check_directory(
            os.path.join(constants.get_directory(), "underground_energies"), force=force
        )

        file_name = os.path.join(
            constants.get_directory(),
            "underground_energies",
            "{0}_{1}_{2}_Underground_Energies_{3}.npy".format(
                constants.get_medium(),
                constants.get_density(),
                constants.get_n_muon(),
                job_array_number,
            ),
        )

        np.save(file_name, u_energies)

        if constants.get_verbose() > 1:
            print("Underground energies written to " + file_name + ".")

    return u_energies

# Load underground energies

def _load_u_energies_from_files(file_name, n_job=1, force=False):

    """
    Load the underground energies resulting from the PROPOSAL Monte Carlo from a file or collection of files stored in data/underground_energies.

    Parameters
    ----------
    file_name : str, optional
        The file name pattern that the underground energy data is stored in. This should end in an underscore so the function can append the job array number.

    n_job : int, optional (default: 1)
        The number of jobs that were run on the computer cluster. Set this to the number of files the underground energies are spread across.

    force : bool, optional (default: False)
        If True, this will force the creation of an underground_energies directory if one does not already exist.

    Returns
    -------
    u_energies : NumPy ndarray
        A two-dimensional array containing lists of underground energies for muons that survived the propagation.
    """

    # Check values

    constants.check_constants(force=force)

    # Check that the directory exists

    if not os.path.exists(constants.get_directory() + "/underground_energies"):

        if constants.get_verbose() >= 1:

            print(
                constants.get_directory()
                + "/underground_energies does not exist. Underground energies not loaded."
            )

        return None

    # Test if the file exists

    if not os.path.isfile(file_name + "_0.npy"):

        if constants.get_verbose() >= 1:

            print(file_name + "_0.npy does not exist. Underground energies not loaded.")

        return None

    # Fill a u_energies array with empty lists that will be able to be extended

    u_energies = np.empty(
        (len(constants.ENERGIES), len(constants.SLANT_DEPTHS)), dtype=object
    )

    for i in np.ndindex(u_energies.shape):
        u_energies[i] = []

    # Loop over all output files and add the contents to u_energies

    for a in tqdm(range(n_job)) if constants.get_verbose() >= 1 else range(n_job):

        u_energies += np.load(file_name + "_" + str(a) + ".npy", allow_pickle=True)

    if constants.get_verbose() > 1:
        print("Loaded underground energies.")

    return u_energies

def calc_survival_probability_tensor(
    seed=0, file_name=None, n_job=1, output=None, force=False
):

    """
    Calculate survival probabilities for the default surface energy grid and slant depths.

    The default surface energy grid is given by constants.ENERGIES, and the default slant depths are given by constants.SLANT_DEPTHS. If the propagation of muons has already been done, this will load the underground energies file (it will load it for n_job = 1; to load for more jobs, call load_u_energies_from_files() directly), unless force is set to True.

    Parameters
    ----------
    seed : int, optional (default: 0)
        The random seed for use in the PROPOSAL propagator.

    file_name : str, optional
        The file name pattern that the underground energy data is stored in. This should end in an underscore so the function can append the job array number.

    n_job : int, optional (default: 1)
        The number of jobs that were run on the computer cluster. Set this to the number of files the underground energies are spread across.

    output : bool, optional (default: taken from constants.get_output())
        If True, an output file will be created to store the results.

    force : bool, optional (default: False)
        If True, this will force the muons to be propagated whether an underground energies file already exists or not.

    Returns
    -------
    survival : NumPy ndarray
        A three-dimensional array containing the survival probabilities.
    """

    # Check values

    if output is None:
        output = constants.get_output()

    # Construct a file name

    file_name_default = os.path.join(
        constants.get_directory(),
        "underground_energies",
        "{0}_{1}_{2}_Underground_Energies".format(
            constants.get_medium(),
            constants.get_density(),
            int(constants.get_n_muon() / n_job),
        ),
    )

    # Check if propagate_muons() should be forced or not
    # If not, check if underground energy files exist
    # If not, ask if muons should be propagated

    if force:

        u_energies = propagate_muons(seed=seed, output=output, force=force)

    else:

        # Check if the user has specified underground energies to load
        # If not, look for the default file name pattern, and check if it exists
        # If so, load the underground energies
        # If not, ask if muons should be propagated

        if file_name is not None:

            u_energies = _load_u_energies_from_files(
                file_name=os.path.join(
                    constants.get_directory(), "underground_energies", file_name
                ),
                n_job=n_job,
                force=force,
            )

        elif os.path.isfile(file_name_default + "_0.npy"):

            u_energies = _load_u_energies_from_files(
                file_name=file_name_default, n_job=n_job, force=force
            )

        else:

            answer = input(
                "No underground energy file currently exists for the set lab, medium, or number of muons. Would you like to create one (y/n)?: "
            )

            if answer.lower() == "y":

                u_energies = propagate_muons(seed=seed, output=output, force=force)

            else:

                print("Underground energies not calculated.")
                print("Survival probabilities not calculated.")

                return

    # Check that the underground energies were loaded

    if u_energies is None:

        print("Survival probabilities not calculated.")

        return

    # Calculate the survival probabilities
    # First index  = Surface energy
    # Second index = Slant depth
    # Third index  = Underground energy

    survival = np.zeros(
        (len(constants.ENERGIES), len(constants.SLANT_DEPTHS), len(constants.ENERGIES))
    )

    if constants.get_verbose() > 1:
        print("Calculating survival probabilities.")

    for i in range(len(constants.ENERGIES)):

        for x in range(len(constants.SLANT_DEPTHS)):

            survival[i, x, :] = np.histogram(
                np.array(u_energies[i, x]), bins=constants.E_BINS
            )[0] / float(constants.get_n_muon())

    if constants.get_verbose() > 1:
        print("Finished calculating survival probabilities.")

    # Write the results to a file

    if output:

        constants.check_directory(
            os.path.join(constants.get_directory(), "survival_probabilities"),
            force=force,
        )

        file_name = os.path.join(
            constants.get_directory(),
            "survival_probabilities",
            "{0}_{1}_{2}_Survival_Probabilities.txt".format(
                constants.get_medium(), constants.get_density(), constants.get_n_muon()
            ),
        )

        file_out = open(file_name, "w")

        for i in range(len(constants.ENERGIES)):

            for x in range(len(constants.SLANT_DEPTHS)):

                for u in range(len(constants.ENERGIES)):

                    file_out.write(
                        "{0:1.14f} {1:1.5f} {2:1.14f} {3:1.14e}\n".format(
                            constants.ENERGIES[i],
                            constants.SLANT_DEPTHS[x],
                            constants.ENERGIES[u],
                            survival[i, x, u],
                        )
                    )

        file_out.close()

        if constants.get_verbose() > 1:
            print("Survival probabilities written to " + file_name + ".")

    return survival

def load_survival_probability_tensor_from_file(force=False):

    """
    Retrieve a survival probability matrix stored in data/survival_probabilities based on the set global parameters.

    The function searches for a file name that matches the set lab, medium, and number of muons. If the file does not exist, prompt the user to run calc_survival().

    Parameters
    ----------
    force : bool
        If True, force the calculation of a new survival probability tensor if required.

    Returns
    -------
    survival : NumPy ndarray
        A two-dimensional array containing the survival probabilities.
    """

    # Define a function to run if there is no survival probability file

    def no_file(force):

        # If the file does not exist, ask the user if they want to run PROPOSAL to create it

        if not force:

            answer = input(
                "No survival probability matrix currently exists for the set lab, medium, or number of muons. Would you like to create one (y/n)?: "
            )

        if force or answer.lower() == "y":

            survival_full = calc_survival_probability_tensor(force=force)

            return survival_full

        else:

            print("Survival probabilities not calculated.")

            return None

    # Construct a file name based on the set lab, medium, and number of muons

    file_name = os.path.join(
        constants.get_directory(),
        "survival_probabilities",
        "{0}_{1}_{2}_Survival_Probabilities.txt".format(
            constants.get_medium(), constants.get_density(), constants.get_n_muon()
        ),
    )

    # Check if the file exists
    # Also check that the file has the correct number of energies and angles

    if os.path.isfile(file_name):

        if constants.get_verbose() > 1:
            print("Loading survival probabilities from " + file_name + ".")

        # If the file exists, read in the survival probabilities from it

        file = open(file_name, "r")
        n_lines = len(file.read().splitlines())

        file.close()

        if n_lines == constants.len_iju:

            survival = np.reshape(
                np.loadtxt(file_name)[:, 3],
                (
                    len(constants.ENERGIES),
                    len(constants.SLANT_DEPTHS),
                    len(constants.ENERGIES),
                ),
            )

            if constants.get_verbose() > 1:
                print("Loaded survival probabilities.")

            return survival

        else:

            return no_file(force=force)

    else:

        return no_file(force=force)

# Read the energies and slant depths from a survival probabilities file

def print_survival_probability_tensor_grids(file_name):

    """Return the slant depths and underground energies in a survival probability file."""

    file_contents = np.loadtxt(
        os.path.join(constants.get_directory(), "survival_probabilities", file_name)
    )

    file_s_energies = np.unique(file_contents[:, 0])
    file_slant_depths = np.unique(file_contents[:, 1])
    file_u_energies = np.unique(file_contents[:, 2])

    print("This file has " + str(len(file_s_energies)) + " surface energies:")
    print(file_s_energies)
    print("This file has " + str(len(file_slant_depths)) + " slant depths:")
    print(file_slant_depths)
    print("This file has " + str(len(file_u_energies)) + " underground energies:")
    print(file_u_energies)

    return


#!/usr/bin/env python
import glob
import os
import os.path
import shutil
import socket
import struct
import sys
import time
import subprocess
import warnings
import re
import types
import tarfile
import tempfile

import numpy as np
import psr_utils
import presto
import sifting

import datafile
import config.searching
import config.processing

# Sifting specific parameters (don't touch without good reason!)
# incoherent power threshold (sigma)
sifting.sigma_threshold = config.searching.sifting_sigma_threshold
# coherent power threshold
sifting.c_pow_threshold = config.searching.sifting_c_pow_threshold
# Fourier bin tolerence for candidate equivalence
sifting.r_err           = config.searching.sifting_r_err
# Shortest period candidates to consider (s)
sifting.short_period    = config.searching.sifting_short_period
# Longest period candidates to consider (s)
sifting.long_period     = config.searching.sifting_long_period
# Power required in at least one harmonic
sifting.harm_pow_cutoff = config.searching.sifting_harm_pow_cutoff

debug = 0

def get_baryv(ra, dec, mjd, T, obs="AO"):
   """
   get_baryv(ra, dec, mjd, T):
     Determine the average barycentric velocity towards 'ra', 'dec'
       during an observation from 'obs'.  The RA and DEC are in the
       standard string format (i.e. 'hh:mm:ss.ssss' and
       'dd:mm:ss.ssss'). 'T' is in sec and 'mjd' is (of course) in MJD.
   """
   tts = psr_utils.span(mjd, mjd+T/86400.0, 100)
   nn = len(tts)
   bts = np.zeros(nn, dtype=np.float64)
   vel = np.zeros(nn, dtype=np.float64)
   presto.barycenter(tts, bts, vel, nn, ra, dec, obs, "DE200")
   avgvel = np.add.reduce(vel)/nn
   return avgvel

def find_masked_fraction(obs):
    """
    find_masked_fraction(obs):
        Parse the output file from an rfifind run and return the
            fraction of the data that was suggested to be masked.
    """
    rfifind_out = obs.basefilenm + "_rfifind.out"
    for line in open(rfifind_out):
        if "Number of  bad   intervals" in line:
            return float(line.split("(")[1].split("%")[0])/100.0
    # If there is a problem reading the file, return 100%
    return 100.0

def get_all_subdms(ddplans):
    """
    get_all_subdms(ddplans):
        Return a sorted array of the subdms from the list of ddplans.
    """
    subdmlist = []
    for ddplan in ddplans:
        subdmlist += [float(x) for x in ddplan.subdmlist]
    subdmlist.sort()
    subdmlist = np.asarray(subdmlist)
    return subdmlist

def find_closest_subbands(obs, subdms, DM):
    """
    find_closest_subbands(obs, subdms, DM):
        Return the basename of the closest set of subbands to DM
        given an obs_info class and a sorted array of the subdms.
    """
    subdm = subdms[np.fabs(subdms - DM).argmin()]
    return "obs.tempdir/%s_DM%.2f.sub[0-6]*"%(obs.basefilenm, subdm)

def timed_execute(cmd, stdout=None, stderr=sys.stderr):
    """
    timed_execute(cmd, stdout=None, stderr=sys.stderr):
        Execute the command 'cmd' after logging the command
            to STDOUT.  Return the wall-clock amount of time
            the command took to execute.

            Output standard output to 'stdout' and standard
            error to 'stderr'. Both are strings containing filenames.
            If values are None, the out/err streams are not recorded.
            By default stdout is None and stderr is combined with stdout.
    """
    # Log command to stdout
    sys.stdout.write("\n'"+cmd+"'\n")
    sys.stdout.flush()

    stdoutfile = False
    stderrfile = False
    if type(stdout) == types.StringType:
        stdout = open(stdout, 'w')
        stdoutfile = True
    if type(stderr) == types.StringType:
        stderr = open(stderr, 'w')
        stderrfile = True

    # Run (and time) the command. Check for errors.
    start = time.time()
    retcode = subprocess.call(cmd, shell=True, stdout=stdout, stderr=stderr)
    if retcode < 0:
        raise PrestoError("Execution of command (%s) terminated by signal (%s)!" % \
                                (cmd, -retcode))
    elif retcode > 0:
        raise PrestoError("Execution of command (%s) failed with status (%s)!" % \
                                (cmd, retcode))
    else:
        # Exit code is 0, which is "Success". Do nothing.
        pass
    end = time.time()

    # Close file objects, if any
    if stdoutfile:
        stdout.close()
    if stderrfile:
        stderr.close()
    return end - start

def get_folding_command(cand, obs):
    """
    get_folding_command(cand, obs):
        Return a command for prepfold for folding the subbands using
            an obs_info instance, and a candidate instance that
            describes the observations and searches.
    """
    # Folding rules are based on the facts that we want:
    #   1.  Between 24 and 200 bins in the profiles
    #   2.  For most candidates, we want to search length = 101 p/pd/DM cubes
    #       (The side of the cube is always 2*M*N+1 where M is the "factor",
    #       either -npfact (for p and pd) or -ndmfact, and N is the number of bins
    #       in the profile).  A search of 101^3 points is pretty fast.
    #   3.  For slow pulsars (where N=100 or 200), since we'll have to search
    #       many points, we'll use fewer intervals in time (-npart 30)
    #   4.  For the slowest pulsars, in order to avoid RFI, we'll
    #       not search in period-derivative.
    zmax = cand.filename.split("_")[-1]
    outfilenm = obs.basefilenm+"_DM%s_Z%s"%(cand.DMstr, zmax)

    # Note:  the following calculations should probably only be done once,
    #        but in general, these calculation are effectively instantaneous
    #        compared to the folding itself
    if config.searching.fold_rawdata:
        # Fold raw data
        foldfiles = obs.filenmstr
        mask = "-mask %s" % (obs.basefilenm + "_rfifind.mask")
    else:
        if config.searching.use_subbands:
            # Fold the subbands
            subdms = get_all_subdms(obs.ddplans)
            subfiles = find_closest_subbands(obs, subdms, cand.DM)
            foldfiles = subfiles
            mask = ""
        else:  # Folding the downsampled PSRFITS files instead
            #
            # TODO: Apply mask!?
            #
            mask = ""
            hidms = [x.lodm for x in obs.ddplans[1:]] + [2000]
            dfacts = [x.downsamp for x in obs.ddplans]
            for hidm, dfact in zip(hidms, dfacts):
                if cand.DM < hidm:
                    downsamp = dfact
                    break
            if downsamp==1:
                foldfiles = obs.filenmstr
            else:
                dsfiles = []
                for f in obs.filenames:
                    fbase = f.rstrip(".fits")
                    dsfiles.append(fbase+"_DS%d.fits"%downsamp)
                foldfiles = ' '.join(dsfiles)
    p = 1.0 / cand.f
    if p < 0.002:
        Mp, Mdm, N = 2, 2, 24
        npart = 50
        otheropts = "-ndmfact 3"
    elif p < 0.05:
        Mp, Mdm, N = 2, 1, 50
        npart = 40
        otheropts = "-pstep 1 -pdstep 2 -dmstep 3"
    elif p < 0.5:
        Mp, Mdm, N = 1, 1, 100
        npart = 30
        otheropts = "-pstep 1 -pdstep 2 -dmstep 1"
    else:
        Mp, Mdm, N = 1, 1, 200
        npart = 30
        otheropts = "-nopdsearch -pstep 1 -pdstep 2 -dmstep 1"

    # If prepfold is instructed to use more subbands than there are rows in the PSRFITS file
    # it doesn't use any data when folding since the amount of data for each part is
    # shorter than the PSRFITS row. However, PRESTO doesn't break up rows.
    # Set npart to the number of rows in the PSRFITS file.
    if npart > obs.numrows:
        npart = obs.numrows

    # Get number of subbands to use
    if obs.backend.lower() == 'pdev':
        nsub = 96
    else:
        nsub = 64
    return "prepfold -noxwin -accelcand %d -accelfile %s.cand -dm %.2f -o %s " \
                "-nsub %d -npart %d %s -n %d -npfact %d -ndmfact %d %s %s" % \
           (cand.candnum, cand.filename, cand.DM, outfilenm, nsub,
            npart, otheropts, N, Mp, Mdm, mask, foldfiles)

class obs_info:
    """
    class obs_info(filenms, resultsdir)
        A class describing the observation and the analysis.
    """
    def __init__(self, filenms, resultsdir):
        # Where to dump all the results
        self.outputdir = resultsdir

        self.filenms = filenms
        self.filenmstr = ' '.join(self.filenms)
        self.basefilenm = os.path.split(filenms[0])[1].rstrip(".fits")
        print "Read info from PSRFITS file", self.filenms
        # Read info from PSRFITS file
        data = datafile.autogen_dataobj(self.filenms)
        # Correct positions in data file headers for WappPsrfitsData
        print "Updating positions"
        if isinstance(data, datafile.WappPsrfitsData):
            data.update_positions()
        print "Positions updated"
        spec_info = data.specinfo
        self.backend = spec_info.backend
        self.MJD = spec_info.start_MJD[0]
        self.ra_string = spec_info.ra_str
        self.dec_string = spec_info.dec_str
        self.orig_N = spec_info.N
        self.dt = spec_info.dt # in sec
        self.BW = spec_info.BW
        self.orig_T = spec_info.T
        # Downsampling is catered to the number of samples per row.
        # self.N = psr_utils.choose_N(self.orig_N)
        self.N = self.orig_N
        self.T = self.N * self.dt
        self.nchan = spec_info.num_channels
        self.samp_per_row = spec_info.spectra_per_subint
        self.fctr = spec_info.fctr
        self.numrows = np.sum(spec_info.num_subint)

        # Determine the average barycentric velocity of the observation
        print "Determine the average barycentric velocity of the observation"
#        self.baryv = get_baryv(self.ra_string, self.dec_string,
#                               self.MJD, self.T, obs="AO")
        self.baryv = 0.0
        # Figure out which host we are processing on
        print "Figure out which host we are processing on"
        self.hostname = socket.gethostname()
        # The fraction of the data recommended to be masked by rfifind
        self.masked_fraction = 0.0
        # The number of candidates folded
        self.num_cands_folded = 0
        # Initialize our timers
        self.rfifind_time = 0.0
        self.downsample_time = 0.0
        self.subbanding_time = 0.0
        self.dedispersing_time = 0.0
        self.FFT_time = 0.0
        self.lo_accelsearch_time = 0.0
        self.hi_accelsearch_time = 0.0
        self.singlepulse_time = 0.0
        self.modindex_time = 0.0    #Added by LGS 30Mar12
        self.sifting_time = 0.0
        self.folding_time = 0.0
        self.total_time = 0.0
        # Inialize some candidate counters
        self.num_sifted_cands = 0
        self.num_folded_cands = 0
        self.num_single_cands = 0
        # Set dedispersion plan
        print "Set dedispersion plan"
        self.set_DDplan()

    def set_DDplan(self):
        """Set the dedispersion plan.

            The dedispersion plans are hardcoded and
            depend on the backend data were recorded with.
        """
        # Generate dedispersion plan
        self.ddplans = []

        # The following code will run the dedispersion planner on demand.
        # Instead, dedispersion plans for WAPP and Mock data are hardcoded.
        #
        # import DDplan2b
        # obs = DDplan2b.Observation(self.dt, self.fctr, self.BW, self.nchan, \
        #                             self.samp_per_row)
        # plan = obs.gen_ddplan(config.searching.lodm, config.searching.hidm, \
        #                       config.searching.numsub, config.searching.resolution)
        # plan.plot(fn=os.path.join(self.outputdir, self.basefilenm+"_ddplan.ps"))
        # print plan
        # for ddstep in plan.DDsteps:
        #     self.ddplans.append(dedisp_plan(ddstep.loDM, ddstep.dDM, ddstep.DMs_per_prepsub, \
        #                    ddstep.numprepsub, ddstep.numsub, ddstep.downsamp))

        if self.backend.lower() == 'pdev':
            # The values here are:       lodm dmstep dms/call #calls #subbands downsamp
            # This line added for debugging. Searches fewer DM so it goes faster
            self.ddplans.append(dedisp_plan(   98.8,  0.1,    76,      1,     96,        1 ))
#            self.ddplans.append(dedisp_plan(   0.0,  0.1,    76,     28,     96,        1 ))
#            self.ddplans.append(dedisp_plan( 212.8,  0.3,    64,     12,     96,        2 ))
#            self.ddplans.append(dedisp_plan( 443.2,  0.3,    76,      4,     96,        3 ))
#            self.ddplans.append(dedisp_plan( 534.4,  0.5,    76,      9,     96,        5 ))
#            self.ddplans.append(dedisp_plan( 876.4,  0.5,    76,      3,     96,        6 ))
#            self.ddplans.append(dedisp_plan( 990.4,  1.0,    76,      10,     96,       10 ))
#            self.ddplans.append(dedisp_plan(1750.4,  2.0,    72,      2,     96,       15 ))
        elif self.backend.lower() == 'wapp':
            # The values here are:       lodm dmstep dms/call #calls #subbands downsamp
            self.ddplans.append(dedisp_plan(   0.0,  0.3,    76,      9,     96,        1 ))
            self.ddplans.append(dedisp_plan( 205.2,  2.0,    76,      5,     96,        5 ))
            self.ddplans.append(dedisp_plan( 965.2, 10.0,    76,      1,     96,       25 ))
        else:
            raise ValueError("No dediserpsion plan for unknown backend (%s)!" % self.backend)

    def write_report(self, filenm):
        report_file = open(filenm, "w")
        report_file.write("---------------------------------------------------------\n")
        report_file.write("Data (%s) were processed on %s\n" % \
                                (', '.join(self.filenms), self.hostname))
        report_file.write("Ending UTC time:  %s\n"%(time.asctime(time.gmtime())))
        report_file.write("Total wall time:  %.1f s (%.2f hrs)\n"%\
                          (self.total_time, self.total_time/3600.0))
        report_file.write("Fraction of data masked:  %.2f%%\n"%\
                          (self.masked_fraction*100.0))
        report_file.write("Number of candidates folded: %d\n"%\
                          self.num_cands_folded)
        report_file.write("---------------------------------------------------------\n")
        report_file.write("          rfifind time = %7.1f sec (%5.2f%%)\n"%\
                          (self.rfifind_time, self.rfifind_time/self.total_time*100.0))
        if config.searching.use_subbands:
            report_file.write("       subbanding time = %7.1f sec (%5.2f%%)\n"%\
                              (self.subbanding_time, self.subbanding_time/self.total_time*100.0))
        else:
            report_file.write("     downsampling time = %7.1f sec (%5.2f%%)\n"%\
                              (self.downsample_time, self.downsample_time/self.total_time*100.0))
        report_file.write("     dedispersing time = %7.1f sec (%5.2f%%)\n"%\
                          (self.dedispersing_time, self.dedispersing_time/self.total_time*100.0))
        report_file.write("     single-pulse time = %7.1f sec (%5.2f%%)\n"%\
                          (self.singlepulse_time, self.singlepulse_time/self.total_time*100.0))
        report_file.write(" modulation index time = %7.1f sec (%5.2f%%)\n"%\
                          (self.modindex_time, self.modindex_time/self.total_time*100.0))
        report_file.write("              FFT time = %7.1f sec (%5.2f%%)\n"%\
                          (self.FFT_time, self.FFT_time/self.total_time*100.0))
        report_file.write("   lo-accelsearch time = %7.1f sec (%5.2f%%)\n"%\
                          (self.lo_accelsearch_time, self.lo_accelsearch_time/self.total_time*100.0))
        report_file.write("   hi-accelsearch time = %7.1f sec (%5.2f%%)\n"%\
                          (self.hi_accelsearch_time, self.hi_accelsearch_time/self.total_time*100.0))
        report_file.write("          sifting time = %7.1f sec (%5.2f%%)\n"%\
                          (self.sifting_time, self.sifting_time/self.total_time*100.0))
        report_file.write("          folding time = %7.1f sec (%5.2f%%)\n"%\
                          (self.folding_time, self.folding_time/self.total_time*100.0))
        report_file.write("---------------------------------------------------------\n")
        report_file.close()

class dedisp_plan:
    """
    class dedisp_plan(lodm, dmstep, dmsperpass, numpasses, numsub, downsamp)
        A class describing a de-dispersion plan for prepsubband in detail.
    """
    def __init__(self, lodm, dmstep, dmsperpass, numpasses, numsub, downsamp):
        self.lodm = float(lodm)
        self.dmstep = float(dmstep)
        self.dmsperpass = int(dmsperpass)
        self.numpasses = int(numpasses)
        self.numsub = int(numsub)
        self.downsamp = int(downsamp)
        # Downsample less for the subbands so that folding
        # candidates is more acurate
        #
        # Turning this off because downsampling factors are not necessarily
        # powers of 2 any more! Also, are we folding from raw data now?
        # -- PL Nov. 26, 2010
        #
        self.sub_downsamp = self.downsamp
        self.dd_downsamp = 1
        # self.sub_downsamp = self.downsamp / 2
        # if self.sub_downsamp==0: self.sub_downsamp = 1
        # The total downsampling is:
        #   self.downsamp = self.sub_downsamp * self.dd_downsamp

        # if self.downsamp==1: self.dd_downsamp = 1
        # else: self.dd_downsamp = 2
        self.sub_dmstep = self.dmsperpass * self.dmstep
        self.dmlist = []  # These are strings for comparison with filenames
        self.subdmlist = []
        for ii in range(self.numpasses):
            self.subdmlist.append("%.2f"%(self.lodm + (ii+0.5)*self.sub_dmstep))
            lodm = self.lodm + ii * self.sub_dmstep
            dmlist = ["%.2f"%dm for dm in \
                      np.arange(self.dmsperpass)*self.dmstep + lodm]
            self.dmlist.append(dmlist)

def main(filenms, workdir, resultsdir):

    # Change to the specified working directory
    os.chdir(workdir)
    print workdir
    job = set_up_job(filenms, workdir, resultsdir)

    print "\nBeginning PALFA search of %s" % (', '.join(job.filenms))
    print "UTC time is:  %s"%(time.asctime(time.gmtime()))
    print "NOTE: This version of the pipeline has been modified by"
    print "   Laura Spitler to do only the single pulse search and"
    print "   also includes the calculation of the modulation index."
    try:
        search_job(job)
    except:
        print "***********************ERRORS!************************"
        print "  Search has been aborted due to errors encountered."
        print "  See error output for more information."
        print "******************************************************"
        raise
    finally:
        clean_up(job)

        # And finish up
        job.total_time = time.time() - job.total_time
        print "\nFinished"
        print "UTC time is:  %s"%(time.asctime(time.gmtime()))

        # Write the job report
        # job.write_report(job.basefilenm+".report")
        job.write_report(os.path.join(job.outputdir, job.basefilenm+".report"))

def set_up_job(filenms, workdir, resultsdir):
    """Change to the working directory and set it up.
        Create a obs_info instance, set it up and return it.
    """
    # Get information on the observation and the job
    print filenms
    job = obs_info(filenms, resultsdir)
    if job.T < config.searching.low_T_to_search:
        raise PrestoError("The observation is too short to search. " \
                            "(%.2f s < %.2f s)" % \
                            (job.T, config.searching.low_T_to_search))
    job.total_time = time.time()

    print "Inside set_up_job for files: ", job.filenms
    # Make sure the output directory (and parent directories) exist
    try:
        os.makedirs(job.outputdir)
    except: pass

    job.workdir = workdir
    # Create a directory to hold all the subbands
    job.tempdir = tempfile.mkdtemp(suffix="_tmp", prefix=job.basefilenm, \
                        dir=config.processing.base_tmp_dir)

    #####
    # Print some info useful for debugging
    print "Initial contents of workdir (%s): " % workdir
    for fn in os.listdir(workdir):
        print "    %s" % fn
    print "Initial contents of resultsdir (%s): " % resultsdir
    for fn in os.listdir(resultsdir):
        print "    %s" % fn
    print "Initial contents of job.tempdir (%s): " % job.tempdir
    for fn in os.listdir(job.tempdir):
        print "    %s" % fn
    sys.stdout.flush()
    #####

    return job

def search_job(job):
    """Search the observation defined in the obs_info
        instance 'job'.
    """
    # Use whatever .zaplist is found in the current directory
    zaplist = glob.glob("*.zaplist")[0]
    print "Using %s as zaplist" % zaplist
    if config.searching.use_subbands and config.searching.fold_rawdata:
        # make a directory to keep subbands so they can be used to fold later
        try:
            os.makedirs(os.path.join(job.workdir, 'subbands'))
        except: pass

    # rfifind the data file
    cmd = "rfifind %s -time %.17g -o %s %s" % \
          (config.searching.datatype_flag, config.searching.rfifind_chunk_time, job.basefilenm,
           job.filenmstr)
    job.rfifind_time += timed_execute(cmd, stdout="%s_rfifind.out" % job.basefilenm)
    maskfilenm = job.basefilenm + "_rfifind.mask"
    # Find the fraction that was suggested to be masked
    # Note:  Should we stop processing if the fraction is
    #        above some large value?  Maybe 30%?
    job.masked_fraction = find_masked_fraction(job)

    # Iterate over the stages of the overall de-dispersion plan
    dmstrs = []
    for ddplan in job.ddplans:

        # Iterate over the individual passes through the data file
        for passnum in range(ddplan.numpasses):
            subbasenm = "%s_DM%s"%(job.basefilenm, ddplan.subdmlist[passnum])

            if config.searching.use_subbands:
                try:
                    os.makedirs(os.path.join(job.tempdir, 'subbands'))
                except: pass

                # Create a set of subbands
		#LGS: -nobary flag added on 19Jan12
                cmd = "prepsubband %s -sub -noweights -noscales -nooffsets -subdm %s -downsamp %d -nsub %d -mask %s " \
                        "-o %s/subbands/%s %s" % \
                        (config.searching.datatype_flag, ddplan.subdmlist[passnum], ddplan.sub_downsamp,
                        ddplan.numsub, maskfilenm, job.tempdir, job.basefilenm,
                        job.filenmstr)
                job.subbanding_time += timed_execute(cmd, stdout="%s.subout" % subbasenm)

                # Now de-disperse using the subbands
		#LGS: -nobary flag added on 19Jan12
                cmd = "prepsubband -noweights -noscales -nooffsets -lodm %.2f -dmstep %.2f -numdms %d -downsamp %d " \
                        "-nsub %d -numout %d -o %s/%s %s/subbands/%s.sub[0-9]*" % \
                        (ddplan.lodm+passnum*ddplan.sub_dmstep, ddplan.dmstep,
                        ddplan.dmsperpass, ddplan.dd_downsamp, ddplan.numsub,
                        psr_utils.choose_N(job.orig_N/ddplan.downsamp),
                        job.tempdir, job.basefilenm, job.tempdir, subbasenm)
                job.dedispersing_time += timed_execute(cmd, stdout="%s.prepout" % subbasenm)

            else:  # Not using subbands
                cmd = "prepsubband -mask %s -lodm %.2f -dmstep %.2f -numdms %d -downsamp %d " \
                        "-numout %d -o %s/%s %s"%\
                        (maskfilenm, ddplan.lodm+passnum*ddplan.sub_dmstep, ddplan.dmstep,
                        ddplan.dmsperpass, ddplan.dd_downsamp*ddplan.sub_downsamp,
                        psr_utils.choose_N(job.orig_N/ddplan.downsamp),
                        job.tempdir, job.basefilenm, job.filenmstr)
                job.dedispersing_time += timed_execute(cmd)

            # Iterate over all the new DMs
            for dmstr in ddplan.dmlist[passnum]:
                dmstrs.append(dmstr)
                basenm = os.path.join(job.tempdir, job.basefilenm+"_DM"+dmstr)
                datnm = basenm+".dat"
                fftnm = basenm+".fft"
                infnm = basenm+".inf"

                # Do the single-pulse search
		#LGS: Added processing speed option 19Jan12
                #LGS: Added cluster search option 09Feb12
                cmd = "single_pulse_search.py %s %s %s -p -m %f -t %f %s"%\
                      (config.searching.singlepulse_flag,\
                       config.searching.singlepulse_cluster,\
                       config.searching.singlepulse_speed,\
                       config.searching.singlepulse_maxwidth, \
                       config.searching.singlepulse_threshold, datnm)
                job.singlepulse_time += timed_execute(cmd)
                try:
                    shutil.move(basenm+".singlepulse", job.workdir)
                    if config.searching.singlepulse_cluster: shutil.move(basenm+".cluster", job.workdir)
                except: pass

                # Move the .inf files
                try:
                    shutil.move(infnm, job.workdir)
                except: pass
                # Remove the .dat and .fft files
                try:
                    os.remove(datnm)
                except: pass

            if config.searching.use_subbands:
                if config.searching.fold_rawdata:
                    # Subband files are no longer needed
                    shutil.rmtree(os.path.join(job.tempdir, 'subbands'))
                else:
                    # Move subbands to workdir
                    for sub in glob.glob(os.path.join(job.tempdir, 'subbands', "*")):
                        shutil.move(sub, os.path.join(job.workdir, 'subbands'))

    #Create the masked version of the tf-data
    cmd = "maskdata -nobary -noweights -nooffsets -noscales -mask %s -o junk.dat %s" % (maskfilenm, job.filenmstr)
    job.modindex_time += timed_execute(cmd)

    #Consolidate all singlepulse files into a single file
    cmd="awk '$1!=\"#\" {print}' %s > %s" % (job.basefilenm+'_DM*.singlepulse', job.basefilenm+'_MF.sp')
    job.modindex_time += timed_execute(cmd)

    if config.searching.singlepulse_cluster:
        cmd="awk '$1!=\"#\" {print $1, $9, $12, $13, $14, $6, $7, $8}' %s > %s" % (job.basefilenm+'_DM*.cluster', job.basefilenm+'_CL.sp')
        job.modindex_time += timed_execute(cmd)

    #Calculate mod index
    cmd="palfa_mi %s raw_data_with_mask.fits %s" % (job.basefilenm+'_MF.sp', job.basefilenm+'_MF.mi')
    job.modindex_time += timed_execute(cmd)

    if config.searching.singlepulse_cluster:
        cmd="palfa_mi %s raw_data_with_mask.fits %s" % (job.basefilenm+'_CL.sp', job.basefilenm+'_CL.mi')
        job.modindex_time += timed_execute(cmd)

    # Make the single-pulse plots
    basedmb = job.basefilenm+"_DM"
    basedme = ".singlepulse "
    # The following will make plots for DM ranges:
    #    0-110, 100-310, 300-1000+
    dmglobs = [basedmb+"[0-9].[0-9][0-9]"+basedme +
               basedmb+"[0-9][0-9].[0-9][0-9]"+basedme +
               basedmb+"10[0-9].[0-9][0-9]"+basedme,
               basedmb+"[12][0-9][0-9].[0-9][0-9]"+basedme +
               basedmb+"30[0-9].[0-9][0-9]"+basedme,
               basedmb+"[3-9][0-9][0-9].[0-9][0-9]"+basedme +
               basedmb+"1[0-9][0-9][0-9].[0-9][0-9]"+basedme]
    dmrangestrs = ["0-110", "100-310", "300-1000+"]
    psname = job.basefilenm+"_singlepulse.ps"
    for dmglob, dmrangestr in zip(dmglobs, dmrangestrs):
        dmfiles = []
        for dmg in dmglob.split():
            dmfiles += glob.glob(dmg.strip())
        # Check that there are matching files and they are not all empty
        if dmfiles and sum([os.path.getsize(f) for f in dmfiles]):
            cmd = 'single_pulse_search.py -t %f -g "%s"' % \
                (config.searching.singlepulse_plot_SNR, dmglob)
            job.singlepulse_time += timed_execute(cmd)
            os.rename(psname,
                        job.basefilenm+"_DMs%s_singlepulse.ps" % dmrangestr)

    # Make the cluster plots
    if config.searching.singlepulse_cluster:
        basedmb = job.basefilenm+"_DM"
        basedme = ".cluster "
        # The following will make plots for DM ranges:
        #    0-110, 100-310, 300-1000+
        dmglobs = [basedmb+"[0-9].[0-9][0-9]"+basedme +
               basedmb+"[0-9][0-9].[0-9][0-9]"+basedme +
               basedmb+"10[0-9].[0-9][0-9]"+basedme,
               basedmb+"[12][0-9][0-9].[0-9][0-9]"+basedme +
               basedmb+"30[0-9].[0-9][0-9]"+basedme,
               basedmb+"[3-9][0-9][0-9].[0-9][0-9]"+basedme +
               basedmb+"1[0-9][0-9][0-9].[0-9][0-9]"+basedme]
        dmrangestrs = ["0-110", "100-310", "300-1000+"]
        psname = job.basefilenm+"_cluster.ps"
        for dmglob, dmrangestr in zip(dmglobs, dmrangestrs):
            dmfiles = []
            for dmg in dmglob.split():
                dmfiles += glob.glob(dmg.strip())
            # Check that there are matching files and they are not all empty
            if dmfiles and sum([os.path.getsize(f) for f in dmfiles]):
                cmd = 'single_pulse_search.py --clust -t %f -g "%s"' % \
                    (config.searching.singlepulse_plot_SNR, dmglob)
                job.singlepulse_time += timed_execute(cmd)
                os.rename(psname,
                        job.basefilenm+"_DMs%s_cluster.ps" % dmrangestr)

    # Now step through the .ps files and convert them to .png and gzip them
    psfiles = glob.glob("*.ps")
    for psfile in psfiles:
        # The '[0]' appeneded to the end of psfile is to convert only the 1st page
        timed_execute("convert -quality 90 %s -background white -flatten -rotate 90 +matte %s" % \
                            (psfile+"[0]", psfile[:-3]+".png"))
        timed_execute("gzip "+psfile)

def clean_up(job):
    """Clean up.
        Tar results, copy them to the results director.
    """
    # Dump search paramters to file
    paramfn = open("search_params.txt", 'w')
    cfgs = config.searching_check.searching.configs
    for key in cfgs:
        paramfn.write("%-25s = %r\n" % (key, cfgs[key].value))
    paramfn.close()

    # Tar up the results files
    tar_suffixes = ["_singlepulse.tgz",
                    "_inf.tgz"]
    tar_globs = ["*.singlepulse",
                 "*_DM[0-9]*.inf"]
    if config.searching.singlepulse_cluster == '--clust':
        tar_suffixes.append("_cluster.tgz")
        tar_globs.append("*.cluster")
    print "Tarring up results"
    for (tar_suffix, tar_glob) in zip(tar_suffixes, tar_globs):
        print "Opening tarball %s" % (job.basefilenm+tar_suffix)
        print "Using glob %s" % tar_glob
        tf = tarfile.open(job.basefilenm+tar_suffix, "w:gz")
        for infile in glob.glob(tar_glob):
            print "    Adding file %s" % infile
            tf.add(infile)
            os.remove(infile)
        tf.close()
    sys.stdout.flush()

    # Copy all the important stuff to the output directory
    resultglobs = ["*rfifind.[bimors]*", "*.ps.gz", "*.tgz", "*.png", \
                    "*.zaplist", "search_params.txt", \
                    "*_merge.out", "*.mi", "*.sp"]

    for resultglob in resultglobs:
            for file in glob.glob(resultglob):
                shutil.move(file, job.outputdir)

    # Remove the tmp directory (in a tmpfs mount)
    try:
        shutil.rmtree(job.tempdir)
    except: pass

class PrestoError(Exception):
    """Error to throw when a PRESTO program returns with
        a non-zero error code.
    """
    pass

if __name__ == "__main__":
    # Arguments to the search program are
    # sys.argv[3:] = data file names
    # sys.argv[1] = working directory name
    # sys.argv[2] = results directory name
    workdir = sys.argv[1]
    resultsdir = sys.argv[2]
    filenms = sys.argv[3:]
    main(filenms, workdir, resultsdir)


"Implements Model"
import numpy as np
from .costed import CostedConstraintSet
from ..nomials import Monomial
from .prog_factories import _progify_fctry, _solve_fctry
from .gp import GeometricProgram
from .sgp import SequentialGeometricProgram
from ..small_scripts import mag
from ..tools.autosweep import autosweep_1d
from ..exceptions import InvalidGPConstraint
from .. import NamedVariables
from ..tools.docstring import expected_unbounded
from .set import add_meq_bounds

class Model(CostedConstraintSet):
    """Symbolic representation of an optimization problem.

    The Model class is used both directly to create models with constants and
    sweeps, and indirectly inherited to create custom model classes.

    Arguments
    ---------
    cost : Posynomial (optional)
        Defaults to `Monomial(1)`.

    constraints : ConstraintSet or list of constraints (optional)
        Defaults to an empty list.

    substitutions : dict (optional)
        This dictionary will be substituted into the problem before solving,
        and also allows the declaration of sweeps and linked sweeps.

    name : str (optional)
        Allows "naming" a model in a way similar to inherited instances,
        and overrides the inherited name if there is one.

    Attributes with side effects
    ----------------------------
    `program` is set during a solve
    `solution` is set at the end of a solve
    """

    # name and num identify a model uniquely
    name = None
    num = None
    # naming holds the name and num evironment in which a model was created
    # this includes its own name and num, and those of models containing it
    naming = None
    program = None
    solution = None

    def __init__(self, cost=None, constraints=None, *args, **kwargs):
        setup_vars = None
        substitutions = kwargs.pop("substitutions", None)  # reserved keyword
        if hasattr(self, "setup"):
            self.cost = None
            with NamedVariables(self.__class__.__name__):
                start_args = [cost, constraints]
                args = tuple(a for a in start_args if a is not None) + args
                cs = self.setup(*args, **kwargs)  # pylint: disable=no-member
                if (isinstance(cs, tuple) and len(cs) == 2
                        and isinstance(cs[1], dict)):
                    constraints, substitutions = cs  # TODO: remove
                else:
                    constraints = cs
                from .. import NAMEDVARS, MODELS, MODELNUMS
                setup_vars = NAMEDVARS[tuple(MODELS), tuple(MODELNUMS)]
                self.name, self.num = MODELS[-1], MODELNUMS[-1]
                self.naming = (tuple(MODELS), tuple(MODELNUMS))
            cost = self.cost  # TODO: remove
        elif args and not substitutions:
            # backwards compatibility: substitutions as third arg
            substitutions, = args

        cost = cost or Monomial(1)
        constraints = constraints or []
        if setup_vars:
            # add all the vars created in .setup to the Model's varkeys
            # even if they aren't used in any constraints
            self.unique_varkeys = frozenset(v.key for v in setup_vars)
        CostedConstraintSet.__init__(self, cost, constraints, substitutions)
        if hasattr(self, "setup") and self.__class__.__doc__:
            if (("Unbounded" in self.__class__.__doc__ or
                 "Bounded by" in self.__class__.__doc__) and
                    "SKIP VERIFICATION" not in self.__class__.__doc__):
                self.verify_docstring()

    gp = _progify_fctry(GeometricProgram)
    sp = _progify_fctry(SequentialGeometricProgram)
    solve = _solve_fctry(_progify_fctry(GeometricProgram, "solve"))
    localsolve = _solve_fctry(_progify_fctry(SequentialGeometricProgram,
                                             "localsolve"))

    def verify_docstring(self):  # pylint:disable=too-many-locals,too-many-branches,too-many-statements
        "Verifies docstring bounds are sufficient but not excessive."
        err = "while verifying %s:\n" % self.__class__.__name__
        bounded, meq_bounded = self.bounded.copy(), self.meq_bounded.copy()
        doc = self.__class__.__doc__
        exp_unbounds = expected_unbounded(self, doc)
        unexp_bounds = bounded.intersection(exp_unbounds)
        if unexp_bounds:  # anything bounded that shouldn't be? err!
            for direction in ["lower", "upper"]:
                badvks = [v for v, d in unexp_bounds if d == direction]
                if not badvks:
                    continue
                badvks = ", ".join(str(v) for v in badvks)
                badvks += (" were" if len(badvks) > 1 else " was")
                err += ("    %s %s-bounded; expected %s-unbounded"
                        "\n" % (badvks, direction, direction))
            raise ValueError(err)
        bounded.update(exp_unbounds)  # if not, treat expected as bounded
        add_meq_bounds(bounded, meq_bounded)  # and add more meqs
        self.missingbounds = {}  # now let's figure out what's missing
        for bound in meq_bounded:  # first add the un-dealt-with meq bounds
            for condition in list(meq_bounded[bound]):
                meq_bounded[bound].remove(condition)
                newcond = condition - bounded
                if newcond and not any(c.issubset(newcond)
                                       for c in meq_bounded[bound]):
                    meq_bounded[bound].add(newcond)
            bsets = " or ".join(str(list(c)) for c in meq_bounded[bound])
            self.missingbounds[bound] = (", but would gain it from any of"
                                         " these sets of bounds: " + bsets)
        # then add everything that's not in bounded
        if len(bounded)+len(self.missingbounds) != 2*len(self.varkeys):
            for key in self.varkeys:
                for bound in ("upper", "lower"):
                    if (key, bound) not in bounded:
                        if (key, bound) not in self.missingbounds:
                            self.missingbounds[(key, bound)] = ""
        if self.missingbounds:  # anything unbounded? err!
            boundstrs = "\n".join("  %s has no %s bound%s" % (v, b, x)
                                  for (v, b), x
                                  in self.missingbounds.items())
            docstring = ("To fix this add the following to %s's"
                         " docstring (you may not need it all):"
                         " \n" % self.__class__.__name__)
            for direction in ["upper", "lower"]:
                mb = [k for (k, b) in self.missingbounds if b == direction]
                if mb:
                    docstring += """
%s Unbounded
---------------
%s
""" % (direction.title(), ", ".join(set(k.name for k in mb)))
            raise ValueError(err + boundstrs + "\n\n" + docstring)

    def as_gpconstr(self, x0):
        "Returns approximating constraint, keeping name and num"
        cs = CostedConstraintSet.as_gpconstr(self, x0)
        cs.name, cs.num = self.name, self.num
        return cs

    def subconstr_str(self, excluded=None):
        "The collapsed appearance of a ConstraintBase"
        return "%s_%s" % (self.name, self.num) if self.name else None

    def subconstr_latex(self, excluded=None):
        "The collapsed appearance of a ConstraintBase"
        return "%s_{%s}" % (self.name, self.num) if self.name else None

    def sweep(self, sweeps, **solveargs):
        "Sweeps {var: values} pairs in sweeps. Returns swept solutions."
        sols = []
        for sweepvar, sweepvals in sweeps.items():
            original_val = self.substitutions.get(sweepvar, None)
            self.substitutions.update({sweepvar: ('sweep', sweepvals)})
            try:
                sols.append(self.solve(**solveargs))
            except InvalidGPConstraint:
                sols.append(self.localsolve(**solveargs))
            if original_val:
                self.substitutions[sweepvar] = original_val
            else:
                del self.substitutions[sweepvar]
        if len(sols) == 1:
            return sols[0]
        return sols

    def autosweep(self, sweeps, tol=0.01, samplepoints=100, **solveargs):
        """Autosweeps {var: (start, end)} pairs in sweeps to tol.

        Returns swept and sampled solutions.
        The original simplex tree can be accessed at sol.bst
        """
        sols = []
        for sweepvar, sweepvals in sweeps.items():
            sweepvar = self[sweepvar].key
            start, end = sweepvals
            bst = autosweep_1d(self, tol, sweepvar, [start, end], **solveargs)
            sols.append(bst.sample_at(np.linspace(start, end, samplepoints)))
        if len(sols) == 1:
            return sols[0]
        return sols

    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
    def debug(self, solver=None, verbosity=1, **solveargs):
        """Attempts to diagnose infeasible models.

        If a model debugs but errors in a process_result call, debug again
        with `process_results=False`
        """
        from .relax import ConstantsRelaxed, ConstraintsRelaxed
        from .bounded import Bounded

        sol = None

        solveargs["solver"] = solver
        solveargs["verbosity"] = verbosity - 1
        solveargs["process_result"] = False

        if verbosity:
            print("< DEBUGGING >")
            print("> Trying with bounded variables and relaxed constants:")

        bounded = Bounded(self)
        if self.substitutions:
            constsrelaxed = ConstantsRelaxed(bounded)
            feas = Model(constsrelaxed.relaxvars.prod()**30 * self.cost,
                         constsrelaxed)
            # NOTE: It hasn't yet been seen but might be possible that
            #       the self.cost component above could cause infeasibility
        else:
            feas = Model(self.cost, bounded)

        try:
            try:
                sol = feas.solve(**solveargs)
            except InvalidGPConstraint:
                sol = feas.localsolve(**solveargs)
            sol["boundedness"] = bounded.check_boundaries(sol)
            if self.substitutions:
                relaxed = get_relaxed([sol(r) for r in constsrelaxed.relaxvars],
                                      constsrelaxed.origvars,
                                      min_return=0 if sol["boundedness"] else 1)
                if verbosity and relaxed:
                    if sol["boundedness"]:
                        print("and these constants relaxed:")
                    else:
                        print("\nSolves with these constants relaxed:")
                    for (_, orig) in relaxed:
                        print("  %s: relaxed from %-.4g to %-.4g"
                              % (orig, mag(constsrelaxed.constants[orig.key]),
                                 mag(sol(orig))))
                    print
            if verbosity:
                print(">> Success!")
        except (ValueError, RuntimeWarning):
            if verbosity:
                print(">> Failure.")
                print("> Trying with relaxed constraints:")

            try:
                constrsrelaxed = ConstraintsRelaxed(self)
                feas = Model(constrsrelaxed.relaxvars.prod()**30 * self.cost,
                             constrsrelaxed)
                try:
                    sol = feas.solve(**solveargs)
                except InvalidGPConstraint:
                    sol = feas.localsolve(**solveargs)
                relaxed = get_relaxed(sol(constrsrelaxed.relaxvars),
                                      range(len(feas[0][0])))
                if verbosity and relaxed:
                    print("\nSolves with these constraints relaxed:")
                    for relaxval, i in relaxed:
                        constraint = feas[0][0][i][0]
                        # substitutions of the final relax value
                        conleft = constraint.left.sub(
                            {constrsrelaxed.relaxvars[i]: relaxval})
                        conright = constraint.right.sub(
                            {constrsrelaxed.relaxvars[i]: relaxval})
                        origconstraint = constrsrelaxed.origconstrs[i]
                        relax_percent = "%i%%" % (0.5+(relaxval-1)*100)
                        print(" %3i: %5s relaxed, from %s %s %s \n"
                              "                     to %s %s %s "
                              % (i, relax_percent, origconstraint.left,
                                 origconstraint.oper, origconstraint.right,
                                 conleft, constraint.oper, conright))
                if verbosity:
                    print("\n>> Success!")
            except (ValueError, RuntimeWarning):
                if verbosity:
                    print(">> Failure")
        if verbosity:
            print
        return sol

def get_relaxed(relaxvals, mapped_list, min_return=1):
    "Determines which relaxvars are considered 'relaxed'"
    sortrelaxed = sorted(zip(relaxvals, mapped_list), key=lambda x: x[0],
                         reverse=True)
    # arbitrarily 1.01 is the min that counts as "relaxed"
    mostrelaxed = max(sortrelaxed[0][0], 1.01)
    for i, (val, _) in enumerate(sortrelaxed):
        if i >= min_return and val <= 1.01 and (val-1) <= (mostrelaxed-1)/10:
            return sortrelaxed[:i]
    return sortrelaxed