from __future__ import annotations

from typing import Any

from staplebridge.chemistry.state import StapleState
from staplebridge.data.vocab import AMINO_ACIDS
from staplebridge.oracles.base import PeptidePriorBase


class MockPeptidePrior(PeptidePriorBase):
    def score_transition(self, old_state: StapleState, new_state: StapleState, context: dict[str, Any] | None = None) -> float:
        del old_state, context
        seq = new_state.sequence_tokens
        valid = sum(1 for t in seq if t in AMINO_ACIDS)
        length_pen = abs(len(seq) - 12) / 12
        return valid / max(len(seq), 1) - length_pen