import glob
import os
import yaml

cur_dir = os.path.dirname(__file__)
geo_dir = os.path.join(cur_dir, "../data/assembly")


class Assembly(object):
    """
    Holds all information required for modeling a basic assembly. Can be inherited to make specific assembly type.
    Current types include:
    Fuel Assembly
    """

    def __init__(self, assemblyInformation):
        self.assemblyDesignation = assemblyInformation[0]
        self.assemblyPosition = assemblyInformation[1]
        self.globalVars = assemblyInformation[2]
        self.core = assemblyInformation[3]
        self.universe = self.globalVars.universe
        self.cellNum = self.globalVars.cellNumber
        self.surfaceNum = self.globalVars.surfaceNumber
        self.materialNum = self.globalVars.materialNumber
        self.xcSet = self.globalVars.xc_set
        self.voidPercent = self.globalVars.void_per
        self.assemblyType = ''
        self.pinsPerAssembly = 0
        self.assemblyPitch = 0.0
        self.ductInnerFlatToFlat = 0.0
        self.ductOuterFlatToFlat = 0.0
        self.ductOuterFlatToFlatMCNPEdge = 0.0
        self.assemblyGap = 0.0
        self.assemblyHeight = 0.0
        self.zPosition = 0.0
        self.coolantMaterial = ''
        self.assemblyMaterial = ''

    def getAssemblyInfo(self, inputs):
        """Assign assembly parameters based on yaml Assembly file."""
        self.pinsPerAssembly = int(inputs['Pins Per Assembly']) if 'Pins Per Assembly' in inputs else 0
        self.assemblyPitch = float(inputs['Assembly Pitch'])
        self.ductInnerFlatToFlat = float(inputs['Duct Inside Flat to Flat']) / 2 if 'Duct Inside Flat to Flat' in \
                                                                                inputs else 0.0
        thickness = float(inputs['Duct Thickness']) if 'Duct Thickness' in inputs else 0.0
        self.ductOuterFlatToFlat = self.ductInnerFlatToFlat + thickness
        self.ductOuterFlatToFlatMCNPEdge = self.ductOuterFlatToFlat * 1.00005
        self.assemblyGap = float(inputs['Assembly Gap']) if 'Assembly Gap' in inputs else 0.0
        self.assemblyHeight = float(inputs['Assembly Height'])
        if self.core is None:
            self.coolantMaterial = inputs['Coolant']
        else:
            self.coolantMaterial = self.core.coolantMaterial
        self.assemblyMaterial = inputs['Assembly Material']
        self.zPosition = float(inputs['Z Position']) if 'Z Position' in inputs else 0.0

    def updateIdentifiers(self, universeTest):
        """Updates cell, surface, material, and universe number to create uniqueness"""
        self.cellNum += 1
        self.surfaceNum += 1
        self.materialNum += 1
        if universeTest:
            self.universe += 1


def assemblyTypeReader(assemblyYamlFile):
    """Reads in the assembly type to determine what type of assembly to build."""
    with open(assemblyYamlFile[0], "r") as mat_file:
        inputs = yaml.safe_load(mat_file)
        assemblyType = inputs['Assembly Type']
    return assemblyType


def getAssemblyLocation(assemblyType):
    """Find the file location for the assembly, and determine if the given assembly exists."""
    assemblyYamlFile = glob.glob(os.path.join(geo_dir, assemblyType + '.yaml'))
    try:
        assert assemblyYamlFile[0][(-(len(assemblyType)+5)):] == '{}.yaml'.format(assemblyType)
    except AssertionError:
        print('No assembly type named {}. Change your assembly type to a previously created assembly, '
              'or create a new assembly.'.format(assemblyType))
    return assemblyYamlFile

import glob
import os
import yaml

cur_dir = os.path.dirname(__file__)
geo_dir = os.path.join(cur_dir, "../data/assembly")


class Assembly(object):
    """
    Holds all information required for modeling a basic assembly. Can be inherited to make specific assembly type.
    Current types include:
    Fuel Assembly
    """

    def __init__(self, assemblyInformation):
        self.assemblyDesignation = assemblyInformation[0]
        self.assemblyPosition = assemblyInformation[1]
        self.globalVars = assemblyInformation[2]
        self.core = assemblyInformation[3]
        self.universe = self.globalVars.universe
        self.cellNum = self.globalVars.cellNumber
        self.surfaceNum = self.globalVars.surfaceNumber
        self.materialNum = self.globalVars.materialNumber
        self.xcSet = self.globalVars.xc_set
        self.voidPercent = self.globalVars.void_per
        self.assemblyType = ''
        self.pinsPerAssembly = 0
        self.assemblyPitch = 0.0
        self.ductInnerFlatToFlat = 0.0
        self.ductOuterFlatToFlat = 0.0
        self.ductOuterFlatToFlatMCNPEdge = 0.0
        self.assemblyGap = 0.0
        self.assemblyHeight = 0.0
        self.zPosition = 0.0
        self.coolantMaterial = ''
        self.assemblyMaterial = ''

    def getAssemblyInfo(self, inputs):
        """Assign assembly parameters based on yaml Assembly file."""
        self.pinsPerAssembly = int(inputs['Pins Per Assembly']) if 'Pins Per Assembly' in inputs else 0
        self.assemblyPitch = float(inputs['Assembly Pitch'])
        self.ductInnerFlatToFlat = float(inputs['Duct Inside Flat to Flat']) / 2 if 'Duct Inside Flat to Flat' in \
                                                                                inputs else 0.0
        thickness = float(inputs['Duct Thickness']) if 'Duct Thickness' in inputs else 0.0
        self.ductOuterFlatToFlat = self.ductInnerFlatToFlat + thickness
        self.ductOuterFlatToFlatMCNPEdge = self.ductOuterFlatToFlat * 1.00005
        self.assemblyGap = float(inputs['Assembly Gap']) if 'Assembly Gap' in inputs else 0.0
        self.assemblyHeight = float(inputs['Assembly Height'])
        if self.core is None:
            self.coolantMaterial = inputs['Coolant']
        else:
            self.coolantMaterial = self.core.coolantMaterial
        self.assemblyMaterial = inputs['Assembly Material']
        self.zPosition = float(inputs['Z Position']) if 'Z Position' in inputs else 0.0

    def updateIdentifiers(self, universeTest):
        """Updates cell, surface, material, and universe number to create uniqueness"""
        self.cellNum += 1
        self.surfaceNum += 1
        self.materialNum += 1
        if universeTest:
            self.universe += 1


def assemblyTypeReader(assemblyYamlFile):
    """Reads in the assembly type to determine what type of assembly to build."""
    with open(assemblyYamlFile[0], "r") as mat_file:
        inputs = yaml.safe_load(mat_file)
        assemblyType = inputs['Assembly Type']
    return assemblyType


def getAssemblyLocation(assemblyType):
    """Find the file location for the assembly, and determine if the given assembly exists."""
    assemblyYamlFile = glob.glob(os.path.join(geo_dir, assemblyType + '.yaml'))
    try:
        assert assemblyYamlFile[0][(-(len(assemblyType)+5)):] == '{}.yaml'.format(assemblyType)
    except AssertionError:
        print('No assembly type named {}. Change your assembly type to a previously created assembly, '
              'or create a new assembly.'.format(assemblyType))
    return assemblyYamlFile

Holds all information required for modeling a basic assembly. Can be inherited to make specific assembly type.     Current types include:     Fuel Assembly Assign assembly parameters based on yaml Assembly file. Updates cell, surface, material, and universe number to create uniqueness Reads in the assembly type to determine what type of assembly to build. Find the file location for the assembly, and determine if the given assembly exists.