#!/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 the "best" i.e. largest common fragment of a set of molecules."""

import argparse
import os
import pathlib
import sys

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

__SCRIPT_NAME__ = pathlib.Path(__file__).absolute().stem
__SCRIPT_DESC__ = "Extract common largest core of molecules."
__SCRIPT_TOOLKITS__ = ["oechem", "oegraphsim"]
__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]",
    )
    parser.add_argument("--help-image", action=HelpPreviewAction)

    # input options
    input_group = parser.add_argument_group("Input options")
    input_group.add_argument(
        "-in",
        "--input",
        type=str,
        required=True,
        metavar="MOL-FILE",
        help="Input molecule file.",
    )
    core_group = parser.add_argument_group("Core identifying options")
    core_group.add_argument(
        "--min-bonds",
        type=int,
        required=False,
        default=5,
        choices=range(1, 10),
        help="Minimum number of bonds in common core",
    )
    core_group.add_argument(
        "--max-bonds",
        type=int,
        required=False,
        default=18,
        choices=range(4, 18),
        help="Maximum number of bonds in common core",
    )
    return parser.parse_args()


def main() -> int:
    """Identify common fragments."""
    args = parse_options()

    # check input file
    ifs = oechem.oemolistream()
    if not ifs.open(args.input):
        oechem.OEThrow.Fatal("Cannot open input file!")

    # read molecules
    mol_list: list[oechem.OEMolBase] = []
    for mol in ifs.GetOEGraphMols():
        oechem.OESuppressHydrogens(mol)
        mol_list.append(oechem.OEGraphMol(mol))

    console = rich.console.Console()
    console.print(f"Number of molecules = {len(mol_list)}")

    if (
        best_core := get_largest_common_core(mol_list, args.min_bonds, args.max_bonds)
    ) is not None:
        console.print(
            f"Largest core fragment of the molecules = {oechem.OEMolToSmiles(best_core)}"
        )
    else:
        console.print("No core fragment is identified!")
    return os.EX_OK


def get_reference_molecule(mol_list: list[oechem.OEMolBase]) -> oechem.OEMolBase | None:
    """Return a reference "smallest" molecule."""
    ref_mol = None
    for mol in mol_list:
        if ref_mol is None or mol.NumAtoms() < ref_mol.NumAtoms():
            ref_mol = mol
    return ref_mol


def get_fragments(
    mol: oechem.OEMolBase, min_bonds: int, max_bonds: int
) -> list[oechem.OEMolBase]:
    """Return a list of all possible fragments of a molecule."""
    frags: list[oechem.OEMolBase] = []
    fp_type = oegraphsim.OEGetFPType(
        f"Tree,ver=2.0.0,size=4096,bonds={min_bonds}-{max_bonds},atype=AtmNum,btype=Order"
    )

    unique_fragments = True
    for ab_set in oegraphsim.OEGetFPCoverage(mol, fp_type, unique_fragments):
        frag_atom_pred = oechem.OEIsAtomMember(ab_set.GetAtoms())

        frag = oechem.OEGraphMol()
        adjust_hcount = True
        oechem.OESubsetMol(frag, mol, frag_atom_pred, adjust_hcount)
        oechem.OEFindRingAtomsAndBonds(frag)
        frags.append(oechem.OEGraphMol(frag))

    return frags


def get_core_fragments(
    mol_list: list[oechem.OEMolBase],
    frags: list[oechem.OEMolBase],
    atom_expr: int = oechem.OEExprOpts_DefaultAtoms,
    bond_expr: int = oechem.OEExprOpts_DefaultBonds,
) -> list[oechem.OEMolBase]:
    """Identify fragments that are present in all molecules."""
    core_frags: list[oechem.OEMolBase] = []
    for frag in frags:
        ss = oechem.OESubSearch(frag, atom_expr, bond_expr)
        if not ss.IsValid():
            continue

        valid_core = True
        for mol in mol_list:
            valid_core = ss.SingleMatch(mol)
            if not valid_core:
                break

        if valid_core:
            core_frags.append(frag)

    return core_frags


def get_largest_common_core(
    mol_list: list[oechem.OEMolBase],
    min_bonds: int,
    max_bonds: int,
    atom_expr: int = oechem.OEExprOpts_DefaultAtoms,
    bond_expr: int = oechem.OEExprOpts_DefaultBonds,
) -> oechem.OEMolBase | None:
    """Enumerate all the fragments to find a largest common core."""
    ref_mol = get_reference_molecule(mol_list)
    if ref_mol is None:
        oechem.OEThrow.Fatal("No reference molecule identified.")
        return None

    frags = get_fragments(ref_mol, min_bonds, max_bonds)
    if len(frags) == 0:
        oechem.OEThrow.Fatal(
            f"No fragment is enumerated with bonds {min_bonds}-{max_bonds}!"
        )

    common_frags = get_core_fragments(mol_list, frags, atom_expr, bond_expr)
    if len(common_frags) == 0:
        oechem.OEThrow.Error("No common fragment is found!")

    largest_core = None
    for frag in common_frags:
        if largest_core is None or get_fragment_score(
            largest_core
        ) < get_fragment_score(frag):
            largest_core = frag
    return largest_core


def get_fragment_score(frag: oechem.OEMolBase) -> float:
    """Score the fragment (larger is better)."""
    score = 0.0
    score += 2.0 * oechem.OECount(frag, oechem.OEAtomIsInRing())
    score += 1.0 * oechem.OECount(frag, oechem.OENotAtom(oechem.OEAtomIsInRing()))
    return score


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())