"""
core.py — Fundamental data structures for the Modular Reasoning Scaffold (MRS).

This module defines the primary reasoning primitives used throughout MRS:
`RecursionNode`, `StateTracker`, and `ConstraintMonitor`.

All entities in this file are lightweight, typed, and designed for 
inference-time execution in small/medium language models.

The goal: provide external cognitive structure the model can reliably use.
"""

from pydantic import BaseModel, Field
from typing import Any, List, Callable


class RecursionNode(BaseModel):
    """
    A single reasoning iteration within MRS.

    Each node represents one step of model-generated transformation, along with
    metadata required to track stability, drift, and transitions.

    Parameters
    ----------
    x : Any
        The input passed to the model for this reasoning step.
    phi : str
        The transformation rule or reasoning directive applied at this step.
    drift : float
        The local coherence deviation (Δ_i), measuring how far the output
        strays from internal consistency or prior steps.
    output : Any
        The model’s produced output for this node.

    Notes
    -----
    RecursionNodes form the core sequence S = {R1, R2, ..., Rn}.  
    They are intentionally minimal: no heavy logic beyond validation.
    """
    x: Any
    phi: str
    drift: float = Field(ge=0.0)
    output: Any


class StateTracker(BaseModel):
    """
    Persistent storage for reasoning outputs across all recursion nodes.

    The StateTracker maintains the cumulative reasoning trajectory, enabling:
    - multi-step decomposition
    - reflective consolidation
    - drift-aware halting decisions

    Parameters
    ----------
    states : List[Any]
        A running list of outputs o_1 ... o_n collected from each node.

    Methods
    -------
    add(state):
        Append a new state to the tracker.
    aggregate():
        Return the full history for inspection or reflective processing.
    """
    states: List[Any] = []

    def add(self, state: Any):
        """Add a new reasoning state to the tracker."""
        self.states.append(state)

    def aggregate(self) -> List[Any]:
        """Return all collected states."""
        return self.states
        

class ConstraintMonitor(BaseModel):
    """
    Drift and stability monitor for MRS reasoning sequences.

    This module computes accumulated drift and determines when recursion
    should halt or branch based on a threshold τ.

    Parameters
    ----------
    threshold : float
        Global drift limit τ.
    history : List[float]
        List of drift values Δ_i across steps.

    Methods
    -------
    update(drift):
        Record new drift value.
    total():
        Return cumulative drift C(S_n).
    is_stable():
        Return True if total drift < τ.
    """
    threshold: float = 1.0
    history: List[float] = []

    def update(self, drift: float):
        """Record a new drift value Δ_i."""
        self.history.append(drift)

    def total(self) -> float:
        """Compute cumulative drift C(S_n)."""
        return sum(self.history)

    def is_stable(self) -> bool:
        """Determine whether recursion may continue."""
        return self.total() < self.threshold
