#!@PYTHON@

# Compoute bond distance
#
# For a given set of elements we go through all pairs and compute the bond length after
# an optimization with a fixed number of steps.


import re, os, os.path, sys, operator
import pyMoleCuilder as mol
from numpy.linalg import norm

import argparse, sys

STEPS = 50

mol.CommandVerbose('1')

parser = argparse.ArgumentParser()
parser.add_argument("--basis-set", type=str, default="6-31G", \
    help="Sets the basis set with mpqc for the CLHF ab-initio computations")
parser.add_argument("elements", nargs=2, \
    help="The pair of elements, e.g., Li H")
parser.add_argument("--output-prefix", type=str, default="", required=True, \
    help="Prefix for output filenames of created pairs")
parser.add_argument("--server-address", type=str, default="", required=True, \
    help="Hostname of the molecuilder_server")
parser.add_argument("--server-port", type=str, default="", required=True, \
    help="Port of the molecuilder_server")
parser.add_argument('--version', '-V', action="store_true", \
    help='Gives version information')

params, _ = parser.parse_known_args()

if params.version:
    # give version and exit
    print(sys.argv[0]+" -- version 0.1")
    sys.exit(0)

elem1 = params.elements[0]
elem2 = params.elements[1]

def createPair(elem1, elem2):
    # we put them in a default distance of 1.5A
    mol.AtomAdd(domain_position = "9.25,10,10", add_atom = elem1)
    mol.SelectionAllMolecules() # enforce adding to present molecule
    mol.AtomAdd(domain_position = "10.75,10,10", add_atom = elem2)
    mol.SelectionAllAtoms()
    mol.BondAdd()
    mol.AtomSaturate(use_outer_shell = "1")
    mol.SelectionClearAllMolecules()

def optimizeGeometry(elem1, elem2):
    prefix = "BondFragment-"+elem1+"_"+elem2
    mol.FragmentationClearFragmentationResults()
    mol.PotentialClearHomologies()
    # create the keyset filen
    mol.SelectionAllAtoms()
    mol.SelectionNotAtomByElement("H")
    with open("./BondFragmentKeySets.dat", "w") as f:
        for i in range(mol.getSelectedAtomCount()):
            f.write(str(i)+" ");
        f.write("\n");
    mol.SelectionAllAtoms()
    mol.wait()
    for step in range(STEPS):
        mol.WorldOutputAs(output_as=params.output_prefix + "-" + elem1 + "_" + elem2 + ".in", force_overwrite="1")
        # parse the full molecule as one single fragment job
        mol.FragmentationAddSelectedAtomsAsFragment(
            add_selected_atoms_as_fragment=params.output_prefix + "-" + elem1 + "_" + elem2,
            output_types=""
        )
        status = mol.wait()
        mol.FragmentationFragmentationAutomation(
            DoLongrange="0",
            server_address=str(params.server_address),
            server_port=str(params.server_port)
        )
        status = mol.wait()
        if not status:
            print("Fragmentation computation failed, STOPPING.", flush=True)
            return False
        mol.FragmentationAnalyseFragmentationResults(
            fragment_prefix=prefix
        )
        mol.wait()
        mol.WorldStepWorldTime("1")
        mol.wait()
        mol.MoleculeForceAnnealing(
            use_bondgraph="1"
        )
        mol.FragmentationClearFragmentationResults()
    return True

def extractNonHydrogenBondLength(elem1, elem2):
    name="length-" + elem1 + "_" + elem2
    mol.SelectionClearAllAtoms()
    if elem1 != "H":
        mol.SelectionAtomByElement(elem1)
    else:
        mol.SelectionAtomByName(elem1+"1")
    if elem2 != "H":
        mol.SelectionAtomByElement(elem2)
    else:
        mol.SelectionAtomByName(elem2+"2")
    mol.GeometryDistanceToVector(distance_to_vector=name)
    mol.wait()
    return norm(mol.get_geometryobject(name), 2)

# original list of element for pairing
elements=["H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg", "Al", "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe"]

# list with elements excluded where force calculation did not work: noble gases, Mn,Fe with H need fine-tuning with smaller step width
elements=["H", "Li", "Be", "B", "C", "N", "O", "F", "Na", "Mg", "Al", "Si", "P", "S", "Cl", "K", "Ca", "Sc", "Ti", "Mn", "Sc", "Ti", "V", "Cr", "Mn", "Fe"]
dysfunctional_pairs=[ "Sc_Ti" ]

# list for testing the code
#elements=["H"] 

mol.ParserSetTremoloAtomdata("Id type x=3 F=3 neighbors=4")
mol.ParserSetParserParameters(set_parser_parameters="mpqc", parser_parameters="theory=CLHF;basis=" + params.basis_set + ";")

element_pair = elem1+"_"+elem2
if element_pair in dysfunctional_pairs:
    print("The combination " + element_pair + " cannot be computed, SKIPPING.", flush=True)
    length = "SKIPPED"
else:
    createPair(elem1, elem2)
    mol.wait()
    status = optimizeGeometry(elem1, elem2)
    if status:
        length = extractNonHydrogenBondLength(elem1, elem2)
        mol.wait()
    else:
        length = "SKIPPED"
print("DISTANCE:   " + elem1 + "  " + elem2 + ": " + str(length), flush=True)

mol.SelectionAllAtoms()
mol.WorldOutputAs(output_as=params.output_prefix + "-" + elem1 + "_" + elem2 + ".data", force_overwrite="1")
mol.wait()