'''Writer for TEM3P input files
'''
import csv
import logging
import os
import string

import smtk

from . import basewriter, cardformat
reload(basewriter)
reload(cardformat)
from cardformat import CardFormat
from . import utils

class Tem3PWriter(basewriter.BaseWriter):
    '''
    '''
    def __init__(self, solver):
        super(Tem3PWriter, self).__init__()
        self.base_indent = '  '  # overwrite default in base object
        self.elastic_materials = dict()
        self.name_prefix = None  # needed to disambiguate atts for ThermoElastic
        self.solver = solver
        if self.is_elastic_solver() or solver == 'tem3p-thermo-elastic':
            # Load the elastic materials table
            source_dir = os.path.dirname(__file__)
            path = os.path.join(source_dir, os.pardir, 'elastic-material.csv')
            if not os.path.exists(path):
                raise Exception('Cannot find materials file at %s' % path)

            with open(path) as csvfile:
                reader = csv.reader(csvfile)
                self.elastic_material = dict()
                for row in reader:
                    name = row[0]
                    # Remove any whitespace and '-', then replace ',' with '-'
                    name = name.replace(' ', '')  # remove space chars
                    name = name.replace('-', '')  # remove dash chars
                    name = name.replace(',','-')  # replace comma with dash
                    #print name
                    self.elastic_materials[name] = row
            #print 'self.elastic_materials:', self.elastic_materials

    def is_elastic_solver(self):
        return self.solver in set(['tem3p-eigen', 'tem3p-elastic', 'tem3p-harmonic'])

    def is_thermal_solver(self):
        return self.solver in set(['tem3p-thermal-linear', 'tem3p-thermal-nonlinear'])

    def write(self, scope):
        self.scope = scope

        if self.solver == 'tem3p-thermo-elastic':
            self.start_command('ThermoElasticProblem', indent='')
            self.scope.output.write('  RunId: 3\n')
            self.finish_command(indent='')

            # Write thermal section
            categories = scope.sim_atts.analysisCategories('TEM3P Thermal Nonlinear')
            scope.categories = list(categories)
            print 'Using categories: %s' % scope.categories
            self.solver = 'tem3p-thermal-nonlinear'
            self.name_prefix = 'Thermal'
            self.write_command()

            # Write elastic section
            categories = scope.sim_atts.analysisCategories('TEM3P Elastic')
            scope.categories = list(categories)
            self.solver = 'tem3p-elastic'
            self.name_prefix = 'Elastic'
            self.write_command()

            # Restore self.solver
            self.solver = 'tem3p-thermo-elastic'
            self.name_prefix = None
            scope.categories = list()
        elif self.is_elastic_solver():
            self.name_prefix = 'Elastic'
            self.write_command()
        elif self.is_thermal_solver():
            self.name_prefix = 'Thermal'
            self.write_command()
        else:
            raise Exception('Unrecognized solver type %s' % self.solver)


    def write_command(self, indent=''):
        '''Internal method to write either ElasticProblem or ThermostaticProblem
        '''
        prefix = None
        if self.is_elastic_solver():
            self.start_command('ElasticProblem', indent='')
            prefix = 'Elastic'
        elif self.is_thermal_solver():
            self.start_command('ThermostaticProblem', indent='')
            prefix='Thermal'
        else:
            raise Exception('Unrecognized solver type %s' % self.solver)
        # Write input mesh and order
        # Write mesh filename (ncdf extension)
        # (If input is .gen file, must convert using acdtool)
        root, ext = os.path.splitext(self.scope.model_file)
        mesh_file = root + '.ncdf'
        logging.info('ncdf model_file %s' % mesh_file)
        self.scope.output.write('  MeshFile: %s\n' % mesh_file)

        # BasisOrder and CurvedSurfaces
        order_att = self.get_attribute('Order')
        args = (self.scope, order_att)
        CardFormat('BasisOrder').write(*args)
        CardFormat('CurvedSurfaces').write(*args)

        # Write remaining parts
        self.write_meshdump()
        self.write_solver()
        if self.is_elastic_solver():
            self.write_elastic_material()
        elif self.is_thermal_solver():
            self.write_thermal_material()
            self.write_thermal_shell()
            self.write_heat_source()
        self.write_boundary_condition()

        if self.solver == 'tem3p-harmonic':
            self.write_harmonic_analysis()

        self.finish_command(indent='')  # Problem


    def write_meshdump(self):
        '''Write the meshdump section
        '''
        if self.solver in ['tem3p-eigen', 'tem3p-elastic']:
            print 'Writing MeshDump section'
            meshdump_att = self.get_attribute('MeshDump')
            self.start_command('MeshDump', indent='  ')
            self.write_standard_items(meshdump_att, skip_list=['Source'], indent='    ')

            deformed_em_item = meshdump_att.itemAtPath(
                'DeformedMeshFiles/WriteDeformedEMMesh', '/')
            if deformed_em_item.isEnabled():
                source_item = meshdump_att.itemAtPath(
                    'DeformedMeshFiles/WriteDeformedEMMesh/Source', '/')
                source_value = source_item.value(0)
                if source_value == 'NERSCDirectory':
                    dir_item = source_item.activeChildItem(0)

                    remote_path = dir_item.value(0).rstrip('/')
                    self.set_symlink(remote_path)
                    mesh_dir = os.path.basename(remote_path)
                    self.scope.output.write('    EMMeshInputDir: %s\n' % mesh_dir)
                else:
                    raise Exception('Unrecognized source type %s' % source_value)
            self.finish_command(indent='  ')  # MeshDump


    def write_solver(self):
        '''Writes the solver specification(s)

        '''
        print 'Writing solver'
        if self.solver == 'tem3p-eigen':
            self.write_standard_instance_att('EigenSolver')
        else:
            # All non-eigen cases use linear solver
            #self.write_standard_instance_att()
            att = self.get_attribute('TEM3PLinearSolver')
            self.start_command('LinearSolver')
            keyword_map = {'Type': 'Solver'}
            self.write_standard_items(att, keyword_table=keyword_map)
            self.finish_command()

        # In addition...
        if self.solver == 'tem3p-harmonic':
            self.write_standard_instance_att('HarmonicAnalysis')

        if self.solver == 'tem3p-thermal-nonlinear':
            self.write_standard_instance_att('PicardSolver')
            keyword_map = {'Type': 'Solver'}
            self.write_standard_instance_att('NonlinearSolver', keyword_table=keyword_map)


    def write_elastic_material(self):
        ''''''
        print 'Writing elastic materials'
        att_list = self.scope.sim_atts.findAttributes('TEM3PElasticMaterial')
        if not att_list:
            return
        att_list.sort(key=lambda att:att.name())

        for att in att_list:
            # Separate command for each model entity (requried?)
            ent_idlist = utils.get_entity_ids(self.scope, att.associations())
            if not ent_idlist:
                print 'Warning: attribute %s not associated to any model entities' % \
                    att.name()
                continue

            type_item = att.findString('Material')
            material_type = type_item.value(0)

            for ent_id in ent_idlist:
                self.start_command('VolumeMaterial')
                self.scope.output.write('    Id: %d\n' % ent_id)
                if material_type == 'Custom':
                    self.write_standard_items(att, silent_list=['Material'])
                elif material_type in self.elastic_materials:
                    material_tuple = self.elastic_materials[material_type]
                    # print 'material_tuple', material_tuple
                    label,density,poisson,youngs,elastic_alpha,electric_conductivity,name,temp = material_tuple
                    self.scope.output.write('    PoissonsRation: %s\n' % poisson)
                    self.scope.output.write('    YoungsModulus: %s\n' % youngs)
                    if density:
                        self.scope.output.write('    Density: %s\n' % density)
                    if elastic_alpha:
                        self.scope.output.write('    ElasticAlpha: %s\n' % elastic_alpha)
                else:
                    raise Exception(
                        'Internal error - Unrecognized material type %s' % material_type)
                self.finish_command()

    def write_thermal_material(self):
        '''Writes ThermalConductivity

        '''
        print 'Writing thermal materials'
        att_list = self.scope.sim_atts.findAttributes('TEM3PThermalMaterial')
        if not att_list:
            return
        att_list.sort(key=lambda att:att.name())
        print att_list

        for att in att_list:
            # Separate command for each model entity (requried?)
            ent_idlist = utils.get_entity_ids(self.scope, att.associations())
            if not ent_idlist:
                print 'Warning: attribute %s not associated to any model entities' % \
                    att.name()
                continue

            material_tuple = self.parse_thermal_conductivity(
                att, 'ConstantThermalConductivity', 'NonlinearMaterial')
            mat_type, mat_value = material_tuple
            mat_value_string = mat_value if mat_type == 'Function' else '%g' % mat_value
            mat_string = '    %s: %s\n' % (mat_type, mat_value_string)

            for ent_id in ent_idlist:
                self.start_command('ThermalConductivity')
                self.scope.output.write('    Id: %d\n' % ent_id)
                self.scope.output.write(mat_string)
                self.finish_command()

    def write_boundary_condition(self):
        ''''''
        print 'Writing boundary conditions'
        if self.is_elastic_solver():
            att_type = 'TEM3PMechanicalBC'
        elif self.is_thermal_solver():
            att_type = 'TEM3PThermalBC'
        else:
            raise Exception('Unsupported analysis type: %s' % self.solver)

        att_list = self.scope.sim_atts.findAttributes(att_type)
        if not att_list:
            return
        att_list.sort(key=lambda att:att.name())

        # Map nonstandard item names to their corresponding ace3p keyword
        keyword_table = {'LFDetuningMethod': 'Method'}

        for att in att_list:
            # Filter by cateogry
            if not att.isMemberOf(self.scope.categories):
                continue

            att_type = att.type()
            # Separate command for each model entity (requried?)
            ent_idlist = utils.get_entity_ids(self.scope, att.associations())
            if not ent_idlist:
                print 'Warning: attribute %s not associated to any model entities' % \
                    att.name()
                continue

            for ent_id in ent_idlist:
                self.start_command('Boundary')
                self.scope.output.write('    Id: %d\n' % ent_id)

                condition_type_lookup = {
                    'TEM3PStructuralNeumann': 'Neumann',
                    'TEM3PStructuralDirchlet': 'Dirichlet',
                    'TEM3PStructuralMixed': 'Mixed',
                    'TEM3PLFDetuning': 'LFDetuning',

                    'TEM3PHarmonicNeumann': 'Neumann',
                    'TEM3PHarmonicDirichlet': 'Dirichlet',
                    'TEM3PHarmonicMixed': 'Mixed',

                    'TEM3PThermalNeumann': 'Neumann',
                    'TEM3PThermalDirichlet': 'Dirichlet',
                    'TEM3PThermalRobin': 'Robin',
                    'TEM3PThermalRFHeating': 'RFHeating'
                }
                condition_type = condition_type_lookup.get(att_type)
                if condition_type is None:
                    raise Exception('Unrecognized BC attribute type %s' % att_type)

                self.scope.output.write('    ConditionType: %s\n' % condition_type)
                if condition_type == 'Mixed':
                    # Special handling for mixed type (space delimited type string)
                    type_item = att.findString('MixedType')
                    type_string = utils.format_vector(type_item, fmt='%s', separator=' ')
                    self.scope.output.write('    MixedType: %s\n' % type_string)
                    CardFormat('MixedValue').write(self.scope, att, indent='    ')
                else:
                    # Other types can use default handling
                    skip_list = ['PlaceHolder', 'NonlinearConvectiveSurface', 'SurfaceResistance']
                    self.write_standard_items(att, skip_list=skip_list, keyword_table=keyword_table)

                # Special cases:

                # For Eigen mode solver, must inject zero values
                if self.solver == 'tem3p-eigen':
                    if condition_type == 'Mixed':
                        self.scope.output.write('    MixedValue: 0., 0., 0.\n')
                    else:
                        self.scope.output.write('    %sValue: 0.\n' % condition_type)

                # Check for LFDetuning case, which requries NERSC directory
                if att.type() == 'TEM3PLFDetuning':
                    path_item = att.findString('NERSCDirectory')
                    remote_path = path_item.value(0).rstrip('/')
                    self.set_symlink(remote_path)
                    basename = os.path.basename(remote_path)
                    self.scope.output.write('    Directory: %s\n' % basename)

                # Check for RFHeating case, which uses Source and optional nonlinear properties
                elif (att.type() == 'TEM3PThermalRFHeating'):
                    path_item = att.findString('NERSCDirectory')
                    remote_path = path_item.value(0).rstrip('/')
                    self.set_symlink(remote_path)
                    basename = os.path.basename(remote_path)
                    self.scope.output.write('    Directory: %s\n' % basename)

                    surface_res_item = att.findString('SurfaceResistance')
                    if surface_res_item.isEnabled() and surface_res_item.isSet(0):
                        surface_res_value = surface_res_item.value(0)
                        if surface_res_value == 'None':
                            continue
                        elif surface_res_value == 'Custom':
                            sr_file_item = surface_res_item.findChild(
                                'SurfaceResistanceFile', smtk.attribute.ACTIVE_CHILDREN)
                            filename = self.use_file(sr_file_item)
                            if filename:
                                self.scope.output.write('    SurfaceResistance: %s\n' % filename)
                        else:
                            # User selected one of the standard SR files
                            source_dir = os.path.dirname(__file__)
                            relative_path = os.path.join(source_dir, os.pardir, 'non-linear-material', surface_res_value)
                            path = os.path.abspath(relative_path)
                            if not os.path.exists(path):
                                print 'Material file not found at:', path
                                raise Exception('Unabled to find standard materials file %s' % surface_res_value)
                            self.scope.output.write('    SurfaceResistance: %s\n' % surface_res_value)
                            self.scope.files_to_upload.add(path)

                conv_surface_item = att.itemAtPath('NonlinearConvectiveSurface', '/')
                if conv_surface_item and conv_surface_item.isEnabled():
                    factor_item = att.itemAtPath('NonlinearConvectiveSurface/RobinFactor', '/')
                    filename = self.use_file(factor_item)
                    if filename:
                        self.scope.output.write('    RobinFactor: %s\n' % filename)

                    value_item = att.itemAtPath('NonlinearConvectiveSurface/RobinValue', '/')
                    filename = self.use_file(value_item)
                    if filename:
                        self.scope.output.write('    RobinValue: %s\n' % filename)

                self.finish_command()


    def write_thermal_shell(self):
        '''Writes ThermalShell elements

        Uses brute force to assign material numbers
        '''
        print 'Writing shells'
        att_list = self.scope.sim_atts.findAttributes('ThermalShell')
        if not att_list:
            return
        att_list.sort(key=lambda att:att.name())

        # For this implementation, assign unique material id to each shell.
        # Number starting at 101
        new_material_id = 101
        material_dict = dict()  # kv == <(type, value), material-id>
        # Traverse atts and write out shells
        for att in att_list:
            # Separate command for each model entity (requried?)
            ent_idlist = utils.get_entity_ids(self.scope, att.associations())
            if not ent_idlist:
                print 'Warning: attribute %s not associated to any model entities' % \
                    att.name()
                continue

            material_tuple = self.parse_thermal_conductivity(
                att, 'ConstantThermalConductivity', 'GeneralThermalConductivity')
            print 'material_tuple:', material_tuple
            material_id = material_dict.get(material_tuple)
            if material_id is None:
                material_id = new_material_id
                new_material_id += 1
                material_dict[material_tuple] = material_id

            # Write shell sections
            for ent_id in ent_idlist:
                self.start_command('Shell')
                self.scope.output.write('    Bd: %d\n' % ent_id)
                self.scope.output.write('    Material: %d\n' % material_id)
                skip_list= ['ConstantThermalConductivity', 'GeneralThermalConductivity']
                self.write_standard_items(att, skip_list=skip_list)
                self.finish_command()

        # Traverse material_dict and write out "material" sections
        for material_tuple in sorted(material_dict.iterkeys()):
            print 'mat_tuple', material_tuple
            mat_type, mat_value = material_tuple
            material_id = material_dict.get(material_tuple)
            mat_value_string = mat_value if mat_type == 'Function' else '%g' % mat_value

            self.start_command('ThermalConductivity')
            self.scope.output.write('    Id: %d\n' % material_id)
            self.scope.output.write('    %s: %s\n' % (mat_type, mat_value_string))
            self.finish_command()

    def write_heat_source(self):
        ''''''
        print 'Writing heat sources'
        att_list = self.scope.sim_atts.findAttributes('HeatSource')
        if not att_list:
            return

        att_list.sort(key=lambda att:att.name())
        for att in att_list:
            # Separate command for each model entity (requried?)
            ent_idlist = utils.get_entity_ids(self.scope, att.associations())
            if not ent_idlist:
                print 'Warning: attribute %s not associated to any model entities' % \
                    att.name()
                continue

            for ent_id in ent_idlist:
                self.start_command('HeatSource')
                self.scope.output.write('    Id: %s\n' % ent_id)
                self.write_standard_items(att, skip_list=['NERSCDirectory', 'ExtVHeatingFile'])

                # NERSC directory for LossyDielectricHeat atts
                path_item = att.findString('NERSCDirectory')
                if path_item is not None:
                    remote_path = path_item.value(0).rstrip('/')
                    self.set_symlink(remote_path)
                    basename = os.path.basename(remote_path)
                    self.scope.output.write('    Directory: %s\n' % basename)

                # ExtVHeatingFile for ExternalVolumeHeating atts
                path_item = att.findString('ExtVHeatingFile')
                if path_item is not None:
                    if not path_item.isSet(0):
                        raise Exception('File not set for Heat Source %s' % att.name())
                    local_path = path_item.value(0)
                    if not os.path.exists(local_path):
                        msg = 'File \"%s\"" not found for Heat Source %s' % (local_path, att.name())
                        raise Exception(msg)

                    basename = os.path.basename(local_path)
                    self.scope.output.write('    ExtVHeatingFile: %s\n' % basename)
                    # Include file in upload to NERSC
                    self.scope.folders_to_upload.add(local_path)


                self.finish_command()


    def write_harmonic_analysis(self):
        print 'Writing harmonic analysis'
        self.write_standard_instance_att('HarmonicAnalysis', silent_list=['Frequency'])


    def parse_thermal_conductivity(self, att, linear_item_name, general_item_name):
        '''Parses the different ways for specifying thermal conductivity

        Returns tuple(type, value) as either ("Value", float) or ("Function", string)
        where the string is a filename.
        '''
        tc_item = None
        if self.solver == 'tem3p-thermal-linear':
            tc_item = att.findDouble(linear_item_name)
            if not tc_item.isSet(0) or not tc_item.isValid():
                raise Exception('Error: thermal conductivity value not valid')
            return ('Value', tc_item.value(0))
        elif self.solver == 'tem3p-thermal-nonlinear':
            tc_item = att.findString(general_item_name)
        if tc_item is None:
            raise Exception('Internal error: Missing thermal conductivity item')

        if not tc_item.isSet(0):
            raise Exception('Thermal material item not set for material attribute: %s' % att.name())

        tc_choice = tc_item.value(0)
        if tc_choice == 'CustomConstant':
            value_item = tc_item.findChild(
                'ConstantThermalConductivity', smtk.attribute.ACTIVE_CHILDREN)
            if not value_item.isSet(0):
                raise Exception('Thermal conductivity value not set for material attribute %s' % att.name())
            #CardFormat('Value').write_item(self.scope, value_item, indent)
            return('Value', value_item.value(0))

        elif tc_choice == 'CustomNonlinear':
            tc_file_item = tc_item.findChild(
                'NonlinearThermalConductivity', smtk.attribute.ACTIVE_CHILDREN)
            if not tc_file_item.isSet(0):
                raise Exception('Thermal conductivity value not set for material attribute %s' % att.name())
            filename = self.use_file(tc_file_item)
            if filename:
                #self.scope.output.write('    Function: %s\n' % filename)
                return('Function', filename)

        else:
            # User has selected a standard nonlinear material in our makeshift library
            source_dir = os.path.dirname(__file__)
            relative_path = os.path.join(source_dir, os.pardir, 'non-linear-material', tc_choice)
            path = os.path.abspath(relative_path)
            if not os.path.exists(path):
                print 'Material file not found at:', path
                raise Exception('Unabled to find standard materials file %s' % tc_choice)
            return('Function', tc_choice)

        # (else)
        raise Exception('Internal error in parse_thermal_conductivity()')


    def set_symlink(self, remote_path):
        '''If symlink already set, checks that remote_path is consistent with it

        '''
        if self.scope.symlink is None:
            print 'setting scope.symlink: %s' % remote_path
            self.scope.symlink = remote_path
            return True

        # (else)
        if remote_path != self.scope.symlink:
            template_string = '''Remote paths dont match:
current symlink: $symlink
remote path:     $remote_path
'''
            tpl = string.Template(template_string)
            msg = tpl.substitute(symlink=self.scope.symlink, remote_path=remote_path)
            raise Exception(msg)
