"""
This module contains some wrappers for using OpenBabel to perform
various operations on BigDFT molecules.

https://open-babel.readthedocs.io/en/latest/UseTheLibrary/Python.html
"""


def convert_atom_to_babel(atom):
    """
    Given a BigDFT atom, this routine creates the open babel
    equivalent.

    Args:
      atom (BigDFT.Atom.Atom): the atom to convert.

    Returns:
      (openbabel.OBAtom): an open babel atom type.
    """
    from openbabel import OBAtom
    a = OBAtom()

    pos = atom.get_position(units="angstroem")
    a.SetVector(pos[0], pos[1], pos[2])
    a.SetAtomicNum(atom.atomic_number)

    if atom.q0 is not None:
        a.SetPartialCharge(atom.q0)

    return a


def convert_frag_to_babel(frag, generate_connectivity=False):
    """
    Convert a BigDFT fragment to an open babel molecule.

    Args:
      frag (BigDFT.Fragments.Fragment): the fragment to convert.
      generate_connectivity (bool): if true, the connectivity will
        be generated by openbabel.

    Returns:
      (openbabel.OBMol): an open babel type molecule.
    """
    from openbabel import OBMol

    # Insert the atoms.
    mol = OBMol()
    for at in frag:
        obat = convert_atom_to_babel(at)
        mol.InsertAtom(obat)

    # Add the bonds.
    if frag.conmat is not None:
        for i in frag.conmat:
            for j in frag.conmat[i]:
                if i > j:
                    mol.AddBond(i+1, j+1, int(frag.conmat[i][j]))
        mol.PerceiveBondOrders()

    return mol


def convert_babel_to_frag(mol):
    """
    Convert an openbabel molecule to a BigDFT fragment.

    Args:
      mol (OBMol): the molecule to compute.

    Returns:
      (BigDFT.Fragments.Fragment): a BigDFT representation of the fragment.
    """
    from BigDFT import Atom as AA
    from BigDFT.Fragments import Fragment
    from openbabel import OBMolAtomIter

    frag = Fragment()
    for obat in OBMolAtomIter(mol):
        sym = AA.number_to_symbol(obat.AtomicNumber())
        pos = [obat.GetX(), obat.GetY(), obat.GetZ()]

        frag += [AA.Atom({sym: pos, "units": "angstroem"})]

    return frag


def compute_fragment_smiles(frag):
    """
    Computes the SMILES representation of a given fragment.

    Args:
      frag (BigDFT.Fragments.Fragment): the fragment to compute the
        representation of.

    Return:
      (str): the smiles representation of this molecule.
    """
    from openbabel import OBConversion

    conv = OBConversion()
    mol = convert_frag_to_babel(frag)
    conv.SetOutFormat("SMI")

    retstr = conv.WriteString(mol)
    retstr = retstr.replace("\n", "")
    retstr = retstr.replace("\t", "")

    return retstr


def compute_fragment_fingerprint(frag, fps="fp2"):
    """
    Computes the fingerprint for a particular fragment.

    Args:
      frag (BigDFT.Fragments.Fragment): the fragment to compute the
        representation of.
      fps (str): the type of finger print to compute.

    Return:
      (openbabel.OBFingerprint): a fingerprint for this fragment.
    """
    from pybel import Molecule

    mol = convert_frag_to_babel(frag)
    pybelmol = Molecule(mol)

    return pybelmol.calcfp(fps)


def optimize_fragment(frag, forcefield="MMFF94", method="SteepestDescent",
                      steps=1000, econv=1e-6, constraints=None):
    """
    Optimize the geometry of a given fragment.

    Args:
      frag (BigDFT.Fragments.Fragment): the fragment to optimize.
      forcefield (str): the type of forcefield to use.
      constraints (list): for each atom, list whether it if frozen or not.

    Returns:
      (BigDFT.Fragments.Fragment): a new fragment with the optimized positions.
    """
    from openbabel import OBForceField, OBMolAtomIter
    from copy import deepcopy

    # Setup the calculation
    ff = OBForceField.FindForceField(forcefield)
    mol = convert_frag_to_babel(frag, generate_connectivity=True)

    ff.Setup(mol)

    if constraints:
        constr = ff.GetConstraints()
        for i in range(0, len(constraints)):
            if constraints[i] is None:
                continue
            elif constraints[i] == "fxyz" or constraints[i] == "f":
                constr.AddAtomConstraint(i+1)
            elif constraints[i] == "fx":
                constr.AddAtomXConstraint(i+1)
            elif constraints[i] == "fy":
                constr.AddAtomYConstraint(i+1)
            elif constraints[i] == "fz":
                constr.AddAtomZConstraint(i+1)
        ff.SetConstraints(constr)

    # Call the optimization routine.
    if method == "SteepestDescent":
        ff.SteepestDescent(steps, econv)
    elif method == "ConjugateGradients":
        ff.ConjugateGradients(steps, econv)
    else:
        raise ValueError("Invalid minimization method.")

    # Extract out the the positions.
    ff.GetCoordinates(mol)
    newfrag = deepcopy(frag)
    for at, obat in zip(newfrag, OBMolAtomIter(mol)):
        at.set_position([obat.GetX(), obat.GetY(), obat.GetZ()],
                        units="angstroem")

    return newfrag


def optimize_system(sys, forcefield="MMFF94", method="SteepestDescent",
                    steps=1000, econv=1e-6,):
    """
    Optimize the geometry of a given system.

    Args:
      sys (BigDFT.Fragments.System): the system to optimize.
      forcefield (str): the type of forcefield to use.

    Returns:
      (BigDFT.Fragments.System): a new system with the optimized positions.
    """
    from BigDFT.Fragments import Fragment
    from copy import deepcopy

    # Convert to one big system.
    fullfrag = Fragment()
    for frag in sys.values():
        for at in frag:
            fullfrag += [at]

    # Compute the connectivity, first the lookup table.
    if sys.conmat is not None:
        counter = 0
        lookup = {}
        for fragid, frag in sys.items():
            for i, at in enumerate(frag):
                lookup[(fragid, i)] = counter
                counter = counter + 1

        # Now the actual connectivity matrix.
        fullfrag.conmat = {}
        for fragid, frag in sys.items():
            for i in range(0, len(frag)):
                t1 = (fragid, i)
                t2list = sys.conmat[fragid][i]
                fullfrag.conmat[lookup[t1]] = {}
                for t2, bo in t2list.items():
                    fullfrag.conmat[lookup[t1]][lookup[t2]] = bo

    # Setup the constraints
    constraint_list = []
    counter = 0
    for frag in sys.values():
        for at in frag:
            constraint_list.append(frag.frozen)
            counter = counter + 1

    # Optimize
    optfrag = optimize_fragment(fullfrag, forcefield=forcefield, method=method,
                                steps=steps, econv=econv,
                                constraints=constraint_list)

    # Convert back
    outsys = deepcopy(sys)
    counter = 0
    for frag in outsys.values():
        for i in range(0, len(frag)):
            frag[i] = optfrag[counter]
            counter = counter + 1

    return outsys
