#!/usr/bin/python3
#
# pKa calculations with APBS
#
# Copyright University College Dublin & Washington University St. Louis 2004-2007
# All rights reserved
#

#
# Get paths
#
debug=False
import optparse
import sys, os

from src.definitions import Definition
from src.forcefield import Forcefield
from src.routines import Routines
from src.protein import getPDBFile, readPDB, Protein, Amino, Nucleic
from src.aa import LIG
from src.errors import PDB2PKAError
from pdb2pka import inputgen_pKa

from pdb2pka.pka_routines import smooth
from pdb2pka import pka_help

def usage(x):
    #
    # Print usage guidelines
    #
    print('Usage: pka.py --ff <forcefield> --lig <ligand in MOL2> --pdie <protein diel cons> --maps <1 for using provided 3D maps; 2 for genereting new maps>')
    print('--xdiel <xdiel maps> --ydiel <ydiel maps> --zdiel <zdiel maps> --kappa <ion-accessibility map> ')
    print('--smooth <st.dev [A] of Gaussian smooting of 3D maps at the boundary, bandthwith=3 st.dev> <pdbfile>')
    print('Force field can be amber, charmm and parse')
    print()
    return

#
# --------------------------------------------------
#

def startpKa():
    """
        Function for starting pKa script from the command line.

        Returns
            protein:    The protein object as generated by PDB2PQR
            routines:   The routines object as generated by PDB2PQR
            forcefield: The forcefield object as generated by PDB2PQR
    """
    print()
    print('PDB2PQR pKa calculations')
    print()

    parser = optparse.OptionParser()

    ##
    ## set optparse options
    ##
    parser.add_option(
        '-v','--verbose',
        dest='verbose',
        action="store_true",
        default=False,
        )
    parser.add_option(
        '--pdie',
        dest='pdie',
        default=8,
        type='int',
        help='<protein dielectric constant>',
        )
    parser.add_option(
        '--sdie',
        dest='sdie',
        default=80,
        type='int',
        help='<solvent dielectric constant>',
        )
    parser.add_option(
        '--ff',
        dest='ff',
        type='choice',
        default='parse',
        choices=("amber","AMBER","charmm","CHARMM","parse","PARSE",),
        help='<force field (amber, charmm, parse)>',
        )
    parser.add_option(
        '--resume',
        dest='resume',
        action="store_true",
        default=False,
        help='resume run from saved state.',
        )
    parser.add_option(
        '--ligand',
        dest='ligand',
        type='str',
        help='<ligand in MOL2 format>',
        )
    parser.add_option(
        '--maps',
        dest='maps',
        default=None,
        type='int',
        help='<1 for using provided 3D maps; 2 for genereting new maps>',
        )
    parser.add_option(
        '--xdiel',
        dest='xdiel',
        default=None,
        type='str',
        help='<xdiel maps>',
        )
    parser.add_option(
        '--ydiel',
        dest='ydiel',
        default=None,
        type='str',
        help='<ydiel maps>',
        )
    parser.add_option(
        '--zdiel',
        dest='zdiel',
        default=None,
        type='str',
        help='<zdiel maps>',
        )
    parser.add_option(
        '--kappa',
        dest='kappa',
        default=None,
        type='str',
        help='<ion-accessibility map>',
        )
    parser.add_option(
        '--smooth',
        dest='sd',
        default=None,
        type='float',
        help='<st.dev [A] of Gaussian smooting of 3D maps at the boundary, bandthwith=3 st.dev>',
        )
    #
    # Cut off energy for calculating non-charged-charged interaction energies
    #
    parser.add_option('--pairene',dest='pairene',type='float',default=1.0,
                      help='Cutoff energy in kT for calculating non charged-charged interaction energies. Default: %default')
    #
    # Options for doing partial calculations
    #
    parser.add_option('--res_energy',
                      dest='desolvation_res',
                      default=[],
                      action='append',
                      type='string',
                      help='Calculate desolvation energy and interaction energy for this residue in its default protonation state. Protonation states can be specified with the --protonation_state argument')
    parser.add_option('--PS_file',dest='PS_file',default='',type='string',action='store',help='Set protonation states according to the pdb2pka protonation state file (option --PS_file)')
    (options,args,) = parser.parse_args()

    ##
    ## parse optparse options
    ##
    ff = options.ff.lower()
    pdie = options.pdie
    verbose = options.verbose
    sdie = options.sdie
    maps = options.maps
    xdiel = options.xdiel
    ydiel = options.ydiel
    zdiel = options.zdiel
    kappa = options.kappa
    sd = options.sd

    #
    # Find the PDB file
    #
    if len(args) != 2:
        parser.error("Usage: pka.py [options] <pdbfile> <output directory>\n")
    input_path = args[0]
    output_path = args[1]

    ligand = None
    if options.ligand is not None:
        try:
            ligand = open(options.ligand, 'rU')
        except IOError:
            print('Unable to find ligand file %s! Skipping...' % options.ligand)

    #Set up the protien object
    #In the standalone version of pdb2pka this is redundent but needed so we emulate the
    #interface needed by pdb2pqr

    pdbfile = getPDBFile(input_path)
    pdblist, errlist = readPDB(pdbfile)
    if len(errlist) != 0 and verbose:
        print("Warning: %s is a non-standard PDB file.\n" %input_path)
        print(errlist)
    #
    # Read the definition file
    #
    myDefinition = Definition()
    #
    #
    # Choose whether to include the ligand or not
    #
    # Add the ligand to the pdb2pqr arrays
    #
    if ligand is None:
        myProtein = Protein(pdblist, myDefinition)
    else:
        from pdb2pka.ligandclean import ligff
        myProtein, _, _ = ligff.initialize(myDefinition, ligand, pdblist, verbose)

    #
    # Call the pre_init function
    #
    return pre_init(protein=myProtein,
                    output_dir=output_path,
                    ff=ff,
                    verbose=verbose,
                    pdie=pdie,
                    sdie=sdie,
                    maps=maps,
                    xdiel=xdiel,
                    ydiel=ydiel,
                    zdiel=zdiel,
                    kappa=kappa,
                    sd=sd,
                    ligand=ligand),options


def pre_init(original_pdb_list=None,
             output_dir=None,
             ff=None,
             verbose=False,
             pdie=8.0,
             sdie=80,
             maps=None,
             xdiel=None,
             ydiel=None,
             zdiel=None,
             kappa=None,
             sd=None,
             ligand=None):
    """This function cleans the PDB and prepares the APBS input file

    Prepares the output folder."""

    #prepare the output directory

    output_dir = os.path.abspath(output_dir)

    try:
        os.makedirs(output_dir)
    except OSError:
        if not os.path.isdir(output_dir):
            raise ValueError('Target directory is a file! Aborting.')

    workspace_dir = os.path.join(output_dir,'workspace')

    try:
        os.makedirs(workspace_dir)
    except OSError:
        if not os.path.isdir(output_dir):
            raise ValueError('Target directory is a file! Aborting.')

    #
    # remove hydrogen atoms
    #

    working_pdb_filename = os.path.join(workspace_dir,'working.pdb')

    pka_help.dump_protein_no_hydrogens(original_pdb_list, working_pdb_filename)
    #
    # Get the PDBfile
    #
    pdbfile = getPDBFile(working_pdb_filename)
    pdblist, errlist = readPDB(pdbfile)

    if verbose:
        print("Beginning PDB2PKA...\n")
    #
    # Read the definition file
    #
    myDefinition = Definition()
    ligand_titratable_groups=None
    #
    #
    # Choose whether to include the ligand or not
    #
    # Add the ligand to the pdb2pqr arrays
    #
    Lig=None
    if ligand is None:
        myProtein = Protein(pdblist, myDefinition)
    else:
        from pdb2pka.ligandclean import ligff
        myProtein, myDefinition, Lig = ligff.initialize(myDefinition, ligand, pdblist, verbose)
    #
    # =======================================================================
    #
    # We have identified the structural elements, now contiue with the setup
    #
    # Print something for some reason?
    #
    if verbose:
        print("Created protein object -")
        print("\tNumber of residues in protein: %s" % myProtein.numResidues())
        print("\tNumber of atoms in protein   : %s" % myProtein.numAtoms())
    #
    # Set up all other routines
    #
    myRoutines = Routines(myProtein, verbose) #myDefinition)
    myRoutines.updateResidueTypes()
    myRoutines.updateSSbridges()
    myRoutines.updateBonds()
    myRoutines.setTermini()
    myRoutines.updateInternalBonds()

    myRoutines.applyNameScheme(Forcefield(ff, myDefinition, None))
    myRoutines.findMissingHeavy()
    myRoutines.addHydrogens()
    myRoutines.debumpProtein()

    #myRoutines.randomizeWaters()
    myProtein.reSerialize()
    #
    # Inject the information on hydrogen conformations in the HYDROGENS.DAT arrays
    # We get this information from ligand_titratable_groups
    #
    from src.hydrogens import hydrogenRoutines
    myRoutines.updateInternalBonds()
    myRoutines.calculateDihedralAngles()
    myhydRoutines = hydrogenRoutines(myRoutines)
    #
    # Here we should inject the info!!
    #
    myhydRoutines.setOptimizeableHydrogens()
    myhydRoutines.initializeFullOptimization()
    myhydRoutines.optimizeHydrogens()
    myhydRoutines.cleanup()
    myRoutines.setStates()

    #
    # Choose the correct forcefield
    #
    myForcefield = Forcefield(ff, myDefinition, None)
    if Lig:
        hitlist, misslist = myRoutines.applyForcefield(myForcefield)
        #
        # Can we get charges for the ligand?
        #
        templist=[]
        ligsuccess=False
        for residue in myProtein.getResidues():
            if isinstance(residue, LIG):
                templist = []
                Lig.make_up2date(residue)
                net_charge=0.0
                print('Ligand',residue)
                print('Atom\tCharge\tRadius')
                for atom in residue.getAtoms():
                    if atom.mol2charge:
                        atom.ffcharge=atom.mol2charge
                    else:
                        atom.ffcharge = Lig.ligand_props[atom.name]["charge"]
                    #
                    # Find the net charge
                    #
                    net_charge=net_charge+atom.ffcharge
                    #
                    # Assign radius
                    #
                    atom.radius = Lig.ligand_props[atom.name]["radius"]
                    print('%s\t%6.4f\t%6.4f' %(atom.name,atom.ffcharge,atom.radius))
                    if atom in misslist:
                        misslist.pop(misslist.index(atom))
                        templist.append(atom)
                    #
                    # Store the charge and radius in the atom instance for later use
                    # This really should be done in a nicer way, but this will do for now
                    #
                    atom.secret_radius=atom.radius
                    atom.secret_charge=atom.ffcharge
                    #
                    #

                charge = residue.getCharge()
                if abs(charge - round(charge)) > 0.01:
                    # Ligand parameterization failed
                    myProtein.residues.remove(residue)
                    raise Exception('Non-integer charge on ligand: %8.5f' %charge)
                else:
                    ligsuccess = 1
                    # Mark these atoms as hits
                    hitlist = hitlist + templist
                #
                # Print the net charge
                #
                print('Net charge for ligand %s is: %5.3f' %(residue.name,net_charge))
        #
        # Temporary fix; if ligand was successful, pull all ligands from misslist
        # Not sure if this is needed at all here ...? (Jens wrote this)
        #
        if ligsuccess:
            templist = misslist[:]
            for atom in templist:
                if isinstance(atom.residue, Amino) or isinstance(atom.residue, Nucleic):
                    continue
                misslist.remove(atom)

    if verbose:
        print("Created protein object (after processing myRoutines) -")
        print("\tNumber of residues in protein: %s" % myProtein.numResidues())
        print("\tNumber of atoms in protein   : %s" % myProtein.numAtoms())
    #
    # Create the APBS input file
    #
    import src.psize
    size=src.psize.Psize()

    method=""
    split=0

    igen = inputgen_pKa.inputGen(working_pdb_filename)
    #
    # For convenience
    #
    igen.pdie = pdie
    print('Setting protein dielectric constant to ',igen.pdie)
    igen.sdie=sdie
    igen.maps=maps
    if maps==1:
        print("Using dielectric and mobile ion-accessibility function maps in PBE")
        if xdiel:
            igen.xdiel = xdiel
        else:
            raise PDB2PKAError('X dielectric map is missing')
        if ydiel:
            igen.ydiel = ydiel
        else:
            raise PDB2PKAError("Y dielectric map is missing\n")
        if zdiel:
            igen.zdiel = zdiel
        else:
            raise PDB2PKAError("Z dielectric map is missing\n")

        print('Setting dielectric function maps: %s, %s, %s'%(igen.xdiel,igen.ydiel,igen.zdiel))

        if kappa:
            igen.kappa = kappa
        else:
            raise PDB2PKAError("Mobile ion-accessibility map is missing\n")

        print('Setting mobile ion-accessibility function map to: ',igen.kappa)

        if sd:
            xdiel_smooth, ydiel_smooth, zdiel_smooth = smooth(xdiel,ydiel,zdiel)
            igen.xdiel = xdiel_smooth
            igen.ydiel = ydiel_smooth
            igen.zdiel = zdiel_smooth
    #
    # Return all we need
    #
    return output_dir, myProtein, myRoutines, myForcefield,igen, ligand_titratable_groups, maps, sd

#
# --------------
#

if __name__ == "__main__":

    (output_dir, protein, routines, forcefield,apbs_setup, ligand_titratable_groups, maps, sd), options = startpKa()
    from pdb2pka import pka_routines
    mypkaRoutines = pka_routines.pKaRoutines(protein, routines, forcefield, apbs_setup, output_dir, maps, sd,
                                             restart=not options.resume, pairene=options.pairene)

    print('Doing full pKa calculation')
    mypkaRoutines.runpKa()