<commit_before>"""PROPKA calculations."""
import math
import propka.protonate
import propka.bonds
from propka.lib import warning


# TODO - this file should be broken into three separate files:
# * calculations.py - includes basic functions for calculating distances, etc.
# * hydrogen.py - includes bonding and protonation functions
# * energy.py - includes energy functions (dependent on distance functions)


# TODO - the next set of functions form a distinct "module" for distance calculation


# Maximum distance used to bound calculations of smallest distance
MAX_DISTANCE = 1e6


def squared_distance(atom1, atom2):
    """Calculate the squared distance between two atoms.

    Args:
        atom1:  first atom for distance calculation
        atom2:  second atom for distance calculation
    Returns:
        distance squared
    """
    dx = atom2.x - atom1.x
    dy = atom2.y - atom1.y
    dz = atom2.z - atom1.z
    res = dx*dx+dy*dy+dz*dz
    return res


def distance(atom1, atom2):
    """Calculate the distance between two atoms.

    Args:
        atom1:  first atom for distance calculation
        atom2:  second atom for distance calculation
    Returns:
        distance
    """
    return math.sqrt(squared_distance(atom1, atom2))


def get_smallest_distance(atoms1, atoms2):
    """Calculate the smallest distance between two groups of atoms.

    Args:
        atoms1:  atom group 1
        atoms2:  atom group 2
    Returns:
        smallest distance between groups
    """
    res_dist = MAX_DISTANCE
    res_atom1 = None
    res_atom2 = None
    for atom1 in atoms1:
        for atom2 in atoms2:
            dist = squared_distance(atom1, atom2)
            if dist < res_dist:
                res_dist = dist
                res_atom1 = atom1
                res_atom2 = atom2
    return [res_atom1, math.sqrt(res_dist), res_atom2]


# TODO - the next set of functions form a distinct "module" for hydrogen addition


def setup_bonding_and_protonation(parameters, molecular_container):
    """Set up bonding and protonation for a molecule.

    NOTE - the unused `parameters` argument is required for compatibility in version.py
    TODO - figure out why there is a similar function in version.py

    Args:
        parameters:  not used
        molecular_container:  molecule container.
    """
    # make bonds
    my_bond_maker = setup_bonding(molecular_container)
    # set up ligand atom names
    set_ligand_atom_names(molecular_container)
    # apply information on pi electrons
    my_bond_maker.add_pi_electron_information(molecular_container)
    # Protonate atoms
    if molecular_container.options.protonate_all:
        PROTONATOR = propka.protonate.Protonate(verbose=False)
        PROTONATOR.protonate(molecular_container)


def setup_bonding(molecular_container):
    """Set up bonding for a molecular container.

    TODO - figure out why there is a similar function in version.py

    Args:
        molecular_container:  the molecular container in question
    Returns:
        BondMaker object
    """
    my_bond_maker = propka.bonds.BondMaker()
    my_bond_maker.find_bonds_for_molecules_using_boxes(molecular_container)
    return my_bond_maker


def set_ligand_atom_names(molecular_container):
    """Set names for ligands in molecular container.

    Args:
        molecular_container:  molecular container for ligand names
    """
    for name in molecular_container.conformation_names:
        molecular_container.conformations[name].set_ligand_atom_names()


def add_arg_hydrogen(residue):
    """Adds Arg hydrogen atoms to residues according to the 'old way'.

    Args:
        residue:  arginine residue to protonate
    Returns:
        list of hydrogen atoms<commit_msg>Restore 3.0 compatibility functions.<commit_after>"""PROPKA calculations."""
import math
import propka.protonate
import propka.bonds
from propka.lib import warning, info


# TODO - this file should be broken into three separate files:
# * calculations.py - includes basic functions for calculating distances, etc.
# * hydrogen.py - includes bonding and protonation functions
# * energy.py - includes energy functions (dependent on distance functions)


# TODO - the next set of functions form a distinct "module" for distance calculation


# Maximum distance used to bound calculations of smallest distance
MAX_DISTANCE = 1e6


def squared_distance(atom1, atom2):
    """Calculate the squared distance between two atoms.

    Args:
        atom1:  first atom for distance calculation
        atom2:  second atom for distance calculation
    Returns:
        distance squared
    """
    dx = atom2.x - atom1.x
    dy = atom2.y - atom1.y
    dz = atom2.z - atom1.z
    res = dx*dx+dy*dy+dz*dz
    return res


def distance(atom1, atom2):
    """Calculate the distance between two atoms.

    Args:
        atom1:  first atom for distance calculation
        atom2:  second atom for distance calculation
    Returns:
        distance
    """
    return math.sqrt(squared_distance(atom1, atom2))


def get_smallest_distance(atoms1, atoms2):
    """Calculate the smallest distance between two groups of atoms.

    Args:
        atoms1:  atom group 1
        atoms2:  atom group 2
    Returns:
        smallest distance between groups
    """
    res_dist = MAX_DISTANCE
    res_atom1 = None
    res_atom2 = None
    for atom1 in atoms1:
        for atom2 in atoms2:
            dist = squared_distance(atom1, atom2)
            if dist < res_dist:
                res_dist = dist
                res_atom1 = atom1
                res_atom2 = atom2
    return [res_atom1, math.sqrt(res_dist), res_atom2]


# TODO - the next set of functions form a distinct "module" for hydrogen addition


def setup_bonding_and_protonation(parameters, molecular_container):
    """Set up bonding and protonation for a molecule.

    NOTE - the unused `parameters` argument is required for compatibility in version.py
    TODO - figure out why there is a similar function in version.py

    Args:
        parameters:  not used
        molecular_container:  molecule container.
    """
    # make bonds
    my_bond_maker = setup_bonding(molecular_container)
    # set up ligand atom names
    set_ligand_atom_names(molecular_container)
    # apply information on pi electrons
    my_bond_maker.add_pi_electron_information(molecular_container)
    # Protonate atoms
    if molecular_container.options.protonate_all:
        PROTONATOR = propka.protonate.Protonate(verbose=False)
        PROTONATOR.protonate(molecular_container)


def setup_bonding(molecular_container):
    """Set up bonding for a molecular container.

    Args:
        molecular_container:  the molecular container in question
    Returns:
        BondMaker object
    """
    my_bond_maker = propka.bonds.BondMaker()
    my_bond_maker.find_bonds_for_molecules_using_boxes(molecular_container)
    return my_bond_maker


def setup_bonding_and_protonation_30_style(parameters, molecular_container):
    """Set up bonding for a molecular container.

    Args:
        parameters:  parameters for calculation
        molecular_container:  the molecular container in question
    Returns:
        BondMaker object
    """
    # Protonate atoms
    protonate_30_style(molecular_container)
    # make bonds
    bond_maker = propka.bonds.BondMaker()
    bond_maker.find_bonds_for_molecules_using_boxes(molecular_container)
    return bond_maker


def protonate_30_style(molecular_container):
    """Protonate the molecule.

    Args:
        molecular_container:  molecule
    """
    for name in molecular_container.conformation_names:
        info('Now protonating', name)
        # split atom into residues
        curres = -1000000
        residue = []
        o_atom = None
        c_atom = None
        for atom in molecular_container.conformations[name].atoms:
            if atom.res_num != curres:
                curres = atom.res_num
                if len(residue) > 0:
                    #backbone
                    [o_atom, c_atom] = add_backbone_hydrogen(
                        residue, o_atom, c_atom)
                    #arginine
                    if residue[0].res_name == 'ARG':
                        add_arg_hydrogen(residue)
                    #histidine
                    if residue[0].res_name == 'HIS':
                        add_his_hydrogen(residue)
                    #tryptophan
                    if residue[0].res_name == 'TRP':
                        add_trp_hydrogen(residue)
                    #amides
                    if residue[0].res_name in ['GLN', 'ASN']:
                        add_amd_hydrogen(residue)
                    residue = []
            if atom.type == 'atom':
                residue.append(atom)


def set_ligand_atom_names(molecular_container):
    """Set names for ligands in molecular container.

    Args:
        molecular_container:  molecular container for ligand names
    """
    for name in molecular_container.conformation_names:
        molecular_container.conformations[name].set_ligand_atom_names()


def add_arg_hydrogen(residue):
    """Adds Arg hydrogen atoms to residues according to the 'old way'.

    Args:
        residue:  arginine residue to protonate
    Returns:
        list of hydrogen atoms
{'fraction_non_alphanumeric': 0.06498974414336608, 'fraction_numerical': 0.010363813019540105, 'mean_word_length': 2.645808736717828, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2321812', 'n_tokens_mistral': 2733, 'n_tokens_neox': 2502, 'n_words': 790}