#!/usr/bin/env python

import alphaspace2 as al
import mdtraj
import os
import argparse
import logging

def alphaspace_main():
    parser = argparse.ArgumentParser()

    parser.add_argument("-v", "--verbose", action="store_true", help="increase output verbosity")
    parser.add_argument("-p", "--protein", help="protein pdb file", required=True)
    parser.add_argument("-l", "--ligand", help="ligand pdb file", required=False)
    parser.add_argument("-qt", "--pdbqt", help="protein pdbqt file", required=False)
    parser.add_argument("-o", "--output", help="output directory", required=False, default='.')
    parser.add_argument("-H","--keepH", help="remove H at run time, default True", required=False, action='store_true')
    parser.add_argument("-O","--overwrite", help="if overwrite", default=False, action='store_true')

    parser.add_argument("--writepdb", help="whether to write the protein and ligand pdb to the pdb_out folder",
                        action="store_true", default=False)
    parser.add_argument("-c", "--config", help="config file", required=False)

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)

    if args.config:
        import configparser
        config = configparser.ConfigParser()
        config.read(args.config)

    if args.ligand:
        receptor = mdtraj.load(args.protein)
        binder = mdtraj.load(args.ligand)
    else:
        complex = mdtraj.load(args.protein)
        molecules = complex.topology.find_molecules()

        if molecules >= 2:
            logging.debug('No ligand was found, using second molecule in protein file')
            receptor = complex.atom_slice(atom_indices=molecules[0])
            binder = complex.atom_slice(atom_indices=molecules[1])
        else:
            logging.debug('No ligand was found, mapping all protein surface')
            receptor = complex
            binder = None


    if not args.keepH:
        receptor = receptor.atom_slice([atom.index for atom in receptor.topology.atoms if atom.element.symbol != 'H'])

    if args.pdbqt:
        al.annotateVinaAtomTypes(pdbqt=args.pdbqt, receptor=receptor)

    # Initialize a snapshot object, this will contain the receptor and the binder informations
    ss = al.Snapshot()
    ss.beta_cluster_dist = 1.6
    ss.contact_cutoff = 1.6
    ss.pocket_cluster_dist = 4.7
    # Run the snapshot object by feeding it receptor and binder mdtraj objects.
    ss.run(receptor=receptor, binder=binder)

    # Write out files to PDB
    out_path = args.output

    if args.writepdb:
        if not os.path.isdir(os.path.join(out_path, 'pdb_out')):
            os.makedirs(os.path.join(out_path, 'pdb_out'))
        receptor.save(os.path.join(out_path, 'pdb_out', 'prot.pdb'))
        if binder:
            binder.save(os.path.join(out_path, 'pdb_out', 'lig.pdb'))
    al.write_snapshot(folder_path=out_path, snapshot=ss, receptor=receptor)
    al.write_chimera_scripts(folder=out_path)



if __name__ == '__main__':
    alphaspace_main()
