import os

def slicer(string):
    out = []
    string = string.strip()
    string = string.split(" ")
    for i in string:
        if i != "":
            out.append(i)
    return out

# Declaration of variables
filelist = []
thrprint = 0

# Generation of output directory and path
os.mkdir("OUTPUT")
path = os.getcwd() + r"\OUTPUT"

# Request of descriptors of interest
char = input("Do you want to collect IBO charges?(y/n): ")
wbi = input("Do you want to collect IBO WBIs?(y/n): ")
comp = input("Do you want to collect orbital compositions?(y/n): ")
qtaim = input("Do you want to collect QTAIM charges?(y/n): ")

if char == "y":
    char = True
if wbi == "y":
    wbi = True
if comp == "y":
    comp = True
if qtaim == "y":
    qtaim = True

# Definition of output files contents
outp_char = ["File Name;M;C(Me1);H1(Me1);H2(Me1);H3(Me1);C(Me2);H1(Me2);H2(Me2);H3(Me2);Total\n"]
outp_wbi = ["File Name;WBI(M-Me1);WBI(M-Me2)\n"]
outp_comp = ["File Name;M;C(Me1);M;C(Me2)\n"]
outp_qtaim = ["File Name;M;C(Me1);H1(Me1);H2(Me1);H3(Me1);C(Me2);H1(Me2);H2(Me2);H3(Me2);Total\n"]

# Generation of a list containing IBO file names
if char == True or wbi == True or comp == True:
    for x in os.listdir():
        if x.endswith(".scf-log.txt"):
            filelist.append(x)
        
    # Descriptors colletion
    for x in filelist:
        outp_comp.append(x + ";")
        outp_char.append(x + ";")
        outp_wbi.append(x + ";")
        if comp == True:
            outp_comp.append("N/A;")
            outp_comp.append("N/A;")
            outp_comp.append("N/A;")
            outp_comp.append("N/A\n")
        inp = open(x, "r")
        inp_content = inp.readlines()
        for i in range(len(inp_content)):
            if "AB" in inp_content[i] and " 2  " in inp_content[i][8:] and " 1  " in inp_content[i][8:]:
                for j in range(len(slicer(inp_content[i]))-4):
                    if slicer(inp_content[i])[j+4] == "2":
                        outp_comp[-3] = slicer(inp_content[i])[j+5] + ";"
                    if slicer(inp_content[i])[j+4] == "1":
                        outp_comp[-4] = slicer(inp_content[i])[j+5] + ";"
            if "AB" in inp_content[i] and " 6  " in inp_content[i][8:] and " 1  " in inp_content[i][8:]:
                for j in range(len(slicer(inp_content[i]))-4):
                    if slicer(inp_content[i])[j+4] == "6":
                        outp_comp[-1] = slicer(inp_content[i])[j+5] + "\n"
                    if slicer(inp_content[i])[j+4] == "1":
                        outp_comp[-2] = slicer(inp_content[i])[j+5] + ";"
            if "Bond order analysis:" in inp_content[i]:
                thrprint = inp_content[i][38:47]
                outp_wbi.append("<" + str(thrprint) + ";")
                outp_wbi.append("<" + str(thrprint) + "\n")
            if wbi == True and " 1  " in inp_content[i] and "-     2   C " in inp_content[i]:
                outp_wbi[-2] = inp_content[i][28:35] + ";"
            if wbi == True and " 1  " in inp_content[i] and "-     6   C " in inp_content[i]:
                outp_wbi[-1] = inp_content[i][28:35] + "\n"
            if char == True and "Total charge composition:" in inp_content[i]:
                for j in range(9):
                    outp_char.append(inp_content[i+3+j][62:70] + ";")
            if char == True and "-> Total" in str(inp_content[i]):
                outp_char.append(inp_content[i][62:70] + "\n")
        inp.close()
        print("The file " + x + " has been analyzed")

# Generation of a list containing QTAIM file names
if qtaim == True:
    filelist = []
    for x in os.listdir():
        if x.endswith(".eos-log.txt"):
            filelist.append(x)
        
    # Descriptors colletion
    for x in filelist:
        outp_qtaim.append(x + ";")
        inp = open(x, "r")
        inp_content = inp.readlines()
        for i in range(len(inp_content)):
            if char == True and "Assignment of red(O)x-type charge/spin vs. (P)artial charge/spin:" in inp_content[i]:
                for j in range(9):
                    outp_qtaim.append(inp_content[i+3+j][66:74] + ";")
            if "-> Total" in str(inp_content[i]):
                outp_qtaim.append(inp_content[i][66:74] + "\n")
        inp.close()
        print("The file " + x + " has been analyzed")

# Generation of CSV files in the output directory
os.chdir(path)
if char == True:
    outp = open("IBO_charges.txt", "w")
    outp.writelines(outp_char)
    outp.close()
    os.rename("IBO_charges.txt", "IBO_charges.csv")
if wbi == True:
    outp = open("IBO_WBIs.txt", "w")
    outp.writelines(outp_wbi)
    outp.close()
    os.rename("IBO_WBIs.txt", "IBO_WBIs.csv")
if comp == True:
    outp = open("IBO_bond_composition.txt", "w")
    outp.writelines(outp_comp)
    outp.close()
    os.rename("IBO_bond_composition.txt", "IBO_bond_composition.csv")
if qtaim == True:
    outp = open("QTAIM_charges.txt", "w")
    outp.writelines(outp_qtaim)
    outp.close()
    os.rename("QTAIM_charges.txt", "QTAIM_charges.csv")