import itertools
import os
import statistics
import subprocess
import sys
from pathlib import Path

import pandas as pd


def checkName(command):
    commandLength = len(command)
    if commandLength > 2 or commandLength < 2:
        print("ERROR: expected 1 argument (run name); got " + str(commandLength - 1) + " instead.")
        return False
    elif command[1] in (item.name for item in Path(os.getcwd()).iterdir()):
        print('ERROR: run name "' + command[1] + '" already exists.')
        return False
    else:
        return True


def readFile(filepath):
    file = open(filepath, "r")
    rawList = file.readlines()
    file.close()
    finalList = [item.replace("\n", "") for item in rawList]
    return finalList


def checkLetter(letter):
    residueList = {
        "A",
        "C",
        "D",
        "E",
        "F",
        "G",
        "H",
        "I",
        "K",
        "L",
        "M",
        "N",
        "P",
        "Q",
        "R",
        "S",
        "T",
        "V",
        "W",
        "Y",
    }
    return letter in residueList


def checkMutationList(mutationList, proteinList):
    for i in mutationList:
        print("Checking mutation " + i)
        mutation = i.split(",")
        for j in mutation:
            j = j.split(".")
            if not checkMutation(j, proteinList):
                return False
    return True


def checkMutation(mutation, proteinList):
    if len(mutation) != 3 and len(mutation) != 4:
        print("ERROR: expected 3 or 4 parameters; got " + str(len(mutation)) + " instead.")
    else:
        chain = mutation[0]
        if ord(chain) < 65 or ord(chain) > len(proteinList) + 64:
            print("ERROR: " + chain + " is not a valid chain.")
        else:
            wildtypeResidue = mutation[1]
            if not checkLetter(wildtypeResidue):
                print("ERROR: " + wildtypeResidue + " is not a valid wildtype residue.")
            else:
                mutationPosition = mutation[2]
                if not mutationPosition.isnumeric():
                    print("ERROR: " + mutationPosition + " is not a valid mutation position.")
                else:
                    mutationPosition = int(mutationPosition)
                    if mutationPosition > len(proteinList[ord(chain) - 65]):
                        print(
                            "ERROR: "
                            + str(mutationPosition)
                            + " exceeds the length of chain "
                            + chain
                            + "."
                        )
                    else:
                        wildtypeResidueFromProtein = proteinList[ord(chain) - 65][
                            mutationPosition - 1
                        ]
                        if wildtypeResidue != wildtypeResidueFromProtein:
                            print(
                                "ERROR: expected wildtype residue "
                                + wildtypeResidueFromProtein
                                + "; got "
                                + wildtypeResidue
                                + " instead."
                            )
                        else:
                            if len(mutation) == 4:
                                mutationResidue = mutation[3]
                                if not checkLetter(mutationResidue):
                                    print(
                                        "ERROR: "
                                        + mutationResidue
                                        + " is not a valid mutant residue."
                                    )
                                elif mutationResidue == wildtypeResidue:
                                    print(
                                        "ERROR: mutant residue is the same as the wildtype residue."
                                    )
                                else:
                                    return True
                            else:
                                return True


def saturateList(mutationList, residueList):

    saturatedMutationList = []

    for i in mutationList:
        i = i.split(",")
        saturationMutations = []
        saturationMutationsIndex = []

        for j in range(len(i)):
            if len(i[j].split(".")) == 3:
                saturationMutations.append(i[j])
                saturationMutationsIndex.append(j)

        permutations = itertools.product(residueList[0], repeat=len(saturationMutations))

        for permutation in permutations:
            mutations = ",".join(
                f"{saturationMutations[j]}.{permutation[j]}"
                for j in range(len(saturationMutations))
            )
            mutations = mutations.split(",")

            checkMutations = i

            for k in range(len(saturationMutations)):
                checkMutations[saturationMutationsIndex[k]] = mutations[k]

            attach = True

            for k in checkMutations:
                k = k.split(".")
                if k[1] == k[3]:
                    attach = False
                    break

            if attach:
                saturatedMutationList.append(",".join(checkMutations))

    return saturatedMutationList


def makeDirectory(runName):
    os.mkdir(os.getcwd() + "/" + runName)
    os.chdir(os.getcwd() + "/" + runName)
    os.mkdir(os.getcwd() + "/mutations")
    os.mkdir(os.getcwd() + "/alphafold")


def generateFasta(mutation, mutantProtein):
    mutantProtein = ":".join(protein.strip() for protein in mutantProtein)
    file = open(os.getcwd() + "/mutations/" + mutation + ".fasta", "w")
    _ = file.write(">" + mutation + "\n" + mutantProtein)
    file.close()


def mutateProteinList(mutationList, proteinList):
    for i in mutationList:
        print("Generating " + i + ".fasta")
        mutantion = i.split(",")
        mutantProtein = proteinList.copy()
        for j in mutantion:
            j = j.split(".")

            chain = j[0]
            wildtypeResidue = j[1]
            mutationPosition = int(j[2])
            mutationResidue = j[3]

            mutantProtein[ord(chain) - 65] = (
                mutantProtein[ord(chain) - 65][: mutationPosition - 1]
                + mutationResidue
                + mutantProtein[ord(chain) - 65][mutationPosition:]
            )

        generateFasta(i, mutantProtein)


def alphafold(mutationList):
    print("Finished creating fasta files; now creating alphafold structures.")
    for i in mutationList:
        print("Generating " + i + " pdbs")
        command = (
            "colabfold_batch "
            + os.getcwd()
            + "/mutations/"
            + i
            + ".fasta "
            + os.getcwd()
            + "/alphafold/"
            + i
        )
        result = os.system(command)
        if result != 0:
            return False
    return True


def prodigy(mutationList, runName):

    print("Finished creating alphafold structures; now calculating KD values.")
    kdTable = []
    kdTable.append(
        [
            "Mutation",
            "KD1 (uM)",
            "KD2 (uM)",
            "KD3 (uM)",
            "KD4 (uM)",
            "KD5 (uM)",
            "avg KD (uM)",
            "SD",
        ]
    )

    for i in mutationList:
        pdbs = sorted(
            list(
                Path(
                    os.path.expanduser("~") + "/Documents/pipeline/" + runName + "/alphafold/" + i
                ).rglob("*.pdb")
            )
        )
        kdRow = []
        kdRow.append(i)
        for pdb in pdbs:
            os.chdir(os.path.expanduser("~") + "/prodigy/")
            command = (
                "source venv/bin/activate && cd "
                + os.path.expanduser("~")
                + "/Documents/pipeline/"
                + runName
                + "/alphafold/"
                + i
                + " && prodigy "
                + str(pdb)
            )
            result = subprocess.run(command, shell=True, capture_output=True, text=True)
            if result.stderr != "":
                print("ERROR: " + result.stderr)
                return
            else:
                kd = float(result.stdout.split("\n")[12][53:]) * 1000000
                kdRow.append(kd)

        kdRow.append(statistics.mean(kdRow[1:6]))
        kdRow.append(statistics.stdev(kdRow[1:6]))

        kdTable.append(kdRow)

    os.chdir(os.path.expanduser("~") + "/Documents/pipeline/" + runName)

    df = pd.DataFrame(kdTable)
    df.to_excel(runName + ".xlsx", index=False, header=False)

    print("Finished calculating KD values; XLSX file generated.")


def main():
    if checkName(sys.argv):
        runName = sys.argv[1]

        mutationList = readFile(os.getcwd() + "/input/mutations.txt")
        proteinList = readFile(os.getcwd() + "/input/protein.txt")
        residueList = readFile(os.getcwd() + "/input/residues.txt")

        if checkMutationList(mutationList, proteinList):

            print("Input mutations are valid; now creating fasta files.")
            saturatedMutationList = saturateList(mutationList, residueList)
            makeDirectory(runName)
            mutateProteinList(saturatedMutationList, proteinList)
            if alphafold(saturatedMutationList):
                prodigy(saturatedMutationList, runName)


if __name__ == "__main__":
    main()
