inference_engine_v6.py

"""
inference_engine_v6.py — Hatch v6: Positional Terminal De-weighting
====================================================================

v6 adds Positional Weighting on top of v5's architecture. All other
state machine logic (Hysteresis, Super-Hatch, AUC weights, CD thresholds)
is identical to v5.

Key change from v5:
  - Each window's weighted score is multiplied by a positional weight w(i)
    before comparison to T_fill and T_drain.
  - Terminal windows (first and last TERMINAL_FRACTION of the sequence)
    have reduced weight, preventing "End-Frustration" false DISORDERED calls.
  - The positional weight uses a smooth sigmoid: w(i) = 1 - alpha * sigmoid_decay(i, L)
    where alpha controls the maximum de-weighting at the termini.

Motivation (End-Frustration problem):
  Many folded proteins have disordered N/C-terminal tails (signal peptides,
  purification tags, intrinsically disordered regions). At W=30, these tails
  generate fill events that push the cup toward overflow, causing false
  DISORDERED predictions on otherwise folded proteins. Positional de-weighting
  reduces the contribution of terminal windows without hard-cutting them.

Positional weight formula:
  For a sequence of length L with n_windows windows:
    pos_weight(i) = 1 - alpha * (
        exp(-i / (TERMINAL_FRACTION * n_windows)) +
        exp(-(n_windows - 1 - i) / (TERMINAL_FRACTION * n_windows))
    )
  where:
    i = window index (0 to n_windows-1)
    alpha = max de-weighting factor (0.0 = no effect, 1.0 = full suppression)
    TERMINAL_FRACTION = fraction of sequence treated as terminal (e.g., 0.15)

  This gives a smooth U-shaped weight curve: 1.0 in the middle, reduced at ends.

Effect on state machine:
  - If pos_weight * score_w < T_fill  → fill event (cup += 1)
  - If pos_weight * score_w >= T_drain → potential drain event
  - Super-Hatch fires on the raw score (not positionally weighted) — a perfect
    window is a perfect window regardless of position.

v6 performance (grid_search_v6 validation, 400 sequences):
  Baseline v5: Mean F1=75.40%, F1 Folded=77.00%, F1 Disordered=73.80%
  v6 target:   Mean F1 > 77%
"""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Optional

import numpy as np

from inference_engine_v5 import (
    HatchClassifierV5,
    V5_OPTIMIZED_THRESHOLDS,
    V5_FOLDED_IS_HIGH,
    V5_DEFAULT_CONFIG,
    MAX_WEIGHTED_SCORE_V5,
    _compute_features,
)
from inference_engine_v4 import DEFAULT_WEIGHTS


# ─── Positional weighting parameters ─────────────────────────────────────────

TERMINAL_FRACTION = 0.15   # fraction of sequence treated as terminal
DEFAULT_ALPHA     = 0.40   # max de-weighting at the termini (40% reduction)


def positional_weights(n_windows: int,
                       terminal_fraction: float = TERMINAL_FRACTION,
                       alpha: float = DEFAULT_ALPHA) -> np.ndarray:
    """
    Compute per-window positional weights.

    Returns a 1D array of shape (n_windows,) with values in [1-alpha, 1.0].
    The weight is 1.0 in the middle of the sequence and (1 - alpha) at the ends.
    """
    if n_windows <= 0:
        return np.array([], dtype=np.float64)

    decay_len = max(terminal_fraction * n_windows, 1.0)
    i = np.arange(n_windows, dtype=np.float64)

    # Exponential decay from both ends
    left_decay  = np.exp(-i / decay_len)
    right_decay = np.exp(-(n_windows - 1 - i) / decay_len)

    weights = 1.0 - alpha * (left_decay + right_decay)
    # Clip to [1 - alpha, 1.0] to prevent negative weights
    weights = np.clip(weights, 1.0 - alpha, 1.0)
    return weights


# ─── HatchClassifierV6 ────────────────────────────────────────────────────────

class HatchClassifierV6:
    """
    Hatch v6: Positional Terminal De-weighting + all v5 features.

    Adds a smooth positional weight to each window's score before comparing
    to T_fill and T_drain. Terminal windows contribute less to cup filling,
    reducing false DISORDERED calls on proteins with disordered tails.
    """

    def __init__(
        self,
        thresholds:        Optional[np.ndarray] = None,
        weights:           Optional[np.ndarray] = None,
        window_size:       int   = V5_DEFAULT_CONFIG["window_size"],
        K:                 int   = V5_DEFAULT_CONFIG["K"],
        T_fill:            float = V5_DEFAULT_CONFIG["T_fill"],
        T_drain:           float = V5_DEFAULT_CONFIG["T_drain"],
        N:                 int   = V5_DEFAULT_CONFIG["N"],
        D:                 int   = V5_DEFAULT_CONFIG["D"],
        super_hatch:       bool  = V5_DEFAULT_CONFIG["super_hatch"],
        step:              int   = V5_DEFAULT_CONFIG["step"],
        terminal_fraction: float = TERMINAL_FRACTION,
        alpha:             float = DEFAULT_ALPHA,
    ):
        self.thresholds        = thresholds if thresholds is not None else V5_OPTIMIZED_THRESHOLDS.copy()
        self.weights           = weights if weights is not None else DEFAULT_WEIGHTS.copy()
        self.window_size       = window_size
        self.K                 = K
        self.T_fill            = T_fill
        self.T_drain           = T_drain
        self.N                 = N
        self.D                 = D
        self.super_hatch       = super_hatch
        self.step              = step
        self.terminal_fraction = terminal_fraction
        self.alpha             = alpha

        self._folded_is_high = V5_FOLDED_IS_HIGH.copy()
        self._max_score      = float(self.weights.sum())

    # ── Weighted scoring ──────────────────────────────────────────────────────

    def _score_window(self, features: np.ndarray) -> float:
        folded_conditions = np.where(
            self._folded_is_high,
            features > self.thresholds,
            features < self.thresholds,
        ).astype(np.float64)
        return float(np.dot(self.weights, folded_conditions))

    # ── Fast path (classify only) ─────────────────────────────────────────────

    def classify(self, sequence: str) -> int:
        """Returns 1 (FOLDED) or 0 (DISORDERED)."""
        n = len(sequence)
        if n < self.window_size:
            return 1

        # Precompute positional weights
        positions    = list(range(0, n - self.window_size + 1, self.step))
        n_windows    = len(positions)
        pos_weights  = positional_weights(n_windows, self.terminal_fraction, self.alpha)

        cup_level           = 0
        consecutive_ordered = 0

        for idx, start in enumerate(positions):
            window   = sequence[start: start + self.window_size]
            features = _compute_features(window)
            score_w  = self._score_window(features)
            pw       = pos_weights[idx]

            # Super-Hatch fires on raw score (position-independent)
            if self.super_hatch and score_w >= self._max_score - 1e-9:
                cup_level           = 0
                consecutive_ordered = 0
            elif score_w * pw < self.T_fill:
                # Positionally-weighted score below fill threshold → clog
                cup_level          += 1
                consecutive_ordered = 0
            elif score_w * pw >= self.T_drain:
                # Positionally-weighted score above drain threshold → potential drain
                consecutive_ordered += 1
                if consecutive_ordered >= self.N:
                    cup_level           = max(0, cup_level - self.D)
                    consecutive_ordered = 0

            if cup_level >= self.K:
                return 0  # DISORDERED

        return 1  # FOLDED

    # ── Trace path (full diagnostics) ─────────────────────────────────────────

    def trace(self, sequence: str) -> dict:
        """Returns full diagnostic trace including positional weights."""
        n = len(sequence)
        if n < self.window_size:
            return {"prediction": 1, "overflow_at": None,
                    "positions": [], "entropy_levels": [],
                    "scores_w": [], "pos_weights": [],
                    "effective_scores": [], "fill_events": [],
                    "drain_events": [], "super_hatch_events": []}

        positions_list = list(range(0, n - self.window_size + 1, self.step))
        n_windows      = len(positions_list)
        pos_weights    = positional_weights(n_windows, self.terminal_fraction, self.alpha)

        cup_level           = 0
        consecutive_ordered = 0
        overflow_at         = None

        out_positions      = []
        entropy_levels     = []
        scores_w           = []
        pw_list            = []
        effective_scores   = []
        fill_events        = []
        drain_events       = []
        sh_events          = []

        for idx, start in enumerate(positions_list):
            window   = sequence[start: start + self.window_size]
            features = _compute_features(window)
            score_w  = self._score_window(features)
            pw       = pos_weights[idx]
            eff      = score_w * pw

            is_sh    = False
            is_fill  = False
            is_drain = False

            if self.super_hatch and score_w >= self._max_score - 1e-9:
                cup_level           = 0
                consecutive_ordered = 0
                is_sh = True
            elif eff < self.T_fill:
                cup_level          += 1
                consecutive_ordered = 0
                is_fill = True
            elif eff >= self.T_drain:
                consecutive_ordered += 1
                if consecutive_ordered >= self.N:
                    cup_level           = max(0, cup_level - self.D)
                    consecutive_ordered = 0
                    is_drain = True

            out_positions.append(start)
            entropy_levels.append(cup_level)
            scores_w.append(score_w)
            pw_list.append(pw)
            effective_scores.append(eff)
            fill_events.append(is_fill)
            drain_events.append(is_drain)
            sh_events.append(is_sh)

            if cup_level >= self.K and overflow_at is None:
                overflow_at = idx

        prediction = 0 if overflow_at is not None else 1
        return {
            "prediction":         prediction,
            "overflow_at":        overflow_at,
            "positions":          out_positions,
            "entropy_levels":     entropy_levels,
            "scores_w":           scores_w,
            "pos_weights":        pw_list,
            "effective_scores":   effective_scores,
            "fill_events":        fill_events,
            "drain_events":       drain_events,
            "super_hatch_events": sh_events,
        }

    def describe(self) -> str:
        return (
            f"HatchV6(W={self.window_size}, K={self.K}, "
            f"T_fill={self.T_fill:.2f}, T_drain={self.T_drain:.2f}, "
            f"N={self.N}, D={self.D}, SH={'ON' if self.super_hatch else 'OFF'}, "
            f"alpha={self.alpha:.2f}, term_frac={self.terminal_fraction:.2f})"
        )


# ─── Convenience factory ──────────────────────────────────────────────────────

def load_v6_classifier(
    terminal_fraction: float = TERMINAL_FRACTION,
    alpha:             float = DEFAULT_ALPHA,
    **kwargs,
) -> HatchClassifierV6:
    """Load the default v6 classifier with v5 optimized thresholds."""
    return HatchClassifierV6(
        terminal_fraction = terminal_fraction,
        alpha             = alpha,
        **kwargs,
    )