"""Smoothed MCG3.

Same physical definition as MCG3 in orderp.py -- 0.9 nm guest-guest, 0.6 nm
guest-water, 45 deg two-cone criterion, >= 5 shared waters per bond, >= 3
bonds per node -- with every yes/no decision replaced by a switching
function, so the value is a continuous function of the coordinates.

"""

import numpy as np
import mdtraj as md
from scipy.spatial import cKDTree
from scipy.sparse import csr_matrix, diags
from scipy.sparse.linalg import eigsh
import os
from pyretis.orderparameter import OrderParameter


def switch(value, edge, width):
    t = np.clip((np.asarray(value, dtype=np.float64) - edge) / width,
                0.0, 1.0)
    return 1.0 - (3.0 * t ** 2 - 2.0 * t ** 3)


class _SoftMCGBase(OrderParameter):
    """Shared machinery: the smoothed bond and node weights.

    Both smoothed CVs need the same quantities -- a weight in [0, 1] for
    every guest-guest bond and for every guest being a "node".  They differ
    only in how those are reduced to a single number.
    """

    def __init__(self, description, wr=0.030, wa=6.0, nc=5.0, bc=3.0,
                 wn=0.6, wb=0.6, topology="gromacs_input/conf.gro"):
        super().__init__(description=description)
        self.top = md.load_topology(topology)
        self.wr, self.wa = wr, wa          # distance (nm) and angle (deg) widths
        self.nc, self.bc = nc, bc          # the >=5 and >=3 thresholds
        self.wn, self.wb = wn, wb          # widths of those two thresholds
        self.guests = self.top.select('name CZ or name C2 or name C')
        self.waters = self.top.select('name O or name OW')

    def weights(self, system):
        """Bond weights and node weights for one configuration.

        Returns
        -------
        node_w : ndarray, shape (n_guests,)
            Weight in [0, 1] that each guest satisfies the ">= 3 bonds" node
            criterion.
        bonds : dict
            ``{(i, j): weight}`` with ``i < j``, the weight in [0, 1] that
            guests i and j share a cage face.
        """
        box = np.asarray(system.box.length, dtype=np.float64)
        pos = system.pos.reshape((self.top.n_atoms, 3)).astype(np.float64) % box
        guests, waters = self.guests, self.waters
        wr, wa = self.wr, self.wa

        tree_g = cKDTree(pos[guests], boxsize=box)
        tree_w = cKDTree(pos[waters], boxsize=box)
        near_g = tree_g.query_ball_tree(tree_g, 0.9 + 4 * wr)
        near_w = tree_g.query_ball_tree(tree_w, 0.6 + 4 * wr)

        near_w_arr = [np.asarray(lst, dtype=np.intp) for lst in near_w]
        mask = np.zeros(len(waters), dtype=bool)
        guest_pos = pos[guests]
        water_pos = pos[waters]

        node_w = np.zeros(len(guests))
        bonds = {}
        pair_cache = {}
        for i, guest in enumerate(guests):
            others = [j for j in near_g[i] if j != i]
            if not others:
                continue
            shell_i = near_w_arr[i]
            if shell_i.size:
                mask[shell_i] = True
            total_bond = 0.0
            pos_i = guest_pos[i]
            for j in others:
                key = (i, j) if i < j else (j, i)
                cached = pair_cache.get(key)
                if cached is not None:
                    total_bond += cached
                    continue
                w_j_idx = near_w_arr[j]
                if w_j_idx.size == 0:
                    continue
                shared = w_j_idx[mask[w_j_idx]]
                if shared.size == 0:
                    continue
                pos_j = guest_pos[j]
                axis = pos_i - pos_j
                axis -= box * np.round(axis / box)
                sep = np.linalg.norm(axis)
                w_pos = water_pos[shared]

                v_j = w_pos - pos_j
                v_j -= box * np.round(v_j / box)
                r_j = np.linalg.norm(v_j, axis=1)
                v_i = w_pos - pos_i
                v_i -= box * np.round(v_i / box)
                r_i = np.linalg.norm(v_i, axis=1)
                ang_j = np.degrees(np.arccos(
                    np.clip(np.dot(v_j, axis) / (sep * r_j), -1.0, 1.0)))
                ang_i = np.degrees(np.arccos(
                    np.clip(np.dot(v_i, -axis) / (sep * r_i), -1.0, 1.0)))

                # Effective number of shared cage waters: each of the four
                # criteria contributes a weight rather than a veto.
                n_eff = np.sum(switch(r_j, 0.6, wr) * switch(r_i, 0.6, wr)
                               * switch(ang_j, 45.0, wa)
                               * switch(ang_i, 45.0, wa))
                bond = (switch(sep, 0.9, wr)
                        / (1.0 + np.exp(-(n_eff - self.nc) / self.wn)))
                pair_cache[key] = bond
                total_bond += bond
                if j > i:
                    bonds[(i, j)] = bond
            if shell_i.size:
                mask[shell_i] = False
            node_w[i] = 1.0 / (1.0 + np.exp(-(total_bond - self.bc) / self.wb))
        return node_w, bonds


class MCG3Soft(_SoftMCGBase):

    def __init__(self, **kwargs):
        super().__init__('MCG3Soft', **kwargs)

    def calculate(self, system):
        node_w, _bonds = self.weights(system)
        return [float(node_w.sum())]