#!/usr/bin/python3
import argparse 

parser = argparse.ArgumentParser(description="Make Report Helper Script")
parser.add_argument("-i","--input",help="Path to report folder")

args = parser.parse_args()

if args.input == None:
    print("Error: Please specify path")
    exit(1)

path = args.input
#path = '/Users/navanchauhan/Desktop/nCOV-19/scripts/pymol/test/'

import untangle
from tabulate import tabulate

#import sys
#report = path + "report.md"
#sys.stdout = open(report, 'w')

from os import listdir
from os.path import isfile, join
onlyfiles = [f for f in listdir(path) if isfile(join(path, f))]
image = ""
for x in onlyfiles:
    if '.png' in x and 'UNL' in x:
        image = x
import os
fname = os.path.join(path,'report.xml')

doc = untangle.parse(fname)

hi, hb, wb, sb, ps, pc, hab, mc = 0,0,0,0,0,0,0,0

indexForUNL = 0

for x in doc.report.bindingsite:
    if x.identifiers.longname.cdata == 'UNL':
        break
    else:
        indexForUNL += 1


name = doc.report.pdbid.cdata
#print(("# " + (name.replace("_"," ")).replace("PROTEIN","")), end="\n\n")

print("## Visualisation", end="\n\n")    
print(f'![]({image})', end="\n\n")

print("## Interactions", end="\n\n")

fallback = 0
try:
    if doc.report.bindingsite[indexForUNL].interactions.hydrophobic_interactions.cdata == '':
        print("No Hydrophobic Interactions Found", end="\n\n")
    else:
        print("**Hydrophobic Interactions Found**", end="\n\n")
        hi = 1
except AttributeError:
    fallback=1

if fallback==0:
    if doc.report.bindingsite[indexForUNL].interactions.hydrophobic_interactions.cdata == '':
        print("No Hydrophobic Interactions Found", end="\n\n")
    else:
        print("**Hydrophobic Interactions Found**", end="\n\n")
        hi = 1
    if doc.report.bindingsite[indexForUNL].interactions.hydrogen_bonds.cdata == '':
        print("No Hydrogen Bonds Found", end="\n\n")
    else:
        print("**Hydrogen Bonds Found**", end="\n\n")
        hb = 1
    if doc.report.bindingsite[indexForUNL].interactions.water_bridges.cdata == '':
        print("No Water Bridges Found", end="\n\n")
    else:
        print("**Water Bridges Found**", end="\n\n")
        wb = 1
    if doc.report.bindingsite[indexForUNL].interactions.salt_bridges.cdata == '':
        print("No Salt Bridges Found", end="\n\n")
    else:
        print("**Salt Bridges Found**", end="\n\n")
        sb = 1
    if doc.report.bindingsite[indexForUNL].interactions.pi_stacks.cdata == '':
        print("No Pi Stacks Found", end="\n\n")
    else:
        print("**Pi Stacks Found**", end="\n\n")
        ps = 1
    if doc.report.bindingsite[indexForUNL].interactions.pi_cation_interactions.cdata == '':
        print("No Pi Cation Interactions Found", end="\n\n")
    else:
        print("**Pi Cation Interactions Found**", end="\n\n")
        pc = 1
    if doc.report.bindingsite[indexForUNL].interactions.halogen_bonds.cdata == '':
        print("No Halogen Bonds Found", end="\n\n")
    else:
        print("** Halogen Bonds Found**", end="\n\n")
        hab = 1
    if doc.report.bindingsite[indexForUNL].interactions.metal_complexes.cdata == '':
        print("No Metal Complexes Found", end="\n\n")    
    else:
        print("**Metal Complexes Found**", end="\n\n")
        mc = 1

if fallback == 1:
    if doc.report.bindingsite.interactions.hydrophobic_interactions.cdata == '':
        print("No Hydrophobic Interactions Found", end="\n\n")
    else:
        print("**Hydrophobic Interactions Found**", end="\n\n")
        hi = 1
    if doc.report.bindingsite.interactions.hydrogen_bonds.cdata == '':
        print("No Hydrogen Bonds Found", end="\n\n")
    else:
        print("**Hydrogen Bonds Found**", end="\n\n")
        hb = 1
    if doc.report.bindingsite.interactions.water_bridges.cdata == '':
        print("No Water Bridges Found", end="\n\n")
    else:
        print("**Water Bridges Found**", end="\n\n")
        wb = 1
    if doc.report.bindingsite.interactions.salt_bridges.cdata == '':
        print("No Salt Bridges Found", end="\n\n")
    else:
        print("**Salt Bridges Found**", end="\n\n")
        sb = 1
    if doc.report.bindingsite.interactions.pi_stacks.cdata == '':
        print("No Pi Stacks Found", end="\n\n")
    else:
        print("**Pi Stacks Found**", end="\n\n")
        ps = 1
    if doc.report.bindingsite.interactions.pi_cation_interactions.cdata == '':
        print("No Pi Cation Interactions Found", end="\n\n")
    else:
        print("**Pi Cation Interactions Found**", end="\n\n")
        pc = 1
    if doc.report.bindingsite.interactions.halogen_bonds.cdata == '':
        print("No Halogen Bonds Found", end="\n\n")
    else:
        print("** Halogen Bonds Found**", end="\n\n")
        hab = 1
    if doc.report.bindingsite.interactions.metal_complexes.cdata == '':
        print("No Metal Complexes Found", end="\n\n")    
    else:
        print("**Metal Complexes Found**", end="\n\n")
        mc = 1

if fallback == 0:
    if hi == 1:
        print("## Hydrophobic Interactions",end="\n\n")
        tableBody = []
        tableHeaders = ['No.','Res.','AA','Dist','Ligand Atom','Proton Atom']
        i = 1
        for x in doc.report.bindingsite[indexForUNL].interactions.hydrophobic_interactions.hydrophobic_interaction:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist.cdata)
            l.append(x.ligcarbonidx.cdata)
            l.append(x.protcarbonidx.cdata)
            i += 1
            tableBody.append(l)
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

    if hb == 1:
        print("## Hydrogen Bonds",end="\n\n")
        tableBody = []
        tableHeaders = ['No.','Res.','AA','Dist H-A','Dist D-A','Don Angle','Protisdon?','Sidechain?','D. Atom','A. Atom']
        i = 1
        for x in doc.report.bindingsite[indexForUNL].interactions.hydrogen_bonds.hydrogen_bond:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist_h_a.cdata)
            l.append(x.dist_d_a.cdata)
            l.append(x.don_angle.cdata)
            l.append(x.protisdon.cdata)
            l.append(x.sidechain.cdata)
            l.append((x.donoridx.cdata + "[" + x.donortype.cdata + "]"))
            l.append((x.acceptoridx.cdata + "[" + x.acceptortype.cdata + "]"))
            i += 1
            tableBody.append(l)
            #print(i, x.resnr.cdata, x.restype.cdata, x.dist_h_a.cdata, x.dist_d_a.cdata, x.don_angle.cdata, x.protisdon.cdata, x.sidechain.cdata, (x.donoridx.cdata + "[" + x.donortype.cdata + "]"), (x.acceptoridx.cdata + "[" + x.acceptortype.cdata + "]"))
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

    if sb == 1:
        print("## Salt Bridges",end="\n\n")
        tableBody = []
        tableHeaders = ['Index','Residue','AA','Distance','Protein positive?','Ligand Group','Ligand Atoms']
        i = 1
        for x in doc.report.bindingsite[indexForUNL].interactions.salt_bridges.salt_bridge:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist.cdata)
            l.append(x.protispos.cdata)
            l.append(x.lig_group.cdata)
            atoms = []
            for y in x.lig_idx_list.idx:
                atoms.append(y.cdata)
            l.append(atoms)
            i += 1
            tableBody.append(l)
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

    if pc==1:
        print("## Pi Cation Interactions",end="\n\n")
        tableBody = []
        tableHeaders = ['Index','Residue','AA','Distance','Prot charged?','Atoms']
        i = 1
        for x in doc.report.bindingsite[indexForUNL].interactions.pi_cation_interactions.pi_cation_interaction:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist.cdata)
            l.append(x.offset.cdata)
            l.append(x.protcharged.cdata)
            atoms = []
            for y in x.lig_idx_list.idx:
                atoms.append(y.cdata)
            l.append(atoms)
            i += 1
            tableBody.append(l)
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")
elif fallback == 1:
    if hi == 1:
        print("## Hydrophobic Interactions",end="\n\n")
        tableBody = []
        tableHeaders = ['No.','Res.','AA','Dist','Ligand Atom','Proton Atom']
        i = 1
        for x in doc.report.bindingsite.interactions.hydrophobic_interactions.hydrophobic_interaction:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist.cdata)
            l.append(x.ligcarbonidx.cdata)
            l.append(x.protcarbonidx.cdata)
            i += 1
            tableBody.append(l)
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

    if hb == 1:
        print("## Hydrogen Bonds",end="\n\n")
        tableBody = []
        tableHeaders = ['No.','Res.','AA','Dist H-A','Dist D-A','Don Angle','Protisdon?','Sidechain?','D. Atom','A. Atom']
        i = 1
        for x in doc.report.bindingsite.interactions.hydrogen_bonds.hydrogen_bond:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist_h_a.cdata)
            l.append(x.dist_d_a.cdata)
            l.append(x.don_angle.cdata)
            l.append(x.protisdon.cdata)
            l.append(x.sidechain.cdata)
            l.append((x.donoridx.cdata + "[" + x.donortype.cdata + "]"))
            l.append((x.acceptoridx.cdata + "[" + x.acceptortype.cdata + "]"))
            i += 1
            tableBody.append(l)
            #print(i, x.resnr.cdata, x.restype.cdata, x.dist_h_a.cdata, x.dist_d_a.cdata, x.don_angle.cdata, x.protisdon.cdata, x.sidechain.cdata, (x.donoridx.cdata + "[" + x.donortype.cdata + "]"), (x.acceptoridx.cdata + "[" + x.acceptortype.cdata + "]"))
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

    if sb == 1:
        print("## Salt Bridges",end="\n\n")
        tableBody = []
        tableHeaders = ['Index','Residue','AA','Distance','Protein positive?','Ligand Group','Ligand Atoms']
        i = 1
        for x in doc.report.bindingsite.interactions.salt_bridges.salt_bridge:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist.cdata)
            l.append(x.protispos.cdata)
            l.append(x.lig_group.cdata)
            atoms = []
            for y in x.lig_idx_list.idx:
                atoms.append(y.cdata)
            l.append(atoms)
            i += 1
            tableBody.append(l)
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

    if pc==1:
        print("## Pi Cation Interactions",end="\n\n")
        tableBody = []
        tableHeaders = ['Index','Residue','AA','Distance','Prot charged?','Atoms']
        i = 1
        for x in doc.report.bindingsite.interactions.pi_cation_interactions.pi_cation_interaction:
            l = []
            l.append(i)
            l.append(x.resnr.cdata)
            l.append(x.restype.cdata)
            l.append(x.dist.cdata)
            l.append(x.offset.cdata)
            l.append(x.protcharged.cdata)
            atoms = []
            for y in x.lig_idx_list.idx:
                atoms.append(y.cdata)
            l.append(atoms)
            i += 1
            tableBody.append(l)
        print(tabulate(tableBody, headers=tableHeaders), end="\n\n")

print("## Figures", end="\n\n")

print(f'![](output-back.png)', end="\n\n")
print(f'![](output-front.png)', end="\n\n")
print(f'![](closeup-back.png)', end="\n\n")
print(f'![](closeup-front.png)', end="\n\n")