'''Writer for Track3P input files
'''

from imp import reload
import os

import smtk
from smtk import attribute

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

class Track3PWriter(basewriter.BaseWriter):
    '''
    '''
    def __init__(self):
        super(Track3PWriter, self).__init__()
        self.base_indent = ''

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

        # Write total time
        self.write_inline_value('Domain', item_path='TotalTime')
        self.write_field_scales()
        self.write_domain()
        self.write_inline_value('Domain', item_path='OutputImpacts')
        self.write_emitters()
        self.write_materials()
        # self.write_inline_value('SingleParticleTrajectory')
        self.write_postprocess()

    def write_field_scales(self):
        '''Writes FieldScales, NormalizedField, ParticlesTrajectories commands

        '''
        att = self.get_attribute('FieldScales')
        if not att:
            return

        args = (self.scope, att)
        self.start_command('FieldScales')

        CardFormat('Type', 'Type', '%s').write(*args)
        CardFormat('ScanToken', 'ScanToken', '%d').write(*args)
        scantoken_item = att.findInt('ScanToken')
        #print '\nscantoken_item.value(0): %s\n\n' % scantoken_item.value(0)
        if scantoken_item.value(0) == 0:
            CardFormat('Scale', 'ScanToken/Scale', '%e').write(*args)
        elif scantoken_item.value(0) == 1:
            CardFormat('Minimum', 'ScanToken/Minimum', '%e').write(*args)
            CardFormat('Maximum', 'ScanToken/Maximum', '%e').write(*args)
            CardFormat('Interval', 'ScanToken/Interval', '%e').write(*args)

        self.finish_command()

        # Write NormalizedField command when FieldScales type is FieldGradient
        type_item = att.findString('Type')
        if type_item.value(0) == 'FieldGradient':
            self.start_command('NormalizedField')
            CardFormat('Beta', 'Type/Beta', '%d').write(*args)
            CardFormat('StartPoint', 'Type/StartPoint').write(*args)
            CardFormat('EndPoint', 'Type/EndPoint').write(*args)
            CardFormat('N', 'Type/N').write(*args)
            self.finish_command()

        # Write ParticlesTrajectores when scantoken is 0
        if scantoken_item.value(0) == 0:
            particles_att = self.get_attribute('ParticleTrajectories')
            pargs = (self.scope, particles_att)
            self.start_command('ParticlesTrajectores')
            CardFormat('ParticleFile', 'ParticleFile', '%s').write(*pargs)
            CardFormat('Start', 'Start', '%d').write(*pargs)
            CardFormat('Stop', 'Stop', '%d').write(*pargs)
            CardFormat('Skip', 'Skip', '%d').write(*pargs)
            self.finish_command()

    def write_domain(self):
        '''Writes Domain command
        '''
        att = self.get_attribute('Domain')
        if not att:
            return

        args = (self.scope, att)  # for convenience/reuse below

        # Baseline domain
        self.start_command('Domain')

        # Field data is specified in InputData attribute
        inputdata_att = self.scope.sim_atts.findAttribute('InputData')
        CardFormat('FieldDir', 'Source/NERSCDirectory').write(self.scope, inputdata_att)

        # Rest of data is specified in the Domain attribute
        CardFormat('FrequencyScanID', 'FrequencyScanID', '%d').write(*args)
        CardFormat('ModeID1', 'Mode', '%d').write(*args)
        CardFormat('dt', 'dt', '%f').write(*args)
        CardFormat('MaxImpacts', 'MaxImpacts', '%d').write(*args)
        CardFormat('LowEnergy', 'LowEnergy', '%f').write(*args)
        CardFormat('HighEnergy', 'HighEnergy', '%f').write(*args)
        CardFormat('InitialEnergy', 'InitialEnergy', '%f').write(*args)
        CardFormat('Bins', 'Bins', '%d').write(*args)
        #CardFormat('Voltage', 'Voltage').write(*args)

        # For rf windows
        for i in range(att.numberOfItems()):
            item = att.item(i)
            # print('\"{}\"'.format(item.name()))
            if item.name() == 'SolidRegion':
                item_find = att.find('SolidRegion')
                # print(item.name(), item, item_find)

        solid_region_item = att.findComponent('SolidRegion')
        vacuum_region_item = att.findComponent('VacuumRegion')
        if solid_region_item.isEnabled() or vacuum_region_item.isEnabled():
            if solid_region_item.isEnabled():
                id_list = utils.get_entity_ids(self.scope, solid_region_item)
                if id_list:
                    id_string = ','.join(str(id) for id in id_list)
                    self.scope.output.write('  SolidRegion: {}\n'.format(id_string))
            if vacuum_region_item.isEnabled():
                id_list = utils.get_entity_ids(self.scope, vacuum_region_item)
                if id_list:
                    id_string = ','.join(str(id) for id in id_list)
                    self.scope.output.write('  VacuumRegion: {}\n'.format(id_string))

        # For uniform external magnetic field
        external_field_item = att.findDouble('ExternalMagneticField')
        if external_field_item.isEnabled():
            value_string = utils.format_vector(external_field_item)
            self.scope.output.write('  ExternalMagneticField: %s\n' % value_string)

        # For external magnetic field map
        external_map_item = att.findGroup('MagneticFieldMap')
        if external_map_item.isEnabled():
            # Verify that map file exists
            file_item = external_map_item.find('File')
            filename = self.use_file(file_item, 'MagneticFieldMap')

            self.start_command('MagneticFieldMap', indent='  ', blank_line=False)
            self.scope.output.write('    File: %s\n' % filename)
            kwargs = {'indent': '    '}
            CardFormat('Scaling', 'MagneticFieldMap/Scaling').write(*args, **kwargs)
            CardFormat('Units', 'MagneticFieldMap/Units').write(*args, **kwargs)
            CardFormat('ZOffset', 'MagneticFieldMap/ZOffset').write(*args, **kwargs)

            self.finish_command('  ')  # MagneticFieldMap

        self.finish_command()  # Domain

    def write_inline_value(self, att_type, item_path=None, name=None):
        '''Writes inline command consisting of single value
        '''
        att = self.get_attribute(att_type)
        if not att:
            print('Did not find attribute type %s' % att_type)
            return

        if item_path is None:
            item_path = att_type
        item = att.itemAtPath(item_path)
        if not item:
            print('Did not find item  at path \"%s\"" under attribute %s' % \
                (att_type, item_path))
            return

        if item.type() == smtk.attribute.Item.VoidType:
            value = 'on' if item.isEnabled() else 'off'
        else:
            value = item.value()
        self.scope.output.write('\n')

        # Make sure name is defined
        if name is None and item_path is None:
            name = att_type
        elif name is None:
            name = item_path
        self.scope.output.write('%s: %s\n' % (name, value))

    def write_emitters(self):
        '''Writes Emitter commands
        '''
        # First create a set(int) for all model entities that are pimary materials
        # This will be used to check that emitters are only associated to those entities
        mat_ids = set()
        mat_list = self.scope.sim_atts.findAttributes('Track3PMaterial')
        for mat in mat_list:
            if mat.type() in ['Track3PMaterialMultipactingPrimary', 'Track3PMaterialDarkCurrentPrimary']:
                mat_ids |= set(utils.get_entity_ids(self.scope, mat.associations()))

        # Write all Emitter attributes
        att_list = self.scope.sim_atts.findAttributes('Emitter')
        for att in att_list:
            # Check model associations
            ent_ids = utils.get_entity_ids(self.scope, att.associations())
            if not ent_ids:
                msg = 'Emitter attribute {} not associated to any model entities -- skipping'.format(att.name())
                print(msg)
                self.scope.warning_messages.append(msg)
                continue

            # Check if all entities are primary surfaces
            invalid_ids = set(ent_ids) - mat_ids
            if invalid_ids:
                msg = 'Emitter attribute {} associated with some model entities that are not primary: {}' \
                    .format(att.name, invalid_ids)
                valid_ids = set(ent_ids) | mat_ids
                ent_ids = list(valid_ids)
                if not ent_ids:
                    continue

            entity_list = [str(id) for id in ent_ids]
            entity_string = ', '.join(entity_list)

            self.start_command('Emitter')

            args = (self.scope, att)  # for convenience/reuse below
            CardFormat('Type', 'Type', '%d').write(*args)
            self.scope.output.write('  BoundaryID: %s\n' % entity_string)
            CardFormat('t0', 'Start', '%d').write(*args)
            CardFormat('t1', 'Stop', '%d').write(*args)

            # Bounds
            min_item = att.findDouble('BoundsMin')
            max_item = att.findDouble('BoundsMax')
            self.scope.output.write('  x0: %g\n' % min_item.value(0))
            self.scope.output.write('  x1: %g\n' % max_item.value(0))
            self.scope.output.write('  y0: %g\n' % min_item.value(1))
            self.scope.output.write('  y1: %g\n' % max_item.value(1))
            self.scope.output.write('  z0: %g\n' % min_item.value(2))
            self.scope.output.write('  z1: %g\n' % max_item.value(2))

            CardFormat('SkipTimeSteps', 'Skip', '%d').write(*args)

            type_item = att.findInt('Type')
            if type_item.value(0)  == 7:  # FieldEmission
                CardFormat('N', item_path='Type/N', fmt='%d').write(*args)
                CardFormat('WorkFunction', item_path='Type/WorkFunction', fmt='%s').write(*args)
                CardFormat('Beta', item_path='Type/Beta', fmt='%s').write(*args)

            self.finish_command()

    def write_materials(self):
        '''Writes Material commands
        '''
        # Get Track3P simulation mode
        track3p_sim = None
        if 'Track3P-Multipacting' in self.scope.categories:
            track3p_sim = 'Multipacting'
        elif 'Track3P-DarkCurrent' in self.scope.categories:
            track3p_sim = 'DarkCurrent'
        else:
            self.scope.warning_messages.append('Unrecognized Track3P simulation type')
            return

        att_list = self.scope.sim_atts.findAttributes('Track3PMaterial')

        # Because the same material attribute is used for surfaces and volumes,
        # we need to separate them here so that we can sort and write separately.
        # Using 2 dictionaries to map <attribute, [entity id]>
        surface_dict = dict()
        volume_dict = dict()

        for att in att_list:
            if not utils.passes_categories(att, self.scope.categories):
                continue

            model_ent_item = att.associations()
            assert model_ent_item is not None

            for i in range(model_ent_item.numberOfValues()):
                if not model_ent_item.isSet(i):
                    continue

                ent = model_ent_item.value(i)
                prop_idlist = self.scope.model_resource.integerProperty(ent.id(), 'pedigree id')
                #print 'idlist', idlist
                if prop_idlist:
                    #scope.output.write(' %d' % idlist[0])
                    pedigree_id = prop_idlist[0]
                    dim = ent.dimension()
                    if dim == 2:
                        surface_list = surface_dict.get(att, [])
                        surface_list.append(pedigree_id)
                        surface_dict[att] = surface_list
                    elif dim == 3:
                        volume_list = volume_dict.get(att, [])
                        volume_list.append(pedigree_id)
                        volume_dict[att] = volume_list
                    else:
                        raise Exception('Invalid entity dimension %s' % s)

        # print '{}'.format(surface_dict)
        # print '{}'.format(volume_dict)
        for d,keyword in [(volume_dict,'SolidBlockID'), (surface_dict,'BoundarySurfaceID')]:
            # Convert id list to string and build new dict of <string, att>
            sorted_dict = dict()
            for att, idlist in d.items():
                idlist.sort()
                string_list = [str(i) for i in idlist]
                id_string = ', '.join(string_list)
                sorted_dict[id_string] = att
            #print '{}'.format(sorted_dict)

            # Traverse sorted_dict and write materials
            for id_string, att in sorted_dict.items():
                if 'Primary' in att.type():
                    self.start_command('Material')
                    self.scope.output.write('  Type: Primary\n')
                    self.scope.output.write('  %s: %s\n' % (keyword, id_string))
                    self.finish_command()

                    # Primary type can include optional secondary item
                    secondary_item = att.find('Secondary')
                    if secondary_item is not None and secondary_item.isEnabled():
                        self.start_command('Material')
                        self.scope.output.write('  Type: Secondary\n')
                        self.scope.output.write('  %s: %s\n' % (keyword, id_string))
                        self._write_secondary_material_items(att, track3p_sim)
                        self.finish_command()

                elif 'Secondary' in att.type():
                    self.start_command('Material')
                    self.scope.output.write('  Type: Secondary\n')
                    self.scope.output.write('  %s: %s\n' % (keyword, id_string))
                    self._write_secondary_material_items(att, track3p_sim)
                    self.finish_command()

                elif 'Absorber' in att.type():
                    self.start_command('Material')
                    self.scope.output.write('  Type: Absorber\n')
                    self.scope.output.write('  %s: %s\n' % (keyword, id_string))
                    self.finish_command()

                else:
                    raise Exception('Unrecognized material type %s' % att.type())

    def _write_secondary_material_items(self, att, track3p_sim):
        '''Offloads some of the logic from write_materials()

        '''
        if track3p_sim == 'Multipacting':
            sey_model_number = 1
        elif track3p_sim == 'DarkCurrent':
            sey_model_number = 2
            pivi_item = att.findVoid('FurnhamPivi')
            if pivi_item and pivi_item.isEnabled():
                sey_model_number = 3
        else:
            raise Exception('Unrecognized track3p_sim value: %s' % track3p_sim)
        self.scope.output.write('  Model: %s\n' % sey_model_number)

        sey_item = att.findFile('SEYFilename')
        if sey_item and sey_item.isSet(0):
            local_path = sey_item.value(0)
            # Check that file exists
            if not os.path.exists(local_path):
                msg = 'ERROR: Local data file not found at %s' % local_path
                print(msg)
                self.scope.warning_messages.append(msg)

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

        if sey_model_number == 2:
            args = (self.scope, att)  # for convenience/reuse
            CardFormat('MinimumNumElectrons', item_path='Secondary/MinimumNumElectrons', fmt='%d').write(*args)
            CardFormat('ElasticThreshold', item_path='Secondary/ElasticThreshold', fmt='%d').write(*args)

    def write_postprocess(self):
        '''Writes Postprocess command
        '''
        att = self.get_attribute('Postprocess')
        if not att:
            return

        self.start_command('Postprocess')
        top_toggle_item = att.findGroup('Toggle')
        top_toggle_value = 'on' if top_toggle_item.isEnabled() else 'off'
        self.scope.output.write('  Toggle: %s\n' % top_toggle_value)

        if top_toggle_item.isEnabled():
            args = (self.scope, att)  # for convenience/reuse below
            kwargs = {'indent': '    '}

            # Write ResonantParticles command
            self.start_command('ResonantParticles', '  ')
            res_part_item = top_toggle_item.find('ResonantParticles')
            res_part_value = 'on' if res_part_item.isEnabled() else 'off'
            self.scope.output.write('    Token: %s\n' % res_part_value)
            if res_part_item.isEnabled():
                CardFormat('InitialImpacts',
                    'Toggle/ResonantParticles/InitialImpacts', '%d').write(*args, **kwargs)
                CardFormat('EnergyRange',
                    'Toggle/ResonantParticles/EnergyRange').write(*args, **kwargs)
                CardFormat('PhaseTolerance',
                    'Toggle/ResonantParticles/PhaseTolerance').write(*args, **kwargs)
            self.finish_command('  ')

            # Write EnhancementCounter command
            self.start_command('EnhancementCounter', '  ')
            enh_ctr_item = top_toggle_item.find('EnhancementCounter')
            enh_ctr_value = 'on' if enh_ctr_item.isEnabled() else 'off'
            self.scope.output.write('    Token: %s\n' % enh_ctr_value)
            CardFormat('MinimumEC', 'Toggle/EnhancementCounter/MinimumEC', '%d') \
                .write(*args, **kwargs)

            for i,name in enumerate(['SEYSurface1', 'SEYSurface2']):
                index = i+1
                sey_item = enh_ctr_item.find(name)
                if not sey_item.isEnabled():
                    break

                boundary_item = sey_item.find('BoundarySurface')
                id_list = utils.get_entity_ids(self.scope, boundary_item)
                if not id_list:
                    msg = 'No BoundarySurface entity defined for {}'.format(name)
                    self.scope.warning_messages.append(msg)
                    break
                id_string = ','.join(str(id) for id in id_list)
                self.scope.output.write('    BoundarySurfaceID%d: %s\n' % \
                    (index, id_string))

                file_keyword = 'SEYFileName%d' % index
                file_item = sey_item.find('SEYFilename')
                filename = self.use_file(file_item, name='SEYFileName')
                self.scope.output.write('    %s: %s\n' % (file_keyword, filename))

            # WindowVolume
            volume_item = enh_ctr_item.find('WindowVolume')
            if volume_item is not None and volume_item.isEnabled():
                pass
                id_list = utils.get_entity_ids(self.scope, volume_item)
                if id_list:
                    id_string = ','.join(str(id) for id in id_list)
                    self.scope.output.write('    SolidVoidID: {}\n'.format(id_string))
                else:
                    self.scope.warning_messages.append(
                        'Window Volume enabled but not assigned to any model entity')

            self.finish_command('  ')

            # Write FaradayCup
            fcup_item = top_toggle_item.find('FaradayCup')
            if fcup_item is not None and fcup_item.isEnabled():
                self.start_command('FaradayCup', '  ')
                self.scope.output.write('    Token: on\n')

                group_item = fcup_item.find('BoundarySurfaceGroup')
                id_list = self.get_entity_ids_from_group(group_item, 'BoundarySurface')
                if id_list:
                    id_string = ','.join(str(id) for id in id_list)
                    self.scope.output.write('    BoundaryID: %s\n' % id_string)
                else:
                    self.scope.warning_messages.append('FaradayCup specified with no boundary surfaces')

                self.finish_command('  ')  # FaradayCup

        self.finish_command()  # Postprocess

    def get_entity_ids_from_group(self, group_item, component_item_name):
        """Returns list of pedigree ids for extensible GroupItem containing ComponentItem

        Args:
            group_item: smtk.attribute.GroupItem, extensible
            component_item_name: string, name for subgroup smtk.attribute.ComponentItem

        """
        id_set = set()  # use set to avoid duplication

        for i in range(group_item.numberOfGroups()):
            comp_item = group_item.find(i, component_item_name)
            if not comp_item.isSet():
                continue

            ent = comp_item.value()
            prop_idlist = self.scope.model_resource.integerProperty(ent.id(), 'pedigree id')
            if prop_idlist:
                id_set.add(prop_idlist[0])

        return list(id_set)
