# custom imports
from atom import *
import sys
import pymol 
#import collections
class PDBFile:
    def getContent(self, filename):
        with open(filename) as f:
            return(f.readlines())
        
    def getHeader(self):
        #Metadata = collections.namedtuple("Metadata", ["header", "compound", "source", "author"])
        Metadata = {}
        for line in self.rawLines:
            # no need to continue if meta are complete
            if(len(Metadata) <4):
                if(line[0:10] == "HEADER    "):
                    Metadata['header']=line
                elif(line[0:10] == "COMPND   2"):
                    Metadata['compound']=line
                elif(line[0:10] == "SOURCE   2"):
                    Metadata['source']=line
                elif(line[0:10] == "AUTHOR    "):
                    Metadata['author']=line
            else:
                # if meta are complete, stop parsing
                break     
        return(Metadata)

    def getAtoms(self, filename):
        self.atoms = []
        self.residues = []
        temp_atoms = []
        count_h = 0
        for line in self.rawLines:
            if line.startswith("ATOM" or "HETATM"):
                if(line[76:78].strip()=="H"):
                    count_h+=1
                atom = Atom(atom_id = int(line[6:11].strip()),
                            atom_name = line[12:16].strip(),
                            res_name = line[17:20].strip(),
                            chain_id = line[21:22].strip(),
                            res_seq_nb = int(line[22:26].strip()),
                            coordinates = [float(line[30:38].strip()),
                                     float(line[38:46].strip()),
                                     float(line[46:54].strip()),
                            ])
                self.atoms.append(atom)
                # get the current indice of atom
                i = self.atoms.index(atom)
                # if this is a brand new residue
                if(len(self.atoms)>1
                   and atom.res_seq_nb != self.atoms[i-1].res_seq_nb):
                    self.residues.append(Residue(temp_atoms))
                    temp_atoms=[]
                temp_atoms.append(atom)
        # last residue
        self.residues.append(Residue(temp_atoms))
        # hydrogens should represent in average 50% of total atoms... We use 30% threshold...
        if(count_h/len(temp_atoms)<0.30):
            #if(output_pdb==None):
            print("Need to add hydrogens ! If you want the modified PDB file, please use the -o output.pdb argument")
            self.add_hydrogens(filename)

    def check_hydrogens(self, atoms):
        print("ENTER CHECK HYDROGEN")
        return True

    def add_hydrogens(self, filename, output_pdb=None):
        pymol.finish_launching(['pymol', '-qc'])
        pymol.cmd.load(filename)
        pymol.cmd.select("nitrogens",'name n')  
        pymol.cmd.h_add("nitrogens")
        pymol.stored.pos = []
        pymol.cmd.iterate_state(1, "hydrogens", 'stored.pos.append([name,resi,x,y,z])')
        if(output_pdb!=None):
            pymol.cmd.save(output_file)
        return(pymol.stored.pos)
        
    def __init__(self, filename, output_pdb=None):
        self.rawLines = self.getContent(filename)
        self.Metadata = self.getHeader()
        self.getAtoms(filename)
        # for elem in self.Metadata : 
        #     print(self.Metadata[elem], end="")
        

if __name__ == "__main__":
    if(len(sys.argv)<2):
        print("Not enough arguments! Run with --help to learn more about proper"
              "call structure and parameters.")
    else:
        pdbFile = PDBFile(sys.argv[1])
        print(pdbFile.residues[15].atoms["C"].coord_x)
        print(pdbFile.add_hydrogens(sys.argv[1])[3])