from collections import OrderedDict

from src.chemrecon import Direction
from chem.mol import MolTemplate, Mol
from chem.create_mol import mol_from_struct_entry
from chemrecon.entrygraph.explore import explore
from chemrecon.entrygraph.scoring import Scorer
from chemrecon.schema import (
    Compound, MolStructure
)
from chemrecon.schema.relation_types_procedural.compound_select_structure_proceduralrelation import \
    CompoundSelectStructure
from chemrecon.entrygraph.explorationprotocol import ExplorationProtocol


def get_structures_from_compound(
    entries: Compound | list[Compound],
    consider_first_entry_primary: bool = False,
) -> OrderedDict[MolTemplate, float]:

    # TODO update with new protocol syntax
    raise NotImplementedError()

    entry_set: set[Compound] = set()
    match entries:
        case Compound():
            entry_set = {entries}
        case list():
            entry_set = set(entries)
        case _:
            raise ValueError()

    # Construct entrygraph and run ranking  # TODO set first entry as primary if specified!
    try:
        eg = EG_Structure(
            initial_entries = entry_set
        )
    except ValueError as e:
        return OrderedDict()
    explore(entrygraph = eg, steps = 2)

    # Score
    scoring = scorer_default(eg)
    scoring_out: OrderedDict[Mol, float] = OrderedDict()
    for k, v in scoring.items():
        try:
            scoring_out[mol_from_struct_entry(k)] = v
        except AttributeError:
            continue

    return scoring_out


# Specification
EG_Structure = ExplorationProtocol(
    entry_types = {Compound, MolStructure},
    relation_types = {
        (CompoundSelectStructure, Direction.FORWARDS)
    },
    entry_types_initial = {Compound}
)

scorer_default = Scorer[MolStructure](
    score_entry_type = MolStructure
)