import exodus, Silo

#
# Main inputs to the script (values are just for example data)
#
siloFile = "globe.silo"
matobjName = "mat1"
matsOfInterest = [2, 3]
fieldsOfInterest =["dx", "u"] # one nodal and one zonal
exoFile = "s2ex-out.e"

# Open the silo file
silodb = Silo.Open(siloFile)

# Get the specified Silo material object and all its data
mat = silodb.GetVarInfo(matobjName, 1)

# Get the mesh object associated with this material object.
# including all its raw data (e.g. 1 for second arg)
mesh = silodb.GetVarInfo(mat["meshid"], 1)

# Get the Silo zonelist for this mesh
zl = silodb.GetVarInfo(mesh["zonelist"], 1)

#
# Traverse *both* zonelist and material objects simulatenously
# searching for zones containing materials of interest. Silo's
# zonelist is broken into 'segments' of constant shape size/type.
# Different segments can involve the same shape size/type though.
#
zloff = 0
zoneidx = 0
zoneids = []
newzl = {"nshapes":0, "shapecnt":[], "shapesize":[], "nodelist":[]}
for i in range(0,zl["nshapes"]):
    shsz = zl["shapesize"][i]
    for j in range(0,zl["shapecnt"][i]):
        keepZone = 0
        if mat["matlist"][zoneidx] in matsOfInterest: # zone with one material
            keepZone = 1
        elif mat["matlist"][zoneidx] < 0: # zone with multiple materials
            mixidx = -zoneidx-1
            while mat["mix_next"][mixidx] != 0:
                if mat["mix_mat"][mixidx] in matsOfInterest:
                    keepZone = 1
                    break
                mixidx = mat["mix_next"][mixidx]
        if keepZone:
            zoneids.append(zoneidx)
            k = len(newzl["shapesize"])
            if k == 0 or newzl["shapesize"][k-1] != shsz:
                newzl["shapesize"].append(shsz)
                newzl["shapecnt"].append(1)
                newzl["nshapes"] += 1
            else:
                newzl["shapecnt"][k-1] += 1
            newzl["nodelist"].extend(zl["nodelist"][zloff:zloff+shsz])
        zloff += shsz
        zoneidx += 1

#
# The zoneids list contains the set of zones we want
# Get unique set of node ids the newzl comprises
#
#uniqnl = sorted(list(set(newzl["nodelist"])))
uniqnl = list(set(newzl["nodelist"]))

#
# Map the newzl's nodelist node ids through the new unique set of nodes
#
newzl["nodelist"] = [uniqnl.index(x) for x in newzl["nodelist"]]

#
# Get coords for these nodes from the mesh
#
coord0 = [mesh["coord0"][x] for x in uniqnl]
coord1 = [mesh["coord1"][x] for x in uniqnl]
coord2 = [mesh["coord2"][x] for x in uniqnl]

#
# Map remaining fields of interest capturing all in fdata dict
# Note: If mapping non-scalar fields (e.g. vector or tensors),
# then there will be more 'valueX' entries in 'var' for each
# of the components
#
fdata = {"nodals":{}, "zonals":{}}
for f in fieldsOfInterest:
    var = silodb.GetVarInfo(f, 1)
    if var["centering"] == Silo.DB_ZONECENT:
        fdata["zonals"][f] = [var["value0"][x] for x in zoneids]
    else:
        fdata["nodals"][f] = [var["value0"][x] for x in uniqnl]

#
# Close the silo file
#
silodb.Close()


#
# Util funcs for converting silo to exodus
#
def Silo2ExodusElemType(sdim, nnodes):
    if sdim == 2:
        if nnodes == 3:
            return "TRI"
        elif nnodes == 4:
            return "QUAD"
    elif sdim == 3:
        if nnodes == 4:
            return "TET"
        elif nnodes == 5:
            return "PYR"
        elif nnodes == 6:
            return "WED"
        elif nnodes == 8:
            return "HEX"
    return "UNK"

def ConvertSiloWedgeOrderingToExodus(connarr):
#    s2eWedgeOrder = [5, 2, 0, 3, 4, 1]
    s2eWedgeOrder = [2, 5, 1, 3, 4, 0]
    nzones = len(connarr) / 6
    for i in range(0, nzones):
        nodes = connarr[i*6:i*6+6]
        connarr[i*6:i*6+6] = [nodes[x] for x in s2eWedgeOrder]

# Open the exodus output file and set some sizing information
e = exodus.exodus(exoFile, mode='w', title="generated by Silo2Exodus", \
    array_type='ctype', numDims=mesh["ndims"], numNodes=len(coord0),
    numElems=len(zoneids), numBlocks=newzl["nshapes"], numNodeSets=0, numSideSets=0)

# specify number of elem vars
e.set_element_variable_number(len(fdata["zonals"]))

# specify number of node vars
e.set_node_variable_number(len(fdata["nodals"]))

# Output coordinates
e.put_coords(coord0, coord1, coord2)

# Loop over the zonelist, building up args for put_concat_elem_blk
blkids=[]
elemtypes=[]
elemcnts=[]
elemsizes=[]
attrcnts=[]
for i in range(0,newzl["nshapes"]):
    blkids.append(i+1)
    # Map Silo zone shape info to exodus elem type info
    elemtypes.append(Silo2ExodusElemType(mesh["ndims"], newzl["shapesize"][i]))
    elemcnts.append(newzl["shapecnt"][i])
    elemsizes.append(newzl["shapesize"][i])
    attrcnts.append(0)

# Put out elem block header info
e.put_concat_elem_blk(blkids,elemtypes,elemcnts,elemsizes,attrcnts,0)

# Put out connectivity for each element block
offset=0
for i in range(0,newzl["nshapes"]):
    cnt = newzl["shapecnt"][i]*newzl["shapesize"][i]
    exconn = [nid+1 for nid in newzl["nodelist"][offset:offset+cnt]]
    if Silo2ExodusElemType(mesh["ndims"], newzl["shapesize"][i]) == "WED":
        ConvertSiloWedgeOrderingToExodus(exconn)
    e.put_elem_connectivity(i+1, exconn)
    offset += cnt

# Put out nodal variables
idx=1
for v in fdata["nodals"]:
    e.put_node_variable_name(v,idx)
    e.put_node_variable_values(v, 1, fdata["nodals"][v])
    idx += 1

# Put out zonal (elem) variables (upon each elem block)
idx=1
for v in fdata["zonals"]:
    offset=0
    e.put_element_variable_name(v,idx)
    for i in range(0,newzl["nshapes"]):
        cnt = newzl["shapecnt"][i]
        vals = [zid for zid in fdata["zonals"][v][offset:offset+cnt]]
        e.put_element_variable_values(i+1, v, 1, vals)
        offset += cnt
    idx += 1

e.close()
