'''Writer for Track3P input files
'''

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

        # 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')

        nersc_item = att.itemAtPath('NERSCDirectory')
        remote_path = nersc_item.value(0).rstrip('/')
        if remote_path == '':
            raise Exception('Field Directory not specified in Track3P Analysis panel ')
        field_dir = os.path.basename(remote_path)
        self.scope.symlink = remote_path

        self.scope.output.write('  FieldDir: %s\n' % field_dir)

        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)
        self.finish_command()

        # For rf windows
        solid_region_item = att.findModelEntity('SolidRegion')
        vacuum_region_item = att.findModelEntity('VacuumRegion')
        if solid_region_item.isEnabled() or vacuum_region_item.isEnabled():
            self.start_command('Domain')
            self.scope.output.write('  FieldDir: %s\n' % field_dir)
            if solid_region_item.isEnabled():
                id_list = utils.get_entity_ids(solid_region_item)
                if id_list:
                    id_string = ','.join(str(id) for id in id_list)
                    self.scope.output.write('  SolidRegion: %s' % id_string)
            if vacuum_region_item.isEnabled():
                id_list = utils.get_entity_ids(vacuum_region_item)
                if id_list:
                    id_string = ','.join(str(id) for id in id_list)
                    self.scope.output.write('  VacuumRegion: %s' % id_string)
            self.finish_command()

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


        # 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('Domain')
            self.scope.output.write('  FieldDir: %s\n' % field_dir)

            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
        '''
        # Emitters are assigned to model entities via the material attributes
        mat_list = self.scope.sim_atts.findAttributes('Track3PMaterial')
        for mat_att in mat_list:
            emitter_item = mat_att.find('Emitter')
            if emitter_item is None or not emitter_item.isSet(0):
                continue

            # Get model entities associated to this material att
            entity_string = utils.format_entity_string(self.scope, mat_att)
            if entity_string == '':
                continue

            # (else)
            att = emitter_item.value(0)
            self.start_command('Emitter')
            args = (self.scope, att)  # for convenience/reuse below
            CardFormat('Type', 'Type', '%d').write(*args)
            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))

            self.scope.output.write('  BoundaryID: %s\n' % entity_string)

            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
        '''
        # Selector att indicates simulation as either Multipacting or DarkCurrent
        selector_att = self.get_attribute('Track3PModelSelectorAtt')
        selector_item = selector_att.findString('Track3PModelSelectorItem')
        track3p_sim = selector_item.value(0)
        print ('%s simulation' % track3p_sim)

        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:
            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_ref = model_ent_item.value(i)
                ent = ent_ref.entity()
                prop_idlist = self.scope.model_resource.integerProperty(ent, 'pedigree id')
                #print 'idlist', idlist
                if prop_idlist:
                    #scope.output.write(' %d' % idlist[0])
                    pedigree_id = prop_idlist[0]
                    dim = ent_ref.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.iteritems():
                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.iteritems():
                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:
                        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
                self.scope.logger.addError(msg)
                raise Exception(msg)

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

        if sey_model_number == 2:
            args = (self.scope, att)  # for convenience/reuse
            CardFormat('MinimumNumElectrons', fmt='%d').write(*args)
            CardFormat('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:
                    print 'No BoundarySurface entity'
                    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))
            self.finish_command('  ')

        self.finish_command()  # Postprocess
