Blame for patch_view/patching.py

"""Patching strategies for entropy-based audio segmentation."""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Dict, List

import numpy as np

from .legacy_patching import (
    patch_additive_min,
    patch_additive_monotonic,
    patch_global_threshold,
    patch_multiplicative_min,
    patch_multiplicative_monotonic,
)


# ============================================================================
# OOP Constraint Classes for Multi-Constraint Patching
# ============================================================================

class Constraint(ABC):
    """Base class for patching constraints."""

    def __init__(self, threshold: float, guard: bool = False):
        self.threshold = threshold
        self.guard = guard
        self.reset()

    @abstractmethod
    def reset(self):
        """Reset constraint state when starting new patch."""
        pass

    @abstractmethod
    def evaluate(self, current_value: float, token_idx: int) -> bool:
        """
        Evaluate if constraint is satisfied.

        Args:
            current_value: Current metric value (entropy/perplexity)
            token_idx: Current token index

        Returns:
            True if boundary should be created
        """
        pass

    @abstractmethod
    def update(self, current_value: float, token_idx: int):
        """Update internal state after processing token."""
        pass


class GlobalConstraint(Constraint):
    """Patch when metric exceeds threshold: H(t) > θ"""

    def __init__(self, threshold: float):
        super().__init__(threshold, guard=False)

    def reset(self):
        pass  # No state to track

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        return current_value > self.threshold

    def update(self, current_value: float, token_idx: int):
        pass  # No state to update


class AdditiveMonotonicConstraint(Constraint):
    """Patch when metric increase exceeds threshold: H(t) - H(t-1) > θ"""

    def reset(self):
        self.prev_value = None

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        if self.prev_value is None:
            return False
        return (current_value - self.prev_value) > self.threshold

    def update(self, current_value: float, token_idx: int):
        self.prev_value = current_value


class MultiplicativeMonotonicConstraint(Constraint):
    """Patch when metric ratio exceeds threshold: H(t) / H(t-1) > θ"""

    def reset(self):
        self.prev_value = None

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        if self.prev_value is None or self.prev_value < 1e-6:
            return False
        return (current_value / self.prev_value) > self.threshold

    def update(self, current_value: float, token_idx: int):
        self.prev_value = current_value


class AdditiveMinConstraint(Constraint):
    """Patch when metric increase from patch minimum exceeds threshold: H(t) - min(H) > θ"""

    def reset(self):
        self.patch_min = None

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        if self.patch_min is None:
            return False
        return (current_value - self.patch_min) > self.threshold

    def update(self, current_value: float, token_idx: int):
        if self.patch_min is None:
            self.patch_min = current_value
        else:
            self.patch_min = min(self.patch_min, current_value)


class MultiplicativeMinConstraint(Constraint):
    """Patch when metric ratio to patch minimum exceeds threshold: H(t) / min(H) > θ"""

    def reset(self):
        self.patch_min = None

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        if self.patch_min is None or self.patch_min < 1e-6:
            return False
        return (current_value / self.patch_min) > self.threshold

    def update(self, current_value: float, token_idx: int):
        if self.patch_min is None:
            self.patch_min = current_value
        else:
            self.patch_min = min(self.patch_min, current_value)


class CumulativeSumConstraint(Constraint):
    """Patch AFTER cumulative sum exceeds threshold: Σ H(t) > θ

    Creates boundary after the token that causes sum to exceed threshold.
    The patch INCLUDES the token that exceeded the threshold.
    """

    def reset(self):
        self.cumulative_sum = 0.0

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        # Check AFTER updating cumsum - boundary after exceeding
        temp_sum = self.cumulative_sum + current_value
        return temp_sum > self.threshold

    def update(self, current_value: float, token_idx: int):
        self.cumulative_sum += current_value


class CumulativeMaxConstraint(Constraint):
    """Patch to limit cumulative sum within patches: Σ H(patch) ≈ θ

    Creates boundaries to keep patch cumulative sums near the threshold.
    Attempts to pack tokens together while avoiding very large cumulative sums.

    Note: Due to boundary semantics (boundary token is included in its patch),
    patches may exceed the threshold by one token's worth. This constraint
    provides approximate control over patch size rather than strict upper bounds.

    Example: metric=[2, 3, 4, 2, 3, 4], threshold=8
    - Tokens 0-3: Packed together until cumsum approaches threshold
    → Patch [0,1,2,3] with sum=11 (exceeded, but tries to pack efficiently)
    - Tokens 4-5: New patch
    → Patch [4,5] with sum=7 (below threshold)

    Compare with CumulativeSumConstraint which creates more frequent boundaries.

    This constraint manages its own state and doesn't use the standard reset()
    mechanism since it needs to persist state across resets triggered by other constraints.
    """

    def __init__(self, threshold: float):
        self.threshold = threshold
        # Don't call super().__init__() to avoid calling reset()
        self.cumulative_sum = 0.0
        self._deferred_break = False
        self._deferred_value = 0.0

    def reset(self):
        # Don't reset - we manage state ourselves
        pass

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        # If we previously detected overflow, create boundary now
        if hasattr(self, '_deferred_break') and self._deferred_break:
            return True
        return False

    def update(self, current_value: float, token_idx: int):
        # Initialize state on first call
        if not hasattr(self, 'cumulative_sum'):
            self.cumulative_sum = 0.0
            self._deferred_break = False

        # If we just created a boundary, reset cumsum to the deferred value,
        # then process the current token
        if self._deferred_break:
            self.cumulative_sum = self._deferred_value
            self._deferred_break = False
            # Now process current token (don't return early!)

        # Check if adding current token would overflow
        temp_sum = self.cumulative_sum + current_value
        if temp_sum > self.threshold and token_idx > 0:
            # Defer the break to next iteration
            self._deferred_break = True
            self._deferred_value = current_value
        else:
            # No overflow, add to cumsum
            self.cumulative_sum += current_value
            self._deferred_break = False


class MinimumMetricConstraint(Constraint):
    """Guard: require metric to reach a minimum value before allowing boundaries."""

    def __init__(self, threshold: float):
        super().__init__(threshold, guard=True)

    def reset(self):
        pass

    def evaluate(self, current_value: float, token_idx: int) -> bool:
        return current_value >= self.threshold

    def update(self, current_value: float, token_idx: int):
        pass


# Constraint registry for easy instantiation
CONSTRAINT_REGISTRY = {
    "global": GlobalConstraint,
    "additive": AdditiveMonotonicConstraint,
    "multiplicative": MultiplicativeMonotonicConstraint,
    "additive_min": AdditiveMinConstraint,
    "multiplicative_min": MultiplicativeMinConstraint,
    "cumulative": CumulativeSumConstraint,
    "cumulative_max": CumulativeMaxConstraint,
    "min_metric": MinimumMetricConstraint,
}


class MultiConstraintEvaluator:
    """Evaluates multiple constraints with OR logic."""

    def __init__(self, constraints: Dict[str, float]):
        """
        Args:
            constraints: {"constraint_name": threshold}
                       e.g., {"global": 5.0, "additive": 0.5}
        """
        self.guards = []
        self.breakers = []
        self._all = []
        for name, threshold in constraints.items():
            if name not in CONSTRAINT_REGISTRY:
                raise ValueError(f"Unknown constraint: {name}")
            constraint_class = CONSTRAINT_REGISTRY[name]
            instance = constraint_class(threshold)
            self._all.append(instance)
            if getattr(instance, "guard", False):
                self.guards.append(instance)
            else:
                self.breakers.append(instance)

        if not self.breakers:
            raise ValueError("At least one patch-breaking constraint is required")

    def evaluate_all(self, metric: np.ndarray) -> np.ndarray:
        """
        Evaluate all constraints on metric sequence.

        Args:
            metric: Entropy/perplexity values

        Returns:
            boundary_mask: Boolean array where True = patch boundary
        """
        num_tokens = len(metric)
        boundary_mask = np.zeros(num_tokens, dtype=bool)

        for i, value in enumerate(metric):
            guards_ready = all(
                guard.evaluate(value, i) for guard in self.guards
            ) if self.guards else True

            # OR logic: any constraint can trigger boundary
            should_break = guards_ready and any(
                breaker.evaluate(value, i)
                for breaker in self.breakers
            )

            # Update all constraint states
            # Note: For constraints that create boundary at i,
            # this update is AFTER reset, so the value at i
            # becomes the first value of the next patch
            for constraint in self._all:
                constraint.update(value, i)

            if should_break:
                # Boundary at current token: include it in the NEXT patch.
                boundary_mask[i] = True
                # Reset all constraints for new patch (after assigning current value)
                for constraint in self._all:
                    constraint.reset()

        return boundary_mask

    def to_patches(self, boundary_mask: np.ndarray) -> List[List[int]]:
        """Convert boundary mask to patch list."""
        patches = []
        current_patch: List[int] = []

        for i, is_boundary in enumerate(boundary_mask):
            if is_boundary:
                if current_patch:
                    patches.append(current_patch)
                current_patch = [i]
            else:
                current_patch.append(i)

        if current_patch:
            patches.append(current_patch)

        return patches


# ============================================================================
# Legacy Single-Constraint Functions (for backward compatibility)
# ============================================================================

def patch_global_threshold(metric: np.ndarray, threshold: float) -> List[List[int]]:
    """
    Patch when metric exceeds threshold.

    Args:
        metric: Entropy or perplexity values
        threshold: Threshold value

    Returns:
        List of patches (each patch is list of token indices)
    """
    patches = []
    current_patch = []

    for i, value in enumerate(metric):
        current_patch.append(i)
        if value > threshold:
            patches.append(current_patch)
            current_patch = []

    if current_patch:
        patches.append(current_patch)

    return patches


def patch_additive_monotonic(metric: np.ndarray, threshold: float) -> List[List[int]]:
    """
    Patch when metric increase exceeds threshold.

    Args:
        metric: Entropy or perplexity values
        threshold: Threshold for metric increase

    Returns:
        List of patches
    """
    patches = []
    current_patch = [0]

    for i in range(1, len(metric)):
        delta = metric[i] - metric[i-1]

        if delta > threshold:
            patches.append(current_patch)
            current_patch = [i]
        else:
            current_patch.append(i)

    if current_patch:
        patches.append(current_patch)

    return patches


def patch_multiplicative_monotonic(metric: np.ndarray, threshold: float) -> List[List[int]]:
    """
    Patch when metric ratio exceeds threshold.

    Args:
        metric: Entropy or perplexity values
        threshold: Threshold for metric ratio

    Returns:
        List of patches
    """
    patches = []
    current_patch = [0]

    for i in range(1, len(metric)):
        if metric[i-1] > 1e-6:  # Avoid division by zero
            ratio = metric[i] / metric[i-1]
        else:
            ratio = 1.0

        if ratio > threshold:
            patches.append(current_patch)
            current_patch = [i]
        else:
            current_patch.append(i)

    if current_patch:
        patches.append(current_patch)

    return patches


def patch_additive_min(metric: np.ndarray, threshold: float) -> List[List[int]]:
    """
    Patch when metric increase from minimum in current patch exceeds threshold.

    Args:
        metric: Entropy or perplexity values
        threshold: Threshold for metric increase from patch minimum

    Returns:
        List of patches
    """
    patches = []
    current_patch = [0]
    patch_min = metric[0]

    for i in range(1, len(metric)):
        delta = metric[i] - patch_min

        if delta > threshold:
            patches.append(current_patch)
            current_patch = [i]
            patch_min = metric[i]
        else:
            current_patch.append(i)
            patch_min = min(patch_min, metric[i])

    if current_patch:
        patches.append(current_patch)

    return patches


def patch_multiplicative_min(metric: np.ndarray, threshold: float) -> List[List[int]]:
    """
    Patch when metric ratio to minimum in current patch exceeds threshold.

    Args:
        metric: Entropy or perplexity values
        threshold: Threshold for metric ratio to patch minimum

    Returns:
        List of patches
    """
    patches = []
    current_patch = [0]
    patch_min = metric[0]

    for i in range(1, len(metric)):
        if patch_min > 1e-6:  # Avoid division by zero
            ratio = metric[i] / patch_min
        else:
            ratio = 1.0

        if ratio > threshold:
            patches.append(current_patch)
            current_patch = [i]
            patch_min = metric[i]
        else:
            current_patch.append(i)
            patch_min = min(patch_min, metric[i])

    if current_patch:
        patches.append(current_patch)

    return patches


def apply_patching(
    metric: np.ndarray,
    constraint: str | None = None,
    threshold: float | None = None,
    constraints: Dict[str, float] | None = None,
) -> List[List[int]]:
    """
    Apply patching with single or multiple constraints.

    Args:
        metric: Entropy or perplexity values
        constraint: (Deprecated) Single constraint name for backward compatibility
        threshold: (Deprecated) Single threshold for backward compatibility
        constraints: {"constraint_name": threshold} for multi-constraint patching
                    e.g., {"global": 5.0, "additive": 0.5, "cumulative": 10.0}

    Returns:
        List of patches (each patch is list of token indices)

    Examples:
        # Single constraint (backward compatible)
        patches = apply_patching(metric, "global", 5.0)

        # Multiple constraints (new)
        patches = apply_patching(metric, constraints={"global": 5.0, "additive": 0.5})
    """
    # Backward compatibility: convert old API to new
    if constraint is not None and threshold is not None:
        constraints = {constraint: threshold}

    if constraints is None or len(constraints) == 0:
        raise ValueError("No constraints specified")

    # Use OOP evaluator for multi-constraint support
    evaluator = MultiConstraintEvaluator(constraints)
    boundary_mask = evaluator.evaluate_all(metric)
    return evaluator.to_patches(boundary_mask)


def compute_patch_statistics(patches: List[List[int]]) -> dict:
    """
    Compute statistics about patches.

    Args:
        patches: List of patches

    Returns:
        Dictionary of statistics
    """
    sizes = [len(p) for p in patches]

    return {
        "num_patches": len(patches),
        "avg_size": float(np.mean(sizes)),
        "min_size": int(np.min(sizes)),
        "max_size": int(np.max(sizes)),
        "median_size": float(np.median(sizes)),
        "std_size": float(np.std(sizes)),
        "sizes": sizes,
    }


__all__ = [
    "Constraint",
    "GlobalConstraint",
    "AdditiveMonotonicConstraint",
    "MultiplicativeMonotonicConstraint",
    "AdditiveMinConstraint",
    "MultiplicativeMinConstraint",
    "CumulativeSumConstraint",
    "CumulativeMaxConstraint",
    "MinimumMetricConstraint",
    "MultiConstraintEvaluator",
    "apply_patching",
    "compute_patch_statistics",
    "patch_global_threshold",
    "patch_additive_monotonic",
    "patch_multiplicative_monotonic",
    "patch_additive_min",
    "patch_multiplicative_min",
]