'''Writer for Track3P input files
'''

import os

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__()

    def write(self, scope):
        '''
        '''
        self.scope = scope
        self.write_total_time()
        self.write_field_scales()
        self.write_domain()
        self.write_boolean_command('OutputImpacts')
        self.write_emitters()
        self.write_materials()
        self.write_boolean_command('SingleParticleTrajectory')
        self.write_postprocess()

    def write_total_time(self):
        '''Writes TotalTime command
        '''
        att = self.get_attribute('Time')
        if att:
            self.scope.output.write('\n')
            CardFormat('TotalTime', 'TotalTime', '%d').write(self.scope, att, '')

    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_dir = 'unknown'
        # The FieldDir input is one of 4 cases
        #   1. Run Omega3P as prerequisite
        #   2. Run S3P as prerequisite
        #   3. Upload folder from local file system
        #   4. Use folder already in NERSC system
        if self.scope.domain_source == 'Omega3P':
            field_dir = 'omega3p_results'
        elif self.scope.domain_source == 'S3P':
            field_dir = 's3p_results'
        elif self.scope.domain_source == 'LocalDirectory':
            local_item = att.itemAtPath('InputData/Source/LocalDirectory')
            local_path = local_item.value(0).rstrip('/')
            # Check that folder exists
            if not os.path.exists(local_path):
                msg = 'ERROR: Input data directory not found at %s' % local_path
                self.scope.logger.addError(msg)
                raise Exception(msg)

            field_dir = os.path.basename(local_path)
            # Add flag to upload folder to NERSC
            self.scope.folders_to_upload.add(local_path)
        elif self.scope.domain_source == 'NERSCDirectory':
            nersc_item = att.itemAtPath('InputData/Source/NERSCDirectory')
            remote_path = nersc_item.value(0).rstrip('/')
            field_dir = os.path.basename(remote_path)
            self.scope.symlink = remote_path
        else:
            msg = 'Unrecognized domain source: %s' % self.scope.domain_source
            self.scope.logger.addError(msg)
            raise Exception(msg)

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

        CardFormat('ModeID1', 'InputData/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_boolean_command(self, name):
        '''Writes inline command consisting of single boolean (VoidType)
        '''
        att = self.get_attribute(name)
        if not att:
            return

        item = att.findVoid(name)
        if not item:
            return

        value = 'on' if item.isEnabled() else 'off'
        self.scope.output.write('\n')
        self.scope.output.write('%s: %s\n' % (name, value))

    def write_emitters(self):
        '''Writes Emitter commands
        '''
        att_list = self.scope.sim_atts.findAttributes('Emitter')
        for att in att_list:
            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))

            boundary_string = utils.format_entity_string(self.scope, att)
            self.scope.output.write('  BoundaryID: %s\n' % boundary_string)

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

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

            self.finish_command()

    def write_materials(self):
        '''Writes Material commands
        '''
        surface_att_list = self.scope.sim_atts.findAttributes('Track3PSurface')
        solid_att_list = self.scope.sim_atts.findAttributes('Track3PSolid')
        att_list = surface_att_list + solid_att_list
        id_keyword_lookup = {
            'Track3PSurface': 'BoundarySurfaceID',
            'Track3PSolid': 'SolidBlockID'}
        for att in att_list:
            entity_string = utils.format_entity_string(self.scope, att)
            if entity_string == '':
                message = 'Warning: no model entities assigned to %s attribute, named %s' % \
                    (att.type(), att.name())
                print message
                self.scope.logger.addWarning(message)
                continue

            id_keyword = id_keyword_lookup.get(att.type(), 'Unknown')

            # Surface as Primary is option
            primary_item = att.findVoid('Primary')
            if primary_item.isEnabled():
                self.start_command('Material')
                self.scope.output.write('  Type: Primary\n')
                self.scope.output.write('  %s: %s\n' % (id_keyword, entity_string))
                self.finish_command()

            # Surface must be either Secondary or Absorber
            self.start_command('Material')

            sec_abs_item = att.findString('SecondaryOrAbsorber')
            self.scope.output.write('  Type: %s\n' % sec_abs_item.value(0))
            self.scope.output.write('  %s: %s\n' % (id_keyword, entity_string))

            # For Secondary, write out SEY Model and other specs
            if sec_abs_item.value(0) == 'Secondary':
                sey_model_path = 'SecondaryOrAbsorber/SEYModel'
                sey_model_item = att.itemAtPath(sey_model_path)
                sey_model_number = sey_model_item.value(0)
                self.scope.output.write('  Model: %d\n' % sey_model_number)

                # Write SEY Curve (expression)
                curve_path = '%s/SEYCurve' % sey_model_path
                CardFormat(['Einit', 'Sigma0'], curve_path).write(self.scope, att)

                if sey_model_number == 2:
                    args = (self.scope, att)  # for convenience/reuse
                    min_path = '%s/MinimumNumElectrons' % sey_model_path
                    CardFormat('MinimumNumElectrons', min_path, '%d').write(*args)
                    et_path = '%s/ElasticThreshold' % sey_model_path
                    CardFormat('ElasticThreshold', et_path, '%d').write(*args)

            self.finish_command()  # Material Secondary or Absorber

    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('    Toggle: %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('    Toggle: %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
