"""Hydrocarbon geometry oracle.

The lactam oracles in :mod:`staplebridge.oracles.geometry` (``MockGeometryOracle``,
``CPComposerStapledOracle``) are **not modified**. This is a separate oracle for
the hydrocarbon branch, because the two chemistries have genuinely different
feasibility criteria:

* lactam K-D/E closes at i,i+3 / i,i+4 with a 4.0-6.5 A Cα window;
* hydrocarbon RCM staples span one helical turn (i,i+4) or two (i,i+7), so the
  i,i+7 window sits much further out.

``ctype`` here additionally requires that the anchor residues really are the
anchor monomers the block names, which the lactam oracle has no notion of.
"""

from __future__ import annotations

import math
from typing import Any

from staplebridge.data.schemas import BuildingBlock
from staplebridge.hydrocarbon.actions import FailureReason, validate_hydrocarbon_staple
from staplebridge.hydrocarbon.catalog import is_hydrocarbon_block
from staplebridge.oracles.base import GeometryOracleBase


def _ca_distance(
    coords: list[tuple[float, float, float]] | None, i: int, j: int
) -> float | None:
    """Cα(i)-Cα(j) distance, or ``None`` when coordinates are unavailable."""
    if coords is None or i < 0 or j < 0 or i >= len(coords) or j >= len(coords):
        return None
    xi, yi, zi = coords[i]
    xj, yj, zj = coords[j]
    return math.sqrt((xi - xj) ** 2 + (yi - yj) ** 2 + (zi - zj) ** 2)


class HydrocarbonGeometryOracle(GeometryOracleBase):
    """Feasibility for hydrocarbon staples.

    Args:
        catalog: the hydrocarbon blocks in play, used for topology validation.
        sentinel_cgeom: penalty returned when the staple is not even
            type-feasible. Matches the lactam oracles' convention of a large
            sentinel so that stage-aware geometry can detect it.
    """

    def __init__(
        self,
        catalog: list[BuildingBlock] | None = None,
        sentinel_cgeom: float = 10.0,
    ) -> None:
        self.catalog = list(catalog or [])
        self.sentinel_cgeom = float(sentinel_cgeom)

    def ctype(
        self,
        sequence: list[str],
        anchor_pair: tuple[int, int] | None,
        block: BuildingBlock | None,
        *,
        peptide_ca: Any = None,
    ) -> bool:
        """True when ``(sequence, anchor_pair, block)`` is a legal hydrocarbon staple.

        Sequence-level only; the Cα window is scored by :meth:`cgeom`.
        """
        del peptide_ca
        if block is None or not is_hydrocarbon_block(block):
            return False
        catalog = self.catalog or [block]
        verdict = validate_hydrocarbon_staple(sequence, anchor_pair, block, catalog)
        return verdict is FailureReason.OK

    def cgeom(
        self,
        sequence: list[str],
        anchor_pair: tuple[int, int] | None,
        block: BuildingBlock | None,
        *,
        peptide_ca: Any = None,
    ) -> float:
        """Distance of Cα(i)-Cα(j) from the block's allowed window, in Å.

        Returns ``0.0`` inside the window. Without coordinates, falls back to a
        sequence-only verdict: ``0.0`` when type-feasible, else the sentinel.
        """
        if not self.ctype(sequence, anchor_pair, block):
            return self.sentinel_cgeom
        assert anchor_pair is not None and block is not None  # ctype guarantees

        i, j = anchor_pair
        distance = _ca_distance(peptide_ca, i, j)
        if distance is None:
            return 0.0

        low, high = block.ca_window
        if distance < low:
            return float(low - distance)
        if distance > high:
            return float(distance - high)
        return 0.0