#!/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.

"""Code snippet for ring perception."""

from openeye import oechem, oedepict

from oecookbook.scripts.ring_perception import atoms_in_same_ring

mol = oechem.OEGraphMol()
oechem.OESmilesToMol(mol, "C1CC2CCC(CCCC3CCC4CCC5CCC(C1)CC5C4C3)CC2")
oedepict.OEPrepareDepiction(mol)

width, height = 400, 300
image = oedepict.OEImage(width, height)


def get_macro_cycle_atoms(
    mol: oechem.OEMolBase, min_ring_size: int = 10
) -> list[oechem.OEAtomBase]:
    """Find atoms in rings larger than a specified size."""
    oechem.OEFindRingAtomsAndBonds(mol)
    if oechem.OECount(mol, oechem.OEAtomIsInRing()) == 0:
        return []  # no ring atoms

    small_ring_atoms: set[oechem.OEAtomBase] = set()
    for atom in mol.GetAtoms(oechem.OEAtomIsInRing()):
        if oechem.OEAtomGetSmallestRingSize(atom) <= min_ring_size:
            small_ring_atoms.add(atom)
    macro_cyclic_ring_seeds: set[oechem.OEAtomBase] = (
        set(mol.GetAtoms(oechem.OEAtomIsInRing())) - small_ring_atoms
    )
    if macro_cyclic_ring_seeds == set():
        return []  # no macro-cyclic atoms

    macro_cycle_atoms: set[oechem.OEAtomBase] = set()
    macro_cycle_atoms.update(macro_cyclic_ring_seeds)

    for atom_one in macro_cyclic_ring_seeds:
        for atom_two in small_ring_atoms:
            if atom_two in macro_cycle_atoms:
                continue  # already identified as macro-cyclic
            if atoms_in_same_ring(atom_one, atom_two):
                macro_cycle_atoms.add(atom_two)

    return list(macro_cycle_atoms)


macro_cycle_atoms = get_macro_cycle_atoms(mol, min_ring_size=10)
abset = oechem.OEAtomBondSet()
for atom in macro_cycle_atoms:
    abset.AddAtom(atom)


scale = oedepict.OEScale_AutoScale
opts = oedepict.OE2DMolDisplayOptions(width, height, scale)

disp = oedepict.OE2DMolDisplay(mol, opts)
highlight = oedepict.OEHighlightByBallAndStick(oechem.OEBlueTint)
oedepict.OEAddHighlighting(disp, highlight, abset)

oedepict.OERenderMolecule(image, disp)
oedepict.OEWriteImage("depict_macro_cycle_atoms.svg", image)