# -*- coding: utf-8 -*-
#########################################################
# JWL - updated 2016
#########################################################
# Basic Instructions
#########################################################
"""
1. Set up a template .inp file for the parent in the correct format with correct flags (see example .inp)
2. Change directories and information in lines 34-42 below
3. Change to the right directory with "cd <directory>"
4. Press green run arrow above
"""

#########################################################
# Import Required Modules
#########################################################
import os
import re
import math
import subprocess
from shutil import copyfile

#########################################################
# Failsafe Function - Allows system reset on unexpected occurence
#########################################################

def dead(reason): # indicates script closure with reason
    print reason
    raise SystemExit
    
#########################################################
# User required info - Remember to cd to working directory before running
#########################################################

working_directory = r'S:\james_lewis\WO3_exhaustive' # directory where trial subgroups .strs are stored
seed_file =         working_directory + r'/WO3_start.inp' # pre setup .inp file to be automatically edited for each subgroup
SG_reference_file = working_directory + r'/SGCOM5_april2016.CPP' # List of space groups in format understood by topas this is topas file but "fixed" version with line numbers fixed to crystal class below 
result_file =       working_directory + r'/results_summary.txt' # New file created any old one with the same name is overwritten
topas_directory =   r'C:\topasa_v6' # directory topas where tc.exe is located
basefile = r"topas" # prefix to subgroups .str files generated by ISODISTORT, could be in another directory
thermal_prms = {"W":1.43932, "O":1.66845} # Do not have a dictionary entry if beq not to be changed 
subgroup_numbers_to_run = [x for x in range(5,1428) if x!=7 and x not in range(217,220) and x!=221 and x!=223 and x!=225] # List of subgroup numbers (suffix of ISODISTORT output) to run
number_of_steps = 3 # number of different instructions sets in the topas.inp to run for each subgroup i.e. #ifdef run_1, #ifdef run_2, etc

#########################################################
# Search compilers (to allow editing of ISODISTORT .str files)
# Assumes fixed formats for these lines in .str files as below
#########################################################

Space_group_line = re.compile(r"\s+space_group\s+?")  # anything then "space_group" then any string
mode_site_line = re.compile("\t\tsite") # two tabs then "site"
displacive_line = re.compile(r"\s+prm\s+!a\d+\s+-?\d+\.?\d+'?") # anything then "prm" then "!a" then any integer then any number then any string then any integer then anything then any integer 

#########################################################
# Main Inp File Creation Function
#########################################################

def inp_maker(basefile, file_number, thermal_prms): # this function converts all str files into input files in the seed format
    new_inp_file = (basefile + "%04d" % (file_number) + '.inp')  # Name new .inp file by its subgroup number
    structural_file = (basefile + "%04d" % (file_number) + '.str') # Structural file identified by number as obtained from ISODISTORT
    new_inp_file_stream = open(new_inp_file,'w') # write to new .inp
    structural_file_stream = open(structural_file, 'r') # read ISODISTORT .str
    structure_lines = structural_file_stream.readlines() # list .str by line
    structural_file_stream.close() # close .str
    displacive_modes_in_this_str = 0 # Begin count of number of displacive mode in the .str
    for line in structure_lines: # For each line in .str
        if displacive_line.search(line): # If the line matches displacive regular expression (line 51)
            displacive_modes_in_this_str += 1 # Add one to the displacive mode count
    space_group, crystal_class = get_space_group(structure_lines) # Call function that returns space group and crystal class of subgroup (see line 104-136)
    section_flag = 0 # Start by saying we are not within one of the section flags in the start file
    for i, line in enumerate(open(seed_file).readlines()): # For each line in the start file, numbering them starting with line 1 as i = 0.
        if re.search("""'!End_Lattice_Prm_Section!""", line) or re.search("""'!End_Insert_Sites_Section!""", line): # If either of these two end section flags appear in the line
            section_flag = 0 # Then declare we are no longer within a section flag
        elif section_flag == 1: continue # If we are within a section flag move directly to the next line of the start file
        elif re.search("""'!Lattice_Prm_Section!""", line): # If this expression appears in the line 
            refine_lps(new_inp_file_stream, structure_lines, crystal_class) # Add lattice parameter information to the .inp (see lines 138-162)
            section_flag = 1 # And declare we are within a section flag
        elif re.search("""'!Insert_Sites_Section!""", line): # If this expression appears in the line 
            add_structural_info(new_inp_file_stream, structure_lines, thermal_prms) # Add distortion modes and site information to the .inp (see lines 91-102)
            section_flag = 1 # And declare we are within a section flag
        elif re.search("""'!Insert_Penalties_Section!""", line): # If this expression appears in the line 
            for i in range(0, displacive_modes_in_this_str): # Then for every mode in the .str
                new_inp_file_stream.write("aP(a%d)\n" % (i+1)) # Add a macro aP(mode) - the start file will indicate what penalty this macro applies
        else: # If not in a section flag and the line contains none of the above expressions
            new_inp_file_stream.write(line) # Then right the line unchanged into the new .inp file
    new_inp_file_stream.close() # close the .inp file

###########################################################
# Supplementary Editing Functions
###########################################################

def add_structural_info(new_inp_file_stream, structure_lines, thermal_prms): # Add mode definitions and structural sites
    for i, line in enumerate(structure_lines[13:]): # For lines between 14th line of the .str file and the end
        site = None # Assume it is not a line containing a site
        if displacive_line.search(line): # If the line contains the standard displacive mode line expression (line 51)
            new_inp_file_stream.write(line.replace("prm  !a", "mode_onoff(a").replace("0.00000", ",  0.00000) drand '")) # Then replace the !a mode label with a macro "modeonoff(mode)" label and the drand macro label. The start file will indicate what these macros mean.
        elif mode_site_line.search(line): # If the line matches the standard site line expression (line 50)
            for atom in thermal_prms: # For each different atom type that a thermal parameter is given for (line 40)
                if ("site %s" % (atom)) in line: # If the site has that atom in it 
                    site = atom # Declare the atom type of the site
                    break # And end the "for atom" loop
            new_inp_file_stream.write(line.replace("beq 0.0", "beq %f" % (thermal_prms[site])).replace(";", ";: 1.00000")) # replace the thermal parameter on that line with the thermal parameter found in the thermal prms dictionary
        else: new_inp_file_stream.write(line) # If it is neither a site or displacive mode line then write the line into the new .inp file unchanged
        
def get_space_group(structure_lines): # determine space group and crystal class (in order to determine which lps to refine)
    for line in structure_lines: # for each line in the .str file
        if Space_group_line.search(line): # If the line matches the standard space group line expression (line 49)
            space_group = line.split()[1] # The space group is the second part of the line when it is split by whitespace (the line will read e.g. "   space_group P1   "
            for k, line in enumerate(open(SG_reference_file).readlines()): # search the standardised Topas space group file "SGCOM5.CPP" by line, numbering each line j starting with j=0. This is a fixed copy from April 2016. It shouldnt be changed.
                  if space_group.lower() + ' ' in line: # if lowercase space group is in the line
                      if k==17 or k==18: # direct to relevant space group definition based on line number space group found in (remember line-1 in SGCOM5.CPP as python starts at zero)
                          print 'triclinic space group detected!'
                          crystal_class = "triclinic"
                      elif k>=21 and k<=87:
                          print 'monoclinic space group detected!'
                          crystal_class = "monoclinic"
                      elif k>=90 and k<=180:
                          print 'orthorhombic space group detected!'
                          crystal_class = "orthorhombic"
                      elif k>=183 and k<=271:
                          print 'tetragonal space group detected!'
                          crystal_class = "tetragonal"
                      elif k>=274 and k<=276 or k==279 or k>=282 and k<=290 or k>=293 and k<=296 or k>=301 and k<=305:
                          print 'trigonal space group detected!'
                          crystal_class = "trigonal"
                      elif k>=277 and k<=278 or k>=280 and k<=281 or k>=291 and k<=292 or k>=297 and k<=300 or k>=306 and k<=311:
                          print 'trigonal_hexagonal_setting space group detected'
                          crystal_class = "trigonal_hexagonal_setting"
                      elif k>=314 and k<=343:
                          print 'hexagonal space group detected!'
                          crystal_class = "hexagonal"
                      elif k>=346 and k<=387:
                          print 'cubic space group detected!'
                          crystal_class = "cubic"
                      else:
                          dead('space group not assigned') # System Exit with reason if the space group does not appear in the Topas space group file.
    return space_group, crystal_class # return space group and crystal class found

def refine_lps(new_inp_file_stream, structure_lines, crystal_class): # Add lattice prms with appropriate refinement flags
    if crystal_class == "triclinic": # if triclinic
        refinement_flags = ["lpa", "lpb", "lpc", "alph", "bet", "gam"] # a!=b!=c alpha, beta, gamma != 90
    elif crystal_class == "monoclinic": # if monoclinic
        refinement_flags = ["lpa", "lpb", "lpc", "", "bet", ""] # a!=b!=c alpha=gamma= 90, beta != 90
    elif crystal_class == "orthorhombic":
        refinement_flags = ["lpa", "lpb", "lpc", "", "", ""] # a!=b!=c alpha=beta=gamma= 90
    elif crystal_class == "tetragonal":
        refinement_flags = ["lpa", "lpa", "lpc", "", "", ""] # a=b!=c alpha=beta=gamma= 90
    elif crystal_class == "trigonal":
        refinement_flags = ["lpa", "lpa", "lpc", "", "", ""] # In ISODISTORT all trigonal lattices are in a hexagonal setting with a=b!=c alpha = 90 beta = 90 gamma = 120
    elif crystal_class == "trigonal_hexagonal_setting":
        refinement_flags = ["lpa", "lpa", "lpc", "", "", ""] # a=b!=c alpha = 90 beta = 90 gamma = 120
    elif crystal_class == "hexagonal":
        refinement_flags = ["lpa", "lpa", "lpc", "", "", ""] # a=b!=c alpha = 90 beta = 90 gamma = 120
    elif crystal_class == "cubic":
        refinement_flags = ["lpa", "lpa", "lpa", "", "", ""] # a=b=c alpha=beta=gamma= 90
    new_inp_file_stream.write((structure_lines[4])) # Write the space group (the 5th line of .str file) into the new .inp file
    ###################################################################################
    # Add the refinement flag as indicated above and additional information on min/max values of cell parameter, its step and annealing. Cell_perc is defined in the start file.
    ###################################################################################
    new_inp_file_stream.write((structure_lines[6].replace("\n", "").replace("a", "a %s" % (refinement_flags[0]))) + " del = 0.0001 Val; val_on_continue = Rand((100-cell_perc)/100,(100+cell_perc)/100)*a_start;	min = 0.95*a_start;	max = 1.05*a_start;\n")
    new_inp_file_stream.write((structure_lines[7].replace("\n", "").replace("b", "b %s" % (refinement_flags[1]))) + " del = 0.0001 Val; val_on_continue = Rand((100-cell_perc)/100,(100+cell_perc)/100)*b_start;	min = 0.95*b_start;	max = 1.05*b_start;\n")                
    new_inp_file_stream.write((structure_lines[8].replace("\n", "").replace("c", "c %s" % (refinement_flags[2]))) + " del = 0.0001 Val; val_on_continue = Rand((100-cell_perc)/100,(100+cell_perc)/100)*c_start;	min = 0.95*c_start;	max = 1.05*c_start;\n")                
    new_inp_file_stream.write((structure_lines[9].replace("\n", "").replace("al", "al %s" % (refinement_flags[3]))) + " del 0.01 val_on_continue = Rand((100-cell_perc)/100,(100+cell_perc)/100)*al_start;	min = 0.95*al_start;	max = 1.05*al_start;\n")               
    new_inp_file_stream.write((structure_lines[10].replace("\n", "").replace("be", "be %s" % (refinement_flags[4]))) + " del 0.01 val_on_continue = Rand((100-cell_perc)/100,(100+cell_perc)/100)*be_start;	min = 0.95*be_start;	max = 1.05*be_start;\n")              
    new_inp_file_stream.write((structure_lines[11].replace("\n", "").replace("ga", "ga %s" % (refinement_flags[5]))) + " del 0.01 val_on_continue = Rand((100-cell_perc)/100,(100+cell_perc)/100)*ga_start;	min = 0.95*ga_start;	max = 1.05*ga_start;\n")
    new_inp_file_stream.write("volume vol  0\n")
        
#####################################################################
# Execution Section 
#####################################################################
#####################################################################
# Create Result Files and Folders (in Working Directory)
#####################################################################
result_file_stream = open(result_file, 'w') # Create a new result file with name as given on line 37
result_file_stream.write("subgroup   rwp   a   b   c   al   be   ga   volume\n") # Header the result file with the information that will be collected
result_file_stream.close() # close the result file
Folders_to_create = ["start_inp_files\\", "best_inp_files\\"] # List of folders to create to store the .inp files created and what they look like after refinement
for folder in Folders_to_create: # for folder in above list
    if not os.path.exists(os.path.dirname(working_directory + r"\\%s" % (folder))): # if they don't already exist
        os.makedirs(os.path.dirname(working_directory + r"\\%s" % (folder))) # make them
#####################################################################
# Subgroup Tests
#####################################################################
for subgroup in subgroup_numbers_to_run: # for subgroup in list to run as defined on line 41
    #####################################################################
    # Make inp file
    #####################################################################
    inp_maker(basefile, subgroup, thermal_prms) # Make a new .inp file as defined by the function starting on line 57
    copyfile(working_directory + r'/' + basefile + "%04d" % (subgroup) + r""".inp""", working_directory + r'\start_inp_files/' + basefile + "%04d_start" % (subgroup) + r""".inp""") # save original (unrefined) input file (copy) into start .inp files folder
    for run in range(1,number_of_steps+1): # For each step in the run process as defined on line 42. The start file will indicate the differences between the runs using #ifdef run_x statements.
       #####################################################################
       # Send to topas
       #####################################################################
       print "Running Subgroup Number %d - Run %d..." % (subgroup, run)
       command = topas_directory + r"""\tc.exe """ + basefile + "%04d" % (subgroup) + r""".inp""" + r""" "#define GUI_LINES #define run_%d macro subgroup_number {prm !subgroup %d}""" % (run, subgroup) # Create function for Topas giving the run and subgroup number
       a = subprocess.call(command, creationflags=0x08000000) # Run the file in Topas (creationflag removes console flash)
    #####################################################################
    # Save ouputs
    #####################################################################    
       copyfile(working_directory + r'/' + basefile + "%04d" % (subgroup) + r""".out""", working_directory + r'/' + basefile + "%04d" % (subgroup) + r""".inp""") # turn the .out file generated by topas into next step .inp (copy and overwrite)
    os.rename(working_directory + r'/' + basefile + "%04d" % (subgroup) + r""".out""", working_directory + r'\best_inp_files/' + basefile + "%04d_best" % (subgroup) + r""".inp""") # save the final best (refined) .out file as .inp file (cut and paste)
