"""Writer for TEM3P input files."""
import csv
from imp import reload
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):
        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

    def is_elastic_solver(self):
        return 'TEM3P-Elastic' in self.scope.categories or 'TEM3P-Eigen' in self.scope.categories

    def is_thermal_solver(self):
        return 'TEM3P-Thermal' in self.scope.categories

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

        if self.is_elastic_solver():
            # 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



        start_categories = scope.categories.copy()

        # Check if thermoelastic analysis
        if set(['TEM3P-Elastic', 'TEM3P-Thermal']).issubset(scope.categories):
            self.start_command('ThermoElasticProblem', indent='')
            self.scope.output.write('  RunId: 3\n')
            self.finish_command(indent='')

            # Write thermal section
            scope.categories = start_categories - set(['TEM3P-Elastic'])
            print('Using categories: %s' % scope.categories)
            self.name_prefix = 'Thermal'
            self.write_command(is_thermoelastic=True)

            # Write elastic section
            scope.categories = set(['TEM3P-Elastic'])
            self.name_prefix = 'Elastic'
            self.write_command(is_thermoelastic=True)

            # Restore scope.categories
            self.name_prefix = None
            scope.categories = start_categories
        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 RuntimeError(
                'Unrecognized TEM3P solver for category set {}'.format(scope.categories))


    def write_command(self, indent='', is_thermoelastic=False):
        """Internal method to write either ElasticProblem or ThermostaticProblem."""
        if self.is_elastic_solver():
            self.start_command('ElasticProblem', indent='')
        elif self.is_thermal_solver():
            self.start_command('ThermostaticProblem', indent='')
        else:
            raise Exception('Unable to determine solver type')
        # 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_type = '{}Order'.format(self.name_prefix)
        order_att = self.get_attribute(order_type)
        args = (self.scope, order_att)
        CardFormat('BasisOrder').write(*args)
        CardFormat('CurvedSurfaces').write(*args)

        # Special case for ReferenceT
        if is_thermoelastic and self.name_prefix =='Elastic':
            temp_item = order_att.findDouble('ReferenceT')
            if temp_item is not None and temp_item.isSet(0):
                value = temp_item.value(0)
                # Todo replace next 3 lines with CardFormat.write_value() (class method)
                value_string = '%g' % temp_item.value(0)
                output_string = '%s%s: %s\n' % ('  ', 'ReferenceT', value_string)
                self.scope.output.write(output_string)

        # 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 'TEM3P-Harmonic' in self.scope.categories:
            self.write_harmonic_analysis()

        self.finish_command(indent='')  # Problem


    def write_meshdump(self):
        """Writes the meshdump section."""
        if not self.is_elastic_solver():
            return

        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():
            if self.scope.nersc_directory is None:
                self.scope.warning_messages.append('Writing WriteDeformedEMMesh with no NERSC Directory specified')
            else:
                self.scope.output.write('    EMMeshInputDir: %s\n' % self.scope.nersc_directory)

        self.finish_command(indent='  ')  # MeshDump


    def write_solver(self):
        """Writes the solver specification(s)."""
        print('Writing {} solver'.format(self.name_prefix))
        if 'TEM3P-Eigen' in self.scope.categories:
            self.write_standard_instance_att('EigenSolver')
        else:
            # All non-eigen cases use linear solver
            att_type = '{}LinearSolver'.format(self.name_prefix)
            att = self.get_attribute(att_type)
            self.start_command('LinearSolver')
            keyword_map = {'Type': 'Solver'}
            self.write_standard_items(att, keyword_table=keyword_map)
            self.finish_command()

        # In addition...
        if 'TEM3P-Harmonic' in self.scope.categories:
            self.write_standard_instance_att('HarmonicAnalysis')

        if 'TEM3P-Nonlinear-Thermal' in self.scope.categories:
            keyword_map = {'Type': 'Solver'}
            self.write_standard_instance_att('NonlinearSolver', keyword_table=keyword_map)
            self.write_standard_instance_att('PicardSolver')


    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:
                    is_eigenmode = 'TEM3P-Eigen' in self.scope.categories
                    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('    PoissonsRatio: %s\n' % poisson)
                    self.scope.output.write('    YoungsModulus: %s\n' % youngs)
                    if density is not None and is_eigenmode:
                        self.scope.output.write('    Density: %s\n' % density)
                    if elastic_alpha is not None and not is_eigenmode:
                        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())

        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')
            if material_tuple is not None:
                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 RuntimeError(
                'Unrecognized TEM3P solver for category set {}'.format(self.scope.categories))
        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 utils.passes_categories(att, 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',
                    'TEM3PStructuralDirichlet': '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 = [
                        'NERSCDirectory',
                        'NonlinearConvectiveSurface',
                        'PlaceHolder',
                        '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 'TEM3P-Eigen' in self.scope.categories:
                    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':
                    if self.scope.nersc_directory is None:
                        self.scope.warning_messages.append('Writing LFDetuning with no NERSC Directory specified')
                    else:
                        self.scope.output.write('    Directory: %s\n' % self.scope.nersc_directory)

                # Check for RFHeating case, which uses Source and optional nonlinear properties
                elif (att.type() == 'TEM3PThermalRFHeating'):
                    if self.scope.nersc_directory is None:
                        self.scope.warning_messages.append('Writing RFHeating case with no NERSC Directory specified')
                    else:
                        self.scope.output.write('    Directory: %s\n' % self.scope.nersc_directory)

                    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':
                            pass
                        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)
                                msg = 'Unable to find standard materials file %s' % surface_res_value
                                print(msg)
                                self.scope.warning_messages.append(msg)
                                continue
                            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:
                msg = 'Warning: thermal shell attribute %s not associated to any model entities' % \
                    att.name()
                print(msg)
                self.scope.warning_messages.append(msg)
                continue

            material_tuple = self.parse_thermal_conductivity(
                att, 'ConstantThermalConductivity', 'GeneralThermalConductivity')
            if material_tuple is not None:
                # 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.keys()):
            # 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:
                msg = 'Warning: heat source attribute %s not associated to any model entities' % \
                    att.name()
                print(msg)
                self.scope.warning_messages.append(msg)
                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=['ExtVHeatingFile'])
                if self.scope.nersc_directory is None:
                    self.scope.warning_messages.append('Writing HeatSource with no NERSC Directory specified')
                else:
                    self.scope.output.write('    Directory: %s\n' % self.scope.nersc_directory)

                # ExtVHeatingFile for ExternalVolumeHeating atts
                path_item = att.findFile('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.files_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; or returns None on error
        """
        tc_item = None
        if 'TEM3P-Nonlinear-Thermal' in self.scope.categories:
            tc_item = att.findString(general_item_name)
        elif 'TEM3P-Linear-Thermal' in self.scope.categories:
            tc_item = att.findDouble(linear_item_name)
            if not tc_item.isSet(0) or not tc_item.isValid():
                msg = 'Error: thermal conductivity value not valid'
                print(msg)
                self.scope.warning_messages.append(msg)
                return None
            return ('Value', tc_item.value(0))
        if tc_item is None:
            msg = 'Internal error: Missing thermal conductivity item'
            print(msg)
            self.scope.warning_messages.append(msg)
            return None

        if not tc_item.isSet(0):
            msg = 'Thermal material item not set for material attribute: %s' % att.name()
            print(msg)
            self.scope.warning_messages.append(msg)
            return None

        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):
                msg = 'Thermal conductivity value not set for material attribute %s' % att.name()
                print(msg)
                self.scope.warning_messages.append(msg)
                return None
            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):
                msg = 'Thermal conductivity value not set for material attribute %s' % att.name()
                print(msg)
                self.scope.warning_messages.append(msg)
                return None
            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)
                msg = 'Unabled to find standard materials file %s' % tc_choice
                print(msg)
                self.scope.warning_messages.append(msg)
                return None
            self.scope.files_to_upload.add(path)
            return('Function', tc_choice)

        # (else)
        msg = 'Internal error in parse_thermal_conductivity()'
        print(msg)
        self.scope.warning_messages.append(msg)
        return None
