#################################################################################
#                                                                               #
# This small program reads in a SPEC data file and splits it into individual    #
# scan files. All scans are split apart and so no user input is required.       #
#                                                                               #
# Since sometimes data analysis software has an issue with trailing line in a   #
# data file, there is the option to truncate lines from the end of every scan;  #
# an equal number of lines will be truncated from every scan is this is         #
# selected.                                                                     #
#                                                                               #
# v 1.0 - 2024.02.22                                                            #
#   - Almost a total rewrite.  Now allows for line truncation from the end      #
#   of every scan and should no longer have line duplication issues that        #
#   cropped up in weird edge cases.                                             #
#                                                                               #
# v 0.2 - 2020.4.9                                                              #
#   - modified for use in PREM workshop                                         #
#   - output scan files now include the name of the parent SPEC file            #
#                                                                               #
# v 0.1 - 2018.7.12                                                             #
#   - initial iteration                                                         #
#                                                                               #
# CJP, CHESS                                                                    #
#################################################################################

import re
import os
import ntpath

# Define an empty Python list to store data from file
lines = []

# Open the data file and append all lines to the list
inputFile = input("Enter the path to your SPEC file: ")
print('')
assert os.path.exists(inputFile), "File doesn't seem to exist..."
specName = str(ntpath.basename(inputFile))

excluded = int(input("How many lines do you want to exclude from the end of each scan (default  = 0)? "))
if isinstance(excluded,int) and excluded > 0:
    pass
else:
    excluded = 0

with open(inputFile,'r') as f:
    for line in f:
        lines.append(line)

# Set some things up before beginning...
start = re.compile(r"#S\b")     # Choose the character string that delineates the beginning of a scan
x = 1                           # This is the index that incremements when a new #S is found
scanDict = {}                   # This is the dictionary that will store all of the scan information
scanArray = []                  # This array will assign the scan data to the dictionary

# Search Python list for new scans, then add each scan to a new file
for row in lines:
    if start.search(row) != None:               # Searh each line for the new scan delineator
        scanArray = []                          # Clear the data array each time a new #S is found
        scanArray.append(row)                   # Append this row to the data array
        scanDict[str(x)] = scanArray            # Assign the contents of the scan array to the corresponding index number in the dictionary
        x = x + 1                               # Increment x by 1
    else:
        scanArray.append(row)                   # Also append the row to the data array
        scanDict[str(x-1)] = scanArray          # Associate the data array with the scan index in the dictionary

# Write each scan to its own file, truncating the user-indicated number of trailing lines
counter = 1
while counter < x:
    filename = specName + '_scan_' + str(counter) + '.txt'    # Make a new filename using the 'x' index
    numLines = len(scanDict[str(counter)])
    saveLines = numLines - excluded
    currentLine = 1
    print("\nFor scan #" + str(counter) + " writing " + str(saveLines) + " out of a total of " + str(numLines) + " lines...")
    f = open(filename,'w+')                # Create a new file using the generated 'filename'
    for scanLine in scanDict[str(counter)]:
        if currentLine <= saveLines:
            f.write(scanLine)                           # Write the row to the created file
        currentLine += 1
    print("Finiished writing " + str(filename))
    counter += 1

input('\nFinished!  Press any key to exit...')