"""Various functions to search the ChemRecon database."""

from collections import OrderedDict

import chemrecon
from chemrecon.chem.mol import MolTemplate
from chemrecon.entrygraph.scoring import Scorer

from aam_annotator.mapped_network import MappedMetabolite, MappedReaction
from aam_annotator.problem import NoAAMsFound, NoStructuresFound
from aam_annotator.utils.cache import cache

# Protocol and scoring functions
# --------------------------------------------------------------------------------------------------------------
# For finding structures
# TODO change to CompoundSelectStructure only?
protocol_structure = chemrecon.ExplorationProtocol(
    relation_types={
        chemrecon.CompoundReference,
        chemrecon.CompoundHasMolStructure,
        chemrecon.MolStructureStandardization,
    }
)
scorer_structure: Scorer[chemrecon.MolStructure] = chemrecon.Scorer[
    chemrecon.MolStructure
](score_entry_type=chemrecon.MolStructure)

# For finding AAMs
protocol_aam = chemrecon.ExplorationProtocol(
    relation_types={
        chemrecon.ReactionReference,
        chemrecon.ReactionHasAAM,
    }
)
scorer_aam: Scorer[chemrecon.AAM] = chemrecon.Scorer[chemrecon.AAM](
    score_entry_type=chemrecon.AAM
)


# Search routines
# --------------------------------------------------------------------------------------------------------------
@cache("get_eg_structures.shelve")
def get_eg_structures(
    metabolite: MappedMetabolite, search_depth: int = 3, cache_key: str = ""
) -> OrderedDict[MolTemplate, float]:
    """Return the structures for a given CRN compound"""
    try:
        eg = chemrecon.EntryGraph(initial_entries=metabolite.entries)  # ty:ignore[invalid-argument-type]
    except ValueError:
        print("    -> No initial entries!")
        return OrderedDict()

    chemrecon.explore(eg, protocol=protocol_structure, steps=search_depth)
    eg_result = scorer_structure(eg)

    # Convert to MolTemplate
    candidate_structures: OrderedDict[MolTemplate, float] = OrderedDict()
    for molstructure_entry, score in eg_result.items():
        assert isinstance(molstructure_entry, chemrecon.MolStructure)
        try:
            mol = chemrecon.mol_from_smiles(molstructure_entry.smiles, template=True)
            assert isinstance(mol, MolTemplate)
            candidate_structures[mol] = score
        except ValueError:
            # Could not parse
            pass

    if len(candidate_structures) == 0:
        raise NoStructuresFound(metabolite, eg)

    return candidate_structures


@cache("get_eg_aams.shelve")
def get_eg_aams(
    reaction: MappedReaction, search_depth: int = 3, cache_key: str = ""
) -> OrderedDict[chemrecon.AAM, float]:
    try:
        eg = chemrecon.EntryGraph(initial_entries=reaction.entries)  # ty:ignore[invalid-argument-type]
    except ValueError:
        print("    -> No initial entries!")
        return OrderedDict()

    chemrecon.explore(eg, protocol=protocol_aam, steps=search_depth)
    eg_result = scorer_aam(eg)

    if len(eg_result) == 0:
        raise NoAAMsFound(reaction, eg)

    return eg_result
