#!/usr/bin/env python3
# (C) 2023 Cadence Design Systems, Inc. (Cadence)
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  OPENEYE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.

"""Identify rings in molecule."""

import argparse
import os
import pathlib
import sys

import rich.console
from openeye import oechem
from rich_argparse import HelpPreviewAction, RichHelpFormatter

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Find atoms in same ring."
__SCRIPT_TOOLKITS__ = ["oechem"]
__SCRIPT_CATEGORIES__ = ["cheminformatics"]


def parse_options() -> argparse.Namespace:
    """Set up command line options."""
    parser = argparse.ArgumentParser(
        add_help=True,
        formatter_class=RichHelpFormatter,
        description="[yellow]" + __SCRIPT_DESC__ + "[/yellow]",
    )

    # input options
    input_group = parser.add_argument_group("Input molecule options")
    exclusive_image_group = input_group.add_mutually_exclusive_group(required=True)
    exclusive_image_group.add_argument(
        "-in",
        "--input",
        type=str,
        required=False,
        metavar="MOL-FILE",
        help="Input molecule file.",
    )
    exclusive_image_group.add_argument(
        "-s",
        "--smiles",
        type=str,
        metavar="SMILES",
        help="Input molecule smiles",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)
    return parser.parse_args()


def main() -> int:
    """Identify atoms in same ring."""
    args = parse_options()

    # initialize molecule
    mol = oechem.OEGraphMol()
    if args.input:
        ifs = oechem.oemolistream()
        if not ifs.open(args.input):
            oechem.OEThrow.Fatal("Cannot open input file!")
        if not oechem.OEReadMolecule(ifs, mol):
            oechem.OEThrow.Fatal("Cannot read input file!")
    elif args.smiles:
        if not oechem.OESmilesToMol(mol, args.smiles):
            oechem.OEThrow.Fatal("Cannot initialize molecule from smiles!")

    if not mol.IsValid():
        oechem.OEThrow.Fatal("Input molecule can not be initialized!")

    console = rich.console.Console()
    for atom in mol.GetAtoms(oechem.OEAtomIsInRing()):
        same_ring_atoms = [
            a
            for a in mol.GetAtoms(oechem.OEAtomIsInRing())
            if atoms_in_same_ring(atom, a)
        ]
        if len(same_ring_atoms) != 0:
            same_ring_str = " ".join([str(a) for a in same_ring_atoms])
            console.print(f"atom {atom!s} in the same ring as {same_ring_str}")

    return os.EX_OK


class ChainAtomOrAlreadyTraversed(oechem.OEUnaryAtomPred):
    """Predicate used to find atoms in same ring system."""

    def __init__(self, exclude: list[oechem.OEAtomBase]) -> None:
        """Initialize predicate."""
        oechem.OEUnaryAtomPred.__init__(self)
        self._exclude = exclude

    def __call__(self, atom: oechem.OEAtomBase) -> bool:
        """Evaluate atom."""
        if not atom.IsInRing():
            return False
        return atom in self._exclude

    def CreateCopy(self):  # noqa: ANN201, N802
        """Copy constructor."""
        return ChainAtomOrAlreadyTraversed(self._exclude).__disown__()


def atoms_in_same_ring(  # noqa: PLR0911
    atom_one: oechem.OEAtomBase, atom_two: oechem.OEAtomBase
) -> bool:
    """Determine whether two atoms belong to the same ring system."""
    if not atom_one.IsInRing() or not atom_two.IsInRing():
        # any of them is a chain atom
        return False

    if atom_one == atom_two:
        return True

    if (bond := atom_one.GetBond(atom_two)) is not None and not bond.IsInRing():
        return False

    first_path = list(
        oechem.OEShortestPath(atom_one, atom_two, oechem.OEAtomIsInChain())
    )
    first_path_length = len(first_path)

    if first_path_length == 2:  # noqa: PLR2004
        return True  # neighbors

    if first_path_length == 0:
        return False  # not is same ring system

    smallest_one = oechem.OEAtomGetSmallestRingSize(atom_one)
    smallest_two = oechem.OEAtomGetSmallestRingSize(atom_two)

    if first_path_length > smallest_one and first_path_length > smallest_two:
        return False  # too far away

    # try to find the second shortest different path
    exclude_pred = ChainAtomOrAlreadyTraversed(first_path[1:-1])
    second_path = list(oechem.OEShortestPath(atom_one, atom_two, exclude_pred))
    second_path_length = len(second_path)

    if second_path_length == 0:
        return False  # can not be in the same ring

    if second_path_length > smallest_one and second_path_length > smallest_two:
        return False  # too far away

    sum_ring_size = len(first_path) + len(second_path) - 2
    if sum_ring_size > smallest_one and sum_ring_size > smallest_two:
        return False

    in_ring_one = oechem.OEAtomIsInRingSize(atom_one, sum_ring_size)
    in_ring_two = oechem.OEAtomIsInRingSize(atom_two, sum_ring_size)
    return in_ring_one and in_ring_two


setattr(main, "__SCRIPT_NAME__", __SCRIPT_NAME__)
setattr(main, "__SCRIPT_DESC__", __SCRIPT_DESC__)
setattr(main, "__SCRIPT_TOOLKITS__", __SCRIPT_TOOLKITS__)
setattr(main, "__SCRIPT_CATEGORIES__", __SCRIPT_CATEGORIES__)

if __name__ == "__main__":
    sys.exit(main())