#!/usr/bin/python
# Mesh specific imports                                                                                                                                       
try:
    from itaps import iMesh, iBase

    HAVE_PYTAPS = True
except ImportError:

    HAVE_PYTAPS = False

try:
    from pyne.material import Material, MaterialLibrary
    HAVE_MATERIAL = True
except ImportError:
    HAVE_MATERIAL = False

try:
    from pyne.tally import Tally
    HAVE_TALLY = True
except ImportError:
    HAVE_TALLY = False

try:
    from pyne.particle import is_valid,name
    HAVE_PARTICLE = True
except ImportError:
    HAVE_PARTICLE = False


import string
import argparse
import os

fluka_predefined_mat=["BERYLLIU","BARIUM","BISMUTH","BROMINE","RHENIUM","RUTHERFO","ROENTGEN","HYDROGEN","PHOSPHO","GERMANIU",
"GADOLINI","GALLIUM","HASSIUM","ZINC","HAFNIUM","MERCURY","HELIUM","PRASEODY","PLATINUM","239-PU","LEAD","CARBON", "POTASSIU","OXYGEN","SULFUR","TUNGSTEN","EUROPIUM","MAGNESIU","MOLYBDEN","MANGANES","URANIUM","IRON","NICKEL","NITROGEN",
"SODIUM","NIOBIUM","NEODYMIU","NEON","ZIRCONIU","BORON","COBALT","FLUORINE","CALCIUM","CERIUM","CADMIUM","VANADIUM","CESIUM",
"CHROMIUM","COPPER","STRONTIU","KRYPTON","SILICON","TIN","SAMARIUM","SCANDIUM","ANTIMONY","LANTHANU","CHLORINE","LITHIUM",
"TITANIUM","TERBIUM","99-TC", "TANTALUM","SILVER","IODINE","IRIDIUM","241-AM","ALUMINUM","ARSENIC","ARGON","GOLD","INDIUM",
"YTTRIUM","XENON","HYDROG-1","DEUTERIU","TRITIUM","HELIUM-3","HELIUM-4","LITHIU-6","LITHIU-7","BORON-10","BORON-11","90-SR",
"129-I","124-XE","126-XE","128-XE","130-XE","131-XE","132-XE","134-XE","135-XE","136-XE","135-CS","137-CS", "230-TH","232-TH",
"233-U","234-U","235-U","238-U"]

"""
function that gets all tags on dagmc geometry
------------------------------
filename : the dagmc filename
return vector of tag_values
"""
def get_tag_values(filename):
    dag_vol_names = [] # list of dag volume names (Cubit id)
#    dag_materials = {} # list of material assignments from group 
    dag_properties = set()

    # material tags
    dag_material_tags = []
    # tally tags
    dag_tally_tags=[]
    # tally assignments
    tally_assigns=[]

    # create imesh instance  
    dag_geom = iMesh.Mesh()
    # load the file
    dag_geom.load(filename)

    # get all entities      
    ents = dag_geom.getEntities()
    # create a  mesh set         
    mesh_set = dag_geom.getEntSets()
    # list of volume ent handles 
    mat_list = []
    # get all geom_dimension ents
    geom_list = []
    cat_list = []

    vol_tag = dag_geom.getTagHandle('GEOM_DIMENSION')
    name_tag = dag_geom.getTagHandle('GLOBAL_ID')
    mat_tag = dag_geom.getTagHandle('NAME')
    cat_tag = dag_geom.getTagHandle('CATEGORY')

    # get the list we need    
    for i in mesh_set:
        tags = dag_geom.getAllTags(i)
        for tag in tags:
            if tag == vol_tag:
                geom_list.append(i)
            if tag == mat_tag:
                mat_list.append(i)
            if tag == cat_tag:
                cat_list.append(i)

    # for the 3d entities     
    for entity in geom_list:
        if vol_tag[entity] == 3:
             dag_vol_names.append(str(name_tag[entity]))

    # loop over all the volumes 
    for entity in geom_list:
        # loop over the material sets
        for meshset in mat_list:
            # if volume in set       
            if meshset.contains(entity):
                mat_name = mat_tag[meshset]
                volume_name = name_tag[entity]
                # dag_materials[volume_name]="".join( chr( val ) for val in mat_name )
                # dag_properties.add("".join( chr( val ) for val in mat_name))
                dag_properties.add(tag_to_script(mat_name))
                if 'tally:' in tag_to_script(mat_name):
                    pair = (volume_name,tag_to_script(mat_name))
                    tally_assigns.append(pair)

    # now we have dag properties, create one with materials and one with tallies
    for tag in dag_properties:
        if 'mat:' in tag:
            dag_material_tags.append(tag)
        if 'tally:' in tag:
            dag_tally_tags.append(tag)

#    print dag_vol_names, dag_properties
    return dag_material_tags,tally_assigns


    """

    mesh = iMesh.Mesh()
    mesh.load(filename)
    # get all entities
    ents = mesh.getEntities(iBase.Type.all, iMesh.Topology.triangle)
    # get mesh set
    mesh_sets = mesh.getEntSets()
    # tag_values = []  # list of tag values


    tag_values = []
    found_all_tags = 0
    for s in mesh_sets:
        if found_all_tags == 1:
            break
        # get all the tags
        tags = mesh.getAllTags(s)
        # loop over the tags checking for name
        for t in tags:
            # look for NAME tag
            if t.name == 'NAME':
                # the handle is the tag name
                t_handle = mesh.getTagHandle(t.name)
                # get the data for the tag, with taghandle in element i
                tag = t_handle[s]
                tag_to_script(tag, tag_values)
                # last tag we are done
                if any('impl_complement' in s for s in tag_values):
                    found_all_tags = 1
    print('The groups found in the h5m file are: ')
    
    """

    return tag_values

"""
function to transform the tags into strings
tag : string of the tag to add to tag_list
tag_list : vector of tags in the problem
returns tag_list
"""
def tag_to_script(tag):
    a = []
    # since we have a byte type tag loop over the 32 elements
    for part in tag:
        # if the byte char code is non 0
        if (part != 0):
            # convert to ascii
            a.append(str(unichr(part)))
            # join to end string
            test = ''.join(a)
            # the the string we are testing for is not in the list of found
            # tag values, add to the list of tag_values

    return test

"""
function which loads pyne material library
filename : string of the name of the material library
returns PyNE MaterialLibrary instance
"""
def load_mat_lib(filename):
    mat_lib = MaterialLibrary()
    mat_lib.from_hdf5(
        filename, datapath='/material_library/materials', nucpath='/material_library/nucid')
    return mat_lib

"""
function to check that material group names exist and creates
a list of names and densities if provided
-------------------------------------------------
tag_values - list of tags from dagmc file
mat_lib - pyne material library instance
returns mat_dens_list, a zipped pair of density and material name
"""
def check_matname(tag_values):
    # loop over tags
    g = 0
    mat_list_matname = []  # list of material names
    mat_list_density = []  # list of density if provided in the group names
    # loop over the tags in the file
    for tag in tag_values:
        print tag
        if 'Graveyard' in tag or 'graveyard' in tag:
            g = 1
            continue
        if 'Vacuum' in tag or 'vacuum' in tag:
            continue
    # look for mat, this id's material in group name
        if tag[0:3] == 'mat':
            # split on the basis of "/" being delimiter and split colons from
            # name
            if '/' in tag:
                mat_name = tag.split('/')
                if ':' not in mat_name[0]:
                    raise Exception("Could not find group name in appropriate format; ':' is absent in %s" % tag)
                # list of material name only
                matname = mat_name[0].split(':')
                if matname[1] == '':
                    raise Exception("Could not find group name in appropriate format; material name in blank  %s" % tag)
                if mat_name[1] == '':
                    raise Exception("Could not find group name in appropriate format; extra \'/\' in %s" % tag)
                if ':' not in mat_name[1]:
                    raise Exception("Could not find group name in appropriate format; ':' is absent after the '/' in %s" % tag)
                matdensity = mat_name[1].split(':')
                try:
                    matdensity_test = float(matdensity[1])
                except:
                    raise Exception("Could not find density in appropriate format!; density is not a float in %s" % tag)
                mat_list_density.append(matdensity[1])
            # otherwise we have only "mat:"
            elif ':' in tag:
                splitted_group_name = mat_split(tag)
            else:
                raise Exception(
                    "Couldn\'t find group name in appropriate format; ': is absent' in  %s" % tag)
            mat_list_matname.append(splitted_group_name['material'])
            mat_list_density.append(splitted_group_name['density'])
    if g == 0:
        raise Exception("Graveyard group is missing! You must have a graveyard")
    mat_dens_list = zip(mat_list_matname, mat_list_density)
    # error conditions, no tags found
#    if len(mat_dens_list) == 0:
#        raise Exception("No material group names found, you must have materials")

    return mat_dens_list


"""
splitting group names containing both materials and densities
"""


def mat_dens_split(tag):
    splitted_group_name = {}
    mat_name = tag.split('/')
    if ':' not in mat_name[0]:
        raise Exception(
            "Couldn\'t find group name in appropriate format; ':' is absent in %s" % tag)
    # list of material name only
    matname = mat_name[0].split(':')
    if len(matname) > 2:
        raise Exception(
            "Wrong format for group names! %s. correct: mat:NAME/rho:VALUE or mat:NAME" % tag)
    if matname[1] == '':
        raise Exception(
            "Couldn\'t find group name in appropriate format; wrong material name in %s" % tag)
    splitted_group_name['material'] = matname[1]
    if mat_name[1] == '':
        raise Exception(
            "Couldn\'t find group name in appropriate format; extra \'/\' in %s" % tag)
    if ':' not in mat_name[1]:
        raise Exception(
            "Couldn\'t find group name in appropriate format; ':' is absent after the '/' in %s" % tag)
    matdensity = mat_name[1].split(':')
    try:
        matdensity_test = float(matdensity[1])
    except:
        raise Exception(
            "Couldn\'t find density in appropriate format!; density is not a float in %s" % tag)
    splitted_group_name['density'] = matdensity[1]
    return splitted_group_name

"""
splitting group names containing only materials
"""


def mat_split(tag):
    splitted_group_name = {}
    matname = tag.split(':')
    if len(matname) > 2:
        raise Exception(
            "Wrong format for group names! %s. correct: mat:NAME/rho:VALUE or mat:NAME" % tag)
    if matname[1] == '':
        raise Exception(
            "Couldn\'t find group name in appropriate format; wrong material name in %s" % tag)
    splitted_group_name['material'] = matname[1]
    splitted_group_name['density'] = ''
    return splitted_group_name

"""
function that checks the existence of material names on the PyNE library 
and creates a list of materials with attributes set
-------------------------------------------------------------
material_list : vector of material_name & density pairs
mat_lib : PyNE Material library object
"""
def check_and_create_materials(material_list, mat_lib):
    flukamaterials_list = []
    material_object_list = []
    d = 0
    # loop over materials in geometry
    for g in range(len(material_list)):
        material = material_list[g][0]
        # loop over materials in library
        for key in mat_lib.iterkeys():
            if material == key:
                d = d + 1
                # get the material
                new_mat = mat_lib.get(key)[:]
                flukamaterials_list.append(material)
                copy_metadata(new_mat, mat_lib.get(key))
                # set the mcnp material number and fluka material name
                set_metadata(new_mat, d, flukamaterials_list)

                # rename the material to match the group
                group_name = "mat:"+material_list[g][0]

                if material_list[g][1] is not '':
                    group_name += "/rho:"+material_list[g][1]
                    new_mat.density = float(material_list[g][1])
                                        
                new_mat.metadata['name'] = group_name

                material_object_list.append(new_mat)
                break
            if mat_lib.keys().index(key) == len(mat_lib.keys()) - 1:
                print('Material {%s} doesn\'t exist in pyne material lib' % material)
                print_near_match(material, mat_lib)
                raise Exception('Couldn\'t find exact match in material library for : %s' % material)

    # check that there are as many materials as there are groups
    if d != len(material_list):
        raise Exception("There are insuficient materials")

    # return the list of material objects to write to file
    print material_object_list
    return material_object_list

"""
function to copy the metadata of materials from the PyNE material library
-------------------------------------
material : PyNE material object to copy data into 
material_from_lib : PyNE material objec to copy data from
"""
def copy_metadata(material, material_from_lib):
    # copy metadata from lib to material
    for key in list(material_from_lib.metadata.keys()):
        material.metadata[key] = material_from_lib.metadata[key]

    material.density = material_from_lib.density
    material.mass = material_from_lib.mass
    material.atoms_per_molecule = material_from_lib.atoms_per_molecule
    material.comp = material_from_lib.comp
    return material


"""
function to set the attributes of the materials:
----------------------------------------
mat : PyNE Material Object
number : mcnp material number
flukamat_list : 
returns : PyNE Material Object
"""
def set_metadata(mat, number, flukamat_list):
    mat.metadata['mat_number'] = str(number)
    mat.metadata['fluka_material_index'] = str(number+25)
    fluka_material_naming(mat, flukamat_list)
    return mat


"""
Function to prepare fluka material names:
"""
def fluka_material_naming(material, flukamat_list):
    matf = material.metadata['name']
    matf = ''.join(c for c in matf if c.isalnum())
    L=len(matf)
    if len(matf) > 8:
        matf = matf[0:8]
    else:
        pass
    # if name is in list, change name by appending number
    if (matf.upper() in flukamat_list) or (matf.upper() in fluka_predefined_mat):
        for a in range(len(flukamat_list)+1):
            a = a + 1
            if (a <= 9):
                if (len(matf) == 8):
                    matf = matf.rstrip(matf[-1])
                    matf = matf + str(a)
                elif (len(matf) < 8):
                    matf=matf[0:L]
                    matf=matf + str(a)
            elif (a > 9):  
                if (len(matf) == 8): 
                    for i in range(len(str(a))):
                        matf = matf.rstrip(matf[-1])
                    matf = matf + str(a)
                elif (len(matf) < 8) and (8-len(matf) >= len(str(a))) :
                    matf=matf[0:L]
                    matf=matf + str(a)   
                elif (len(matf) < 8) and (8-len(matf) < len(str(a))) : 
                    difference= len(str(a)) - (8-len(matf)) 
                    for i in range(difference):
                        matf = matf.rstrip(matf[-1])
                    matf = matf + str(a)  
            if (matf.upper() in flukamat_list) or (matf.upper() in fluka_predefined_mat):
                continue
            else:
                flukamat_list.append(matf.upper())
                break
    # otherwise uppercase
    else:
        flukamat_list.append(matf.upper())
    material.metadata['fluka_name'] = matf.upper()
    return material

""" 
function to print near matches to material name
"""
def print_near_match(material, material_library):
    list_of_matches = []
    for item in material_library.iterkeys():
        if (material.lower() in item.lower()) or (material.upper() in item.upper()):
            print("Near matches to %s are :" % material)
            print item
        list_of_matches.append(item)
    return list_of_matches

"""
Function that writes material objects to hdf5 file
-------
material_list: list of PyNE Material Objects
filename: filename to write the objects to
"""
def write_mats_h5m(materials_list, filename):
    new_matlib = MaterialLibrary()
    for material in materials_list:
        # using fluka name as index since this is unique
        new_matlib[material.metadata['name']] = material
    new_matlib.write_hdf5(filename)

"""
Function that writes tally objects to hdf5 file
-------
tally_list: list of volume id tally pairs of PyNE Material Objects
filename: filename to write the objects to
"""
def write_tally_h5m(tally_list,filename):
    # loop over list
    for tally in tally_list:
        name_flux = tally[1]
        
        particle_name = (name_flux.split(':'))[1].split('/')[0]
        if not is_valid(particle_name):
            print "particle: "+particle_name+" in tally "+name_flux+" in volume "+str(tally[0])
            exit
        tally_type = (name_flux.split('/'))[1].split('\n')[0]
        tally_name = particle_name[0]+particle_name[1].upper()+tally_type[0].upper()+tally_type[1].upper()+str(tally[0])
        print particle_name, tally_type, tally_name, tally[0], str(tally[0])
        new_tally = Tally(tally_type,particle_name,tally[0],"Volume",str(tally[0]),tally_name,0.0,1.0)
        new_tally.write_hdf5(filename,"/tally")
"""
function to parse the script, adding options:
defining 
-f  : the .h5m file path
-d  : nuc_data path
-o  : name of the output h5m file "NAME.h5m"
"""
def parsing():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-f', action='store', dest='datafile', help='The path to the .h5m file')
    parser.add_argument('-d', action='store', dest='nuc_data',
                        help='The path to the PyNE materials library nuc_data.h5')
    parser.add_argument(
        '-o', action='store', dest='output', help='The name of the output file ***.h5m')
    args = parser.parse_args()
    if not args.datafile:
        raise Exception('h5m file path not specified!!. [-f] is not set')
    if not args.nuc_data:
        raise Exception('nuc_data file path not specified!!. [-d] is not set')
    if not args.output:
        args.output = 'output.h5m'
    return args

"""
main
"""
def main():
    # parse the script
    args = parsing()
    # first add atomic mass to file
    launch_script = "nuc_data_make -m atomic_mass -o "+args.output
    print launch_script
    os.system(launch_script)
    # get list of tag values
    (tag_values,tally_assigns) = get_tag_values(args.datafile)
    # now load material library
    mat_lib = load_mat_lib(args.nuc_data)
    # check that material tags exist in library # material_list is list of
    # pyne objects in problem
    mat_dens_list = check_matname(tag_values)
    # create material objects from library
    material_object_list = check_and_create_materials(
        mat_dens_list, mat_lib)
    # write materials to file
    if(len(material_object_list) != 0 ):
        write_mats_h5m(material_object_list, args.output)
    
    # now get tally data and write to file
    write_tally_h5m(tally_assigns,args.output)

if __name__ == "__main__":
    main()
