import sys, os, re  # importiert das pythonmodul sys


class HelixExtraction(object):

    def __init__(self, name, atomIndexNumberBegin,
                 atomIndexNumberEnd):  # hier definiere ich, was mein objekt hat (name, anfang, ende)
        self.name = name
        self.atomIndexNumberBegin = atomIndexNumberBegin
        self.atomIndexNumberEnd = atomIndexNumberEnd
        self.outputStream = None

    def begin(self):
        self.outputStream = open(self.name,
                                 "w")  # w sagt, dass er ne neue dateie schreiben soll --> w ist ein parameter für die Methode open

    def addAtomLine(self, atomLine):
        self.outputStream.write(atomLine)

    def end(self):
        self.outputStream.close()


def readPdbFile(inputdirectoryname, helixPdbFileInputName,
                outputdirectoryname):  # in Klammern kommt der Parameter, den wir einführen können. In iesem Fall der Name der Datei
    print(helixPdbFileInputName)
    helixPdbFile = open(
        inputdirectoryname + "\\" + helixPdbFileInputName).readlines()  # das öffnet die Datei und lädt alle Zeilen in die Variable helixPdbFile
    currentHelixExtractionNumber = 0
    helixExtractionList = []  # hier legen wir register/liste an. Ist ursprünglich ler. Liste mit eckigen Klammern. Durch die Klammer erkennt python, dass es sich um Liste handelt
    currentlyExtracting = False
    readingHelixDefinitions = False
    totalNumberOfHelixExtractions = 0

    for pdbLine in helixPdbFile:  # diese Zeile ittariert über jede Zeile im Dokument
        pdbLineTokens = pdbLine.split()  # wir teilen Zeile in ihre Tokens
        if pdbLineTokens[
            0] == "HELIX":  # wenn das erste Token einer Zeile Helix ist, dann beginnen wir mit dem Ausschneiden der Helix, wenn das erste Token Helix ist --> gehts weiter
            print("Neue Helix gefunden - Nummer: " + str(totalNumberOfHelixExtractions))
            readingHelixDefinitions = True
            helixBegin = pdbLineTokens[5]
            helixEnd = pdbLineTokens[8]
            if re.search('[a-zA-Z]', helixBegin) == None:
                helixBegin = helixBegin + "A"

            if re.search('[a-zA-Z]', helixEnd) == None:
                helixEnd = helixEnd + "A"

            helixExtractionList.append(HelixExtraction(
                outputdirectoryname + "\\" + os.path.splitext(helixPdbFileInputName)[0] + "Helix" + str(
                    totalNumberOfHelixExtractions) + ".pdb", helixBegin,
                helixEnd))  # append fügt unserer Liste ein neues Item hinzu
            totalNumberOfHelixExtractions += 1
        if readingHelixDefinitions and pdbLineTokens[0] != "HELIX":
            print("Alle Helixen Eingelesen - Insgesamt" + str(totalNumberOfHelixExtractions))
            break

    for currentHelixExtraction in helixExtractionList:
        currentHelixExtraction.begin()
        print("Beginne nun mit Extrahieren von Helix: " + currentHelixExtraction.name)
        for pdbLine in helixPdbFile:
            pdbLineTokens = pdbLine.split()
            if pdbLineTokens[0] == "ATOM":
                aminoAcidNumberTokenIndex = 5
                if len(pdbLineTokens[2]) >= 7:
                    aminoAcidNumberTokenIndex = 4
                print(pdbLineTokens)

                if re.search('[a-zA-Z]', pdbLineTokens[aminoAcidNumberTokenIndex]) == None:
                    compareSafeAminoAcidNumber = pdbLineTokens[aminoAcidNumberTokenIndex] + "A"

                compareSafeHelixBegin = currentHelixExtraction.atomIndexNumberBegin
                compareSafeHelixEnd = currentHelixExtraction.atomIndexNumberEnd


                if compareSafeHelixBegin <= compareSafeAminoAcidNumber <= compareSafeHelixEnd:  # mit int mache ich aus Text eine Zahl

                    print(pdbLineTokens)
                    print(currentHelixExtraction.atomIndexNumberBegin, currentHelixExtraction.atomIndexNumberEnd)
                    currentHelixExtraction.addAtomLine(pdbLine)
                elif compareSafeAminoAcidNumber > compareSafeHelixEnd:
                    currentHelixExtraction.end()  # close ist funktion auf datei, end ist funktion auf helix objekt
                    print("fertig mit aktueller Helix")

                    print(currentHelixExtraction.atomIndexNumberBegin, currentHelixExtraction.atomIndexNumberEnd)


if __name__ == '__main__':  # main wird als erstes ausgeführt, wenn python programm gestartet wird
    outputdirectoryname = sys.argv[2]
    for fileName in os.listdir(
            sys.argv[1]):  # sys.arg schaut in den kommandozeilenparameter und wir holen uns den 1. von denen
        if fileName.endswith(".pdb"):
            readPdbFile(sys.argv[1], fileName, outputdirectorynam