import numpy as np
import sys
sys.path.append('/homes/rlreed/workspace/unotran/src')
from coarseBounds import computeBounds, Grouping
import pickle
from makeDLPbasis import makeBasis as makeDLP
from makeKLTbasis import makeBasis as makeKLT
import sph
import sph_dgm
import pydgm


def buildGEO(ass_map):
    fine_map = [1]
    coarse_map = [1.26]

    material_map = [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]

    npins = len(ass_map)

    cm = [0.0]
    fm = []
    mm = []
    for i, ass in enumerate(ass_map):
        mm += material_map[ass]
        cm += coarse_map
        fm += fine_map
    cm = np.cumsum(cm)

    return npins, fm, cm, mm

def makeDGMXS(G, refXS, dgmstructure, basisType):
    if 'klt' in basisType:
        makeKLT(basisType, dgmstructure)
    else:
        makeDLP(dgmstructure)

    dgmstructure.fname = '{}_{}'.format(basisType, dgmstructure.fname)

    fname = '_homo.'.join(xs_name.split('.'))
    refXS.write_homogenized_XS(fname)

    nPin, fm, cm, mm = buildGEO(pin_map)

    dgm = sph_dgm.DGMSOLVER(G, fname, fm, cm, mm, nPin, dgmstructure, solveFlag=False)
    pydgm.dgmsolver.initialize_dgmsolver()
    dgm.extractInfo()

    pydgm.dgmsolver.finalize_dgmsolver()
    pydgm.control.finalize_control()

    nCellPerPin = dgm.phi.shape[2] // dgm.npin

    return sph_dgm.XS(G, nCellPerPin, dgm.sig_t, dgm.vsig_f, dgm.chi, dgm.sig_s)

if __name__ == '__main__':
    np.set_printoptions(precision=6)

    G = 44

    dgmstructure = computeBounds(G, 'full', 1, 0.0, 1.3, 60)
    fname = dgmstructure.fname
    xs_name = 'XS/{}gXS.anlxs'.format(G)
    pin_map = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    data_path = 'data2'

    # Get the homogenized cross sections
    refXS = pickle.load(open('{}/refXS_sph_space_{}.p'.format(data_path, G), 'rb'))

    for basis in ['dlp', 'klt_full', 'klt_combine', 'klt_pins_full']:
        dgmstructure.fname = fname
        XS = makeDGMXS(G, refXS, dgmstructure, basis)
        pickle.dump(XS, open('{}/refXS_dgm_{}_{}_h{}.p'.format(data_path, dgmstructure.fname, 'fine_mu', 0), 'wb'))


[((3, 0, 3, 54), 'sys.path.append', 'sys.path.append', ({(3, 16, 3, 53): '"""/homes/rlreed/workspace/unotran/src"""'}, {}), "('/homes/rlreed/workspace/unotran/src')", False, 'import sys\n'), ((28, 9, 28, 22), 'numpy.cumsum', 'np.cumsum', ({(28, 19, 28, 21): 'cm'}, {}), '(cm)', True, 'import numpy as np\n'), ((45, 10, 45, 86), 'sph_dgm.DGMSOLVER', 'sph_dgm.DGMSOLVER', (), '', False, 'import sph_dgm\n'), ((46, 4, 46, 42), 'pydgm.dgmsolver.initialize_dgmsolver', 'pydgm.dgmsolver.initialize_dgmsolver', ({}, {}), '()', False, 'import pydgm\n'), ((49, 4, 49, 40), 'pydgm.dgmsolver.finalize_dgmsolver', 'pydgm.dgmsolver.finalize_dgmsolver', ({}, {}), '()', False, 'import pydgm\n'), ((50, 4, 50, 36), 'pydgm.control.finalize_control', 'pydgm.control.finalize_control', ({}, {}), '()', False, 'import pydgm\n'), ((54, 11, 54, 80), 'sph_dgm.XS', 'sph_dgm.XS', ({(54, 22, 54, 23): 'G', (54, 25, 54, 36): 'nCellPerPin', (54, 38, 54, 47): 'dgm.sig_t', (54, 49, 54, 59): 'dgm.vsig_f', (54, 61, 54, 68): 'dgm.chi', (54, 70, 54, 79): 'dgm.sig_s'}, {}), '(G, nCellPerPin, dgm.sig_t, dgm.vsig_f, dgm.chi, dgm.sig_s)', False, 'import sph_dgm\n'), ((57, 4, 57, 36), 'numpy.set_printoptions', 'np.set_printoptions', (), '', True, 'import numpy as np\n'), ((61, 19, 61, 60), 'coarseBounds.computeBounds', 'computeBounds', ({(61, 33, 61, 34): 'G', (61, 36, 61, 42): '"""full"""', (61, 44, 61, 45): '1', (61, 47, 61, 50): '0.0', (61, 52, 61, 55): '1.3', (61, 57, 61, 59): '60'}, {}), "(G, 'full', 1, 0.0, 1.3, 60)", False, 'from coarseBounds import computeBounds, Grouping\n'), ((34, 8, 34, 40), 'makeKLTbasis.makeBasis', 'makeKLT', ({(34, 16, 34, 25): 'basisType', (34, 27, 34, 39): 'dgmstructure'}, {}), '(basisType, dgmstructure)', True, 'from makeKLTbasis import makeBasis as makeKLT\n'), ((36, 8, 36, 29), 'makeDLPbasis.makeBasis', 'makeDLP', ({(36, 16, 36, 28): 'dgmstructure'}, {}), '(dgmstructure)', True, 'from makeDLPbasis import makeBasis as makeDLP\n')]