class Omega3PTemplate(dict):    #Template class for Omega3P input
    
    def __init__(self, *args, **kwargs):
        super(Omega3PTemplate, self).__init__(*args, **kwargs)
        self['ModelInfo'] = {}
    
    def set_model_file(self, mesh_filename, tolerant=1):
        self['ModelInfo']['File'] = mesh_filename
        self['ModelInfo']['Tolerant'] = tolerant
    
    def add_boundary_condition(self, bc_type, bc_ref):
        if 'BoundaryCondition' not in self['ModelInfo'].keys():
            self['ModelInfo']['BoundaryCondition'] = {}
        if bc_type in ['Electric','Magnetic','Exterior','Impedance','Absorbing', \
                       'Waveguide','Port','Periodic_M','Periodic_S','Theta']:
            self['ModelInfo']['BoundaryCondition'][bc_type] = bc_ref
        else:
            print('Boundary condition type not recognized.')
    
    def add_surface_material(self, sm_ref, sm_sigma=5.8e7):
        scount = 0
        sname = 'SurfaceMaterial'+str(scount)
        while sname in self['ModelInfo'].keys():    #Each material is a dict
            scount += 1
            sname = 'SurfaceMaterial'+str(scount)
        self['ModelInfo'][sname] = {}
        self['ModelInfo'][sname]['ReferenceNumber'] = sm_ref
        self['ModelInfo'][sname]['Sigma'] = sm_sigma
    
    def add_material(self, mat_ref, mat_eps=1.0, mat_mu=1.0):
        mcount = 0
        mname = 'Material'+str(mcount)
        mat_eps = complex(mat_eps)
        mat_mu = complex(mat_mu)
        while mname in self['ModelInfo'].keys():    #Each material is a dict
            mcount += 1
            mname = 'Material'+str(mcount)
        self['ModelInfo'][mname] = {}
        self['ModelInfo'][mname]['Attribute'] = mat_ref
        self['ModelInfo'][mname]['Epsilon'] = mat_eps.real
        self['ModelInfo'][mname]['Mu'] = mat_mu.real
        if mat_eps.imag != 0:
            self['ModelInfo'][mname]['EpsilonImag'] = mat_eps.imag
        if mat_mu.imag != 0:
            self['ModelInfo'][mname]['MuImag'] = mat_mu.imag
    
    def set_finite_element(self, order, curved_flag='on'):
        self['FiniteElement'] = {}
        self['FiniteElement']['Order'] = order
        if curved_flag.lower() in ['on','off']:
            self['FiniteElement']['CurvedSurfaces'] = curved_flag.lower()
        else:
            print('Curved surfaces flag must be set to \'on\' or \'off\'.')
    
    def add_pregion(self, ref, order):
        if 'FiniteElement' not in self.keys():
            self['FiniteElement'] = {}
        self['FiniteElement']['Order'] = 0    #Set order to 0 if pregion is used
        prcount = 0
        prname = 'PRegion'+str(prcount)
        while prname in self.keys():    #Each pregion is its own dict
            prcount += 1
            prname = 'PRegion'+str(prcount)
        self[prname] = {}
        self[prname]['Type'] = 'Material'
        self[prname]['Reference'] = ref
        self[prname]['Order'] = order
    
    def set_eigensolver(self, num_eigs, freq_shift, tol=1e-7, max_iters=1000):
        self['EigenSolver'] = {}
        self['EigenSolver']['NumEigenvalues'] = num_eigs
        self['EigenSolver']['FrequencyShift'] = freq_shift
        self['EigenSolver']['Tolerance'] = tol
        self['EigenSolver']['MaxIterations'] = max_iters
        
    def write_file(self, filename):    #Unpack and write template to file
        lines = self.unpack_template(self)
        with open(filename, 'w') as file:
            file.writelines(lines)
    
    def unpack_template(self, temp, lines=None, depth=0):    #Recursively unpack
        if isinstance(lines, type(None)):
            lines = []
        for key, value in temp.items():
            if key.startswith(('SurfaceMaterial','Material','PRegion')):
                lines.append(''.join(c for c in key if not c.isnumeric()) + ' : ')
            else:
                lines.append(key + ' : ')
            if isinstance(temp[key], dict):    #Unpack dict within dict
                lines.append('{\n')
                self.unpack_template(temp[key], lines, depth+1)
                lines.append('}\n')
                if depth==0:    #Add a blank line between top-level keys
                    lines.append('\n')
            else:
                lines.append(str(value).strip('[').strip('(').replace(',','')
                             .replace('\'','').strip(')').strip(']'))
                lines.append('\n')    #Remove '(', ',', and ')' in value
        return lines

# if __name__ == "__main__":
#     demo_template = Omega3PTemplate()                            #Create demo template
#     demo_template.set_model_file('my_mesh.ncdf')                 #Set model info mesh file
#     demo_template.add_boundary_condition('Electric',(1,2)) #Set boundary conditions for surfaces
#     demo_template.add_boundary_condition('Waveguide',3)
#     demo_template.add_boundary_condition('Magnetic',(4,5))
#     demo_template.add_surface_material(1,3e6)              #Add surface materials with specified sigma
#     demo_template.add_surface_material(2,5e9)
#     demo_template.add_material(3,1.0,2.0)       #Add a material property with specified epsilon and mu
#     demo_template.set_finite_element(2,'on')                     #Set finite element parameters
#     demo_template.add_pregion(2,2)                               #Set region 2 to use 2nd order FE
#     demo_template.add_pregion(3,3)                               #Set region 3 to use 3rd order FE
#     demo_template.set_eigensolver(1,1.3e9)                       #Set eigensolver paramaters
#     demo_template.write_file('demo_o3p.in')                      #Write template to file for use by omega3p