# mmpdb - matched molecular pair database generation and analysis
#
# Copyright (c) 2015-2017, F. Hoffmann-La Roche Ltd.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#    * Redistributions of source code must retain the above copyright
#      notice, this list of conditions and the following disclaimer.
#    * Redistributions in binary form must reproduce the above
#      copyright notice, this list of conditions and the following
#      disclaimer in the documentation and/or other materials provided
#      with the distribution.
#    * Neither the name of F. Hoffmann-La Roche Ltd. nor the names of
#      its contributors may be used to endorse or promote products
#      derived from this software without specific prior written
#      permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from __future__ import print_function

# Determine the environment around the attachment point(s) of the constant
# part using circular fingerprints.

import sys
import hashlib
import binascii

from rdkit import Chem
from rdkit.Chem import AllChem

from collections import namedtuple

if sys.version_info.major == 2:
    def fingerprint_to_text(s):
        return s
else:
    def fingerprint_to_text(s):
        return s.decode("ascii")

# Internal helper class
class EnvironmentCenters(object):
    def __init__(self, mol, atom_ids):
        self.mol = mol
        self.atom_ids = atom_ids

    def __repr__(self):
        return "EnvironmentCenters(%r, %r)" % (self.mol, self.atom_ids)

class EnvironmentFingerprint(object):
    __slots__ = ("radius", "fingerprint")

    def __init__(self, radius, fingerprint):
        self.radius = radius
        self.fingerprint = fingerprint
    

def find_centers(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError("Could not parse the context SMILES %r" % (smiles,))
    
    # Figure out which atoms to use as the centers
    centers = {}
    for atom_idx, atom in enumerate(mol.GetAtoms()):
        if atom.GetAtomicNum() != 0:
            continue
        mapno = atom.GetIntProp("molAtomMapNumber")

        # Don't need to clear the property because RDKit's Morgan fingerprint
        # doesn't use it.
        #atom.ClearProp("molAtomMapNumber")
        assert mapno not in centers, (mapno, smiles)

        centers[mapno] = atom_idx
            
    if not centers:
        raise ValueError("No attachment points found in context SMILES %r" % (smiles,))

    n = len(centers)
    if n > 3:
        raise ValueError("Too many attachment points found in context SMILES %r" % (smiles,))
    if n == 3:
        try:
            ordered_centers = centers[1], centers[2], centers[3]
        except KeyError:
            raise ValueError("Expecting attachment points *:1, *:2, and *:3 in context SMILES %r" % (smiles,))
    elif n == 2:
        try:
            ordered_centers = centers[1], centers[2]
        except KeyError:
            raise ValueError("Expecting attachment points *:1 and *:2 in context SMILES %r" % (smiles,))
    elif n == 1:
        try:
            ordered_centers = centers[1],  # this is a one element tuple
        except KeyError:
            raise ValueError("Expecting attachment point *:1 in context SMILES %r" % (smiles,))

    return EnvironmentCenters(mol, ordered_centers)

# Get each of the atoms counts radii as it's being computed.
def iter_num_atoms_for_radii(centers, max_radius):
    return _iter_num_atoms_for_radii(centers.mol, max_radius, centers.atom_ids)

# My thought is to use this for testing.
def _iter_num_atoms_for_radii(mol, max_radius, start_atoms): 
    atom_objs = list(mol.GetAtoms())
    
    unique_atoms = set(start_atoms)
    assert len(start_atoms) == len(unique_atoms), "duplicate start atom"
    ignore_atoms = set(a for a in start_atoms if not is_heavy_atom(mol.GetAtomWithIdx(a)))
    
    yield (len(unique_atoms) - len(ignore_atoms))
    
    border_atoms = unique_atoms.copy()
    
    for radius in range(max_radius):  # up to and including max_radius
        new_atoms = set()
        
        for atom in border_atoms:
            for neighbor_atom_obj in mol.GetAtomWithIdx(atom).GetNeighbors():
                neighbor_atom = neighbor_atom_obj.GetIdx()
                if neighbor_atom not in unique_atoms:
                    unique_atoms.add(neighbor_atom)
                    new_atoms.add(neighbor_atom)
                    
                    if not is_heavy_atom(neighbor_atom_obj):
                        ignore_atoms.add(neighbor_atom)

        border_atoms = new_atoms
        yield (len(unique_atoms) - len(ignore_atoms))

    
def is_heavy_atom(atom_obj):
    # Exclude '*' and '[H]' atoms (but count '[1H]', '[2H]', etc.)
    eleno = atom_obj.GetAtomicNum()
    if eleno == 0 or (eleno == 1 and atom_obj.GetIsotope() == 0):
        return False
    return True


def find_center_fingerprints(centers, radius):
    center_fps = []
    for atom_idx in centers.atom_ids:
        fp = AllChem.GetMorganFingerprint(centers.mol, radius, fromAtoms=[atom_idx])
        # Hash so everything so they are a constant 32 binary bytes.
        center_fp = hashlib.sha256(fp.ToBinary()).digest()
        center_fps.append(center_fp)
    return center_fps


def _make_fp(*center_fps):
    concat_fp = hashlib.new("sha256")
    for center_fp in center_fps:
        concat_fp.update(center_fp)
        
    sha2_digest = concat_fp.digest()
    env_fp = binascii.b2a_base64(sha2_digest)
    assert env_fp[-2:] == b"=\n", env_fp
    return fingerprint_to_text(env_fp[:-2])  # trim the "=\n"


def find_environment_fingerprint(centers, radius):
    # I want a unique fingerpint for the given environment.

    # The attachment points are in canonical order so I can
    # concatenate the circular fingerprints, making sure that there's
    # only one way to produce a given concatenation. (There's a low
    # chance that simple concatenation of two ToBinary() strings,
    # which are variable length, might be ambiguous with a
    # concatenation of other strings.)

    # Fix the ambiguity problem by using the SHA2 of the ToBinary() so
    # all circular fingerprints are a constant length.

    # The result is far longer than it needs to be. I only need to
    # test two environments for identify. I don't need to know the
    # per-center breakdown. Use sha2 again to solve that problem.

    # TODO: the fingerprints at r=0 are constant based on the number of cuts.
    # Would it be faster to have special code for that case?
    concat_fp = _make_fp(*find_center_fingerprints(centers, radius))
    return concat_fp

# The input is a SMILES for the constant term, like
#   CCC([*:2])CCC([*:1])CCCCCC[*:3]
# The centers are the attachment points 1, 2, and 3.
# 

def compute_constant_environment_from_centers(centers, max_radius=5):
    env_fps = []
    for radius in range(max_radius+1):
        fingerprint = find_environment_fingerprint(centers, radius)
        env_fps.append(
            EnvironmentFingerprint(radius, fingerprint)
            )

    return env_fps


def compute_constant_center_fingerprints(constant_smiles, min_radius=0, max_radius=5):
    # If the constant is unlabeled, get the labeled version.
    if "*" in constant_smiles:
        # Handle differences in * representation in pre/post 2018 versions of RDKit
        if "[*]" in constant_smiles:
            wildcard = "[*]"
            width = 3
        else:
            wildcard = "*"
            width = 1
        # The conversion is direct; it's already canonical, in the order 1, 2, 3.
        # The following works for 1, 2, or 3 attachment points. If n=1 then
        # the last two replace() functions do nothing.
        # The main problem that if the wildcard is "*" then I can't do a
        # simple string substitution otherwise I'll match the previous *.
        i = 0
        i = constant_smiles.find(wildcard, i)
        if i >= 0:
            constant_smiles = constant_smiles[:i] + "[*:1]" + constant_smiles[i+width:]
            i = constant_smiles.find(wildcard, i+5)
            if i >= 0:
                constant_smiles = constant_smiles[:i] + "[*:2]" + constant_smiles[i+width:]
                i = constant_smiles.find(wildcard, i+5)
                if i >= 0:
                    constant_smiles = constant_smiles[:i] + "[*:3]" + constant_smiles[i+width:]

    env_centers = find_centers(constant_smiles)

    all_center_fps = []
    for radius in range(min_radius, max_radius+1):
        all_center_fps.append(find_center_fingerprints(env_centers, radius))

    return all_center_fps

######

# The "reorder" describes the changes in the the core. For example:

_invert_order_table = {
    None: None,
    "1": None,
    
    "12": None,
    "21": (1, 0),
    
    "123": None,
    "132": (0, 2, 1),
    "213": (1, 0, 2),
    "231": (2, 0, 1),
    "312": (1, 2, 0),
    "321": (2, 1, 0),
    }


def compute_possible_environments(center_fps, symmetry_class, reorder=None):
    # This is complex.
    # The center_fps are based on the constant/context fragments, which have the
    # attachment orders of 1, 2, 3.

    # The attachment orders of the core SMILES is arbitrary. Not only
    # might the (semi-)canonical SMILES be  [*:2]C([*:3])N[*:1],
    # but the attachments might be reordered into one of the n! forms,
    # like [*:1]C([*:3])N[*:2] or [*:2]C([*:1])N[*:3]. I need to
    # unorder the permutations to work with the 'true' numbering.

    # Once I have that, I need to find all the ways to match the
    # core attachments the context. Consider:
    #   [*:2]CN([*:3])N[*:1]
    # This has a symmetry class of "122" because the *:1 is in its
    # own symmetry group and the *:2 and *:3 are in the same group.
    # (The symmetry group labels are in attachment number order, not
    # position in the SMILES.)

    # If this is attached to "[:*1]N.[:*2]N.[:*3]P" then there's
    # some ambiguity. The environment context should be the same
    # as if it were attached to "[:*1]N.[:*2]P.[:*3]N", because the
    # :*2 and *:3 are symmetric.

    # This means I can have to enumerate all of the possiblities, to
    # find all of the possible environments.

    invert_order = _invert_order_table[reorder]
    if invert_order is not None:
        center_fps = [center_fps[i] for i in invert_order]

    fps = set()
    if symmetry_class in ("1", "12", "123"):
        fps.add(_make_fp(*center_fps))
    elif symmetry_class == "11":
        fps.add(_make_fp(center_fps[0], center_fps[1]))
        fps.add(_make_fp(center_fps[1], center_fps[0]))

    elif symmetry_class == "111":
        fps.add(_make_fp(center_fps[0], center_fps[1], center_fps[2]))
        fps.add(_make_fp(center_fps[0], center_fps[2], center_fps[1]))
        fps.add(_make_fp(center_fps[1], center_fps[0], center_fps[2]))
        fps.add(_make_fp(center_fps[1], center_fps[2], center_fps[0]))
        fps.add(_make_fp(center_fps[2], center_fps[0], center_fps[1]))
        fps.add(_make_fp(center_fps[2], center_fps[1], center_fps[0]))

    elif symmetry_class == "112":
        fps.add(_make_fp(center_fps[0], center_fps[1], center_fps[2]))
        fps.add(_make_fp(center_fps[1], center_fps[0], center_fps[2]))

    elif symmetry_class == "122":
        fps.add(_make_fp(center_fps[0], center_fps[1], center_fps[2]))
        fps.add(_make_fp(center_fps[0], center_fps[2], center_fps[1]))

    elif symmetry_class == "121":
        fps.add(_make_fp(center_fps[0], center_fps[1], center_fps[2]))
        fps.add(_make_fp(center_fps[2], center_fps[1], center_fps[0]))

    else:
        raise AssertionError("I forgot one: %r" % (symmetry_class,))

    return list(fps)

## def generate_all_possible_constant_fingerprints(constant_smiles):
##     # The n circular fingerprints around each of the n attachment points, for each possible radius
##     all_center_fps = compute_constant_center_fingerprints(constant_smiles)
##     n = constant_smiles.count("*")
##     symclass = "1" * n
##     fingerprints = set()
##     for center_fps in all_center_fps:
##         fingerprints.update(compute_possible_environments(center_fps, symclass))
##     return fingerprints

def get_all_possible_fingerprints(all_center_fps, symmetry_class, permutation):
    possible_envs = set()
    for center_fps in all_center_fps:
        possible_envs.update(compute_possible_environments(
            center_fps, symmetry_class, permutation))
    return sorted(possible_envs)
