from __future__ import annotations

from typing import Any

from staplebridge.data.schemas import BuildingBlock
from staplebridge.oracles.base import BlockPriorBase


def _motif_satisfied(seq: list[str], i: int, j: int, motif: dict | None) -> bool:
    if motif is None:
        return True
    spacings = motif.get("spacings")
    if spacings != "any" and spacings is not None and (j - i) not in spacings:
        return False
    if motif.get("i_aa") and seq[i] not in motif["i_aa"]:
        return False
    if motif.get("j_aa") and seq[j] not in motif["j_aa"]:
        return False
    return True


class MockBlockPrior(BlockPriorBase):
    def score_block(
        self,
        sequence: list[str],
        anchor_pair: tuple[int, int] | None,
        block: BuildingBlock | None,
        context: dict[str, Any] | None = None,
    ) -> float:
        del context
        if block is None:
            return -0.3
        if anchor_pair is None:
            return -1.0
        i, j = anchor_pair
        score = 0.0
        score += 0.5 * block.synthetic_accessibility_score
        score -= 0.2 * block.cost_score
        if _motif_satisfied(sequence, i, j, block.motif):
            score += 1.0
        else:
            score -= 1.0
        return score