class CubitTemplate:    #Template definition for Cubit input file
    
    def __init__(self, *args):
        self.file = args[0]    #First argument is Cubit input file name
        if len(args)>1:    #Optional second argument is output file name
            self.file_out = args[1]
        else:
            self.file_out = args[0]
        self.parse_cubit_file()
        
    def parse_cubit_file(self):    #Read in Cubit input file
        with open(self.file) as file:
            self.lines = file.readlines()
        numrows = len(self.lines)    #Number of rows in input file
        self.ncflag = [0] * numrows    #Flag to store non-commented lines
        for i in range(numrows):
            if self.lines[i].strip() != '':    #Ignore empty rows
                if not self.lines[i].startswith('##'):  #Ignore comment rows
                    self.ncflag[i] = i+1
        self.ncflag = [i-1 for i in self.ncflag if i != 0]
    
    def get_value(self, key):  #Read value of first instance of variable
        value = None
        for i in range(len(self.ncflag)):
            line = self.lines[self.ncflag[i]]
            flag = line.replace(' ','').find(key + '=')
            if flag != -1:    #If variable name declaration found in file
                indx = line.find('=')    #Index of '=' sign following name
                if line.strip().startswith('#{'):    #If APREPRO code line
                    value = line[indx+1::].replace('}','')
                else:
                    value = line[indx+1::].strip()
                break
        if flag == -1:
            print('Error: \'' + key + '=\' not found in Cubit file, ' \
                  + 'value \'None\' returned.')
            return None
        else:
            return eval(value)
    
    def set_value(self, kwargs):    #Set value of first instance of variable
        for key, value in kwargs.items():
            for i in range(len(self.ncflag)):
                line = self.lines[self.ncflag[i]]
                flag = line.replace(' ','').find(key + '=')
                if flag != -1:    #If variable name declaration found in file
                    indx = line.find('=')    #Index of '=' sign following name
                    if line.strip().startswith('#{'):    #If APREPRO code line
                        new_line = line[0:indx] + '=' + str(value) + '}\n'
                    else:
                        new_line = line[0:indx] + '=' + str(value) + '\n'
                    self.lines[self.ncflag[i]] = new_line
                    break
            if flag == -1:
                print('Error: \'' + key + '\' not defined in Cubit file, ' \
                      + 'no action taken.')
    
    def set_export(self, name, format_='genesis', opts='overwrite'):
        for i in range(len(self.ncflag)):
            words = self.lines[self.ncflag[i]].split()
            if words[0] == 'export':
                new_line = ['export', format_, '"' + name + '"']
                if opts is not None:
                    if not isinstance(opts, list):
                        opts = [opts]
                    new_line = new_line + [opt for opt in opts]
                self.lines[self.ncflag[i]] = ' '.join(new_line + ['\n'])
                return
        print('Error: no export command found in Cubit journal.')
        
    def write_file(self, *file_out):
        if isinstance(file_out[0], str):
            with open(file_out[0], 'w') as file:
                file.writelines(self.lines)
        else:
            with open(self.file_out, 'w') as file:
                file.writelines(self.lines)            
                
# if __name__ == "__main__":  #Used for debugging
#     demo_template = CubitTemplate('uw-srf-gun-w-loop.jou')
#     print('acc_gap is ' + str(demo_template.get_value('acc_gap')))
#     demo_template.set_value({'acc_gap':200})
#     print('acc_gap is ' + str(demo_template.get_value('acc_gap')))
#     demo_template.set_export('SRF-GUN-X.gen')
#     print(demo_template.file_out)
#     demo_template.write_file()