# -*- coding: utf-8 -*-
"""Order parameter for the water autoionization study (2019).

RECONSTRUCTION NOTICE
---------------------
The original study was run with the now-retired TISMOL program, whose
``H2OMAX DISTANCE`` reaction-coordinate source code is *not* part of these
source files.  This module is therefore a faithful-as-possible
*reconstruction*, not the original code.

The order parameter is defined here as the largest, over all hydrogen atoms,
of the distance from that hydrogen to its nearest oxygen -- i.e. the
"most detached proton".  When every proton is covalently bound this is the
typical O-H bond length (~1.0 Angstrom); as a proton transfers during
autoionization the relevant distance grows.  This is physically consistent
with water autoionization and with the published interface set (~1.07 to
10 Angstrom), but it has NOT been verified to reproduce TISMOL's exact
``H2OMAX`` values.  Confirm against the original definition before using the
results quantitatively.

Coordinate access mirrors the working CP2K order parameters in this archive
(``particles.pos`` in internal/Angstrom units, periodic distances via
``system.box.pbc_dist_coordinate``).
"""
import numpy as np
from pyretis.orderparameter import OrderParameter


class MaxDetachedProton(OrderParameter):
    """Maximum over H atoms of the distance to the nearest O atom.

    Attributes
    ----------
    periodic : boolean
        Apply periodic boundary conditions to the O-H distances.
    """

    def __init__(self, periodic=True):
        """Initialise the order parameter.

        Parameters
        ----------
        periodic : boolean, optional
            Apply periodic boundary conditions to the O-H distances.
        """
        super().__init__(
            description='Water autoionization: max detached-proton O-H distance',
            velocity=False,
        )
        self.periodic = periodic

    @staticmethod
    def _oxygen_hydrogen_indices(particles):
        """Return the oxygen and hydrogen atom indices from the names."""
        names = [str(name) for name in particles.name]
        oxygen = [idx for idx, name in enumerate(names) if name.startswith('O')]
        hydrogen = [idx for idx, name in enumerate(names) if name.startswith('H')]
        return oxygen, hydrogen

    def calculate(self, system):
        """Calculate the order parameter.

        Parameters
        ----------
        system : object like :py:class:`.System`
            The object holding the positions and box.

        Returns
        -------
        out : list with one float
            The maximum, over hydrogens, of the distance to the nearest
            oxygen.
        """
        particles = system.particles
        oxygen, hydrogen = self._oxygen_hydrogen_indices(particles)

        max_min_distance = 0.0
        for h_idx in hydrogen:
            nearest = None
            for o_idx in oxygen:
                delta = particles.pos[h_idx] - particles.pos[o_idx]
                if self.periodic:
                    delta = system.box.pbc_dist_coordinate(delta)
                distance = np.sqrt(np.dot(delta, delta))
                if nearest is None or distance < nearest:
                    nearest = distance
            if nearest is not None and nearest > max_min_distance:
                max_min_distance = nearest
        return [max_min_distance]