class AcdtoolTemplate(dict):    #Template class for Omega3P input
    
    def __init__(self, *args, **kwargs):
        super(AcdtoolTemplate, self).__init__(*args, **kwargs)
        self['RFField'] = {}
    
    def set_result_file(self, omega3p_file='omega3p_results'):
        self['RFField']['ResultDir'] = omega3p_file
    
    def set_mode_info(self, freq_scan=0, mode_id=0):
        self['RFField']['FreqScanID'] = freq_scan
        self['RFField']['ModeID'] = mode_id
    
    def set_symmetry(self, xsymmetry='none', ysymmetry='none'):
        self['RFField']['xsymmetry'] = xsymmetry
        self['RFField']['ysymmetry'] = ysymmetry
    
    def set_gradient(self, x0, y0, gz1, gz2, npoint, fmnx, fmny, fmnz, 
                     gradient=-1, cavityBeta=1.0, revPflow=0):
        self['RFField']['gradient'] = gradient
        self['RFField']['cavityBeta'] = cavityBeta
        self['RFField']['reversePowerFlow'] = revPflow
        self['RFField']['x0'] = x0
        self['RFField']['y0'] = y0
        self['RFField']['gz1'] = gz1
        self['RFField']['gz2'] = gz2
        self['RFField']['npoint'] = npoint
        self['RFField']['fmnx'] = fmnx
        self['RFField']['fmny'] = fmny
        self['RFField']['fmnz'] = fmnz
        
    def set_openpmd(self, nx, ny, nz, x1, x2, y1, y2, z1, z2, ionoff=1):
        self['OpenPMD_IMPACT'] = {}
        self['OpenPMD_IMPACT']['ionoff'] = ionoff
        self['OpenPMD_IMPACT']['nx'] = nx
        self['OpenPMD_IMPACT']['ny'] = ny
        self['OpenPMD_IMPACT']['nz'] = nz
        self['OpenPMD_IMPACT']['x1'] = x1
        self['OpenPMD_IMPACT']['x2'] = x2
        self['OpenPMD_IMPACT']['y1'] = y1
        self['OpenPMD_IMPACT']['y2'] = y2
        self['OpenPMD_IMPACT']['z1'] = z1
        self['OpenPMD_IMPACT']['z2'] = z2
        
    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 depth == 0:
                lines.append(key + '\n')
            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 = AcdtoolTemplate()
#     demo_template.set_result_file()
#     demo_template.set_mode_info()
#     demo_template.set_symmetry('magnetic','magnetic')
#     demo_template.set_gradient(0.0, 0.0, 0.0, 0.3, 1000, 20, 20, 100)
#     demo_template.set_openpmd(101,101,301,-0.015,0.015,-0.015,0.015,0.0,0.3)
#     demo_template.write_file('demo_rfpost.in')