#!/usr/bin/env python3
# helix_cdc/morph_guidance.py
from dataclasses import dataclass
from typing import List, Tuple, Dict, Optional
import torch
from transformers import LogitsProcessor

@dataclass(frozen=True)
class MorphRule:
    src_text: str
    dst_text: str
    bias: float = 2.6
    forbid_src_cont: bool = True
    forbid_bias: float = -2.6

def _enc(tok, text: str) -> List[int]:
    # leading space keeps BPE alignment sane
    return tok.encode(" " + text, add_special_tokens=False)

def compile_rules(tok, rules: List[MorphRule]):
    compiled = []
    for r in rules:
        # two variants to survive boundary conditions
        src = _enc(tok, r.src_text)
        src_ns = tok.encode(r.src_text, add_special_tokens=False)
        dst = _enc(tok, r.dst_text)
        if not src or not dst:
            continue
        compiled.append((src, dst, r))
        if src_ns != src and src_ns:
            compiled.append((src_ns, dst, r))
    return compiled

class MorphBiasProcessor(LogitsProcessor):
    """
    Robust morph guidance with suffix matching for both src and dst.
    - If context ends with full src => bias toward dst[0], block src continuation
    - If context ends with partial dst => nudge toward next dst token
    - If context ends with partial src => block src continuation
    """
    def __init__(self, tokenizer, compiled_rules, se=None, depth_target: Optional[float]=None, debug: bool=False):
        self.tok = tokenizer
        self.rules = compiled_rules
        self.se = se
        self.depth_target = depth_target or None
        self.debug = debug
        self._events: Dict[int, list] = {}

    def _match_suffix(self, ids: List[int], pat: List[int]) -> int:
        # returns matched suffix length [0..len(pat)]
        L = min(len(ids), len(pat))
        for k in range(L, 0, -1):
            if ids[-k:] == pat[:k]:
                return k
        return 0

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
        B, T = input_ids.shape
        V = scores.shape[-1]
        for b in range(B):
            ids = input_ids[b].tolist()

            # If we have Se, feed it the pre-sampling logits for this step
            if self.se is not None:
                self.se.step(ids, scores[b])

            # Dynamic bias scale based on depth gap
            bias_scale = 1.0
            if self.se is not None and self.depth_target:
                D = self.se.metrics().get("D", 1.0)
                # if depth below target, lean a bit more; clamp to avoid over-steer
                bias_scale = max(0.7, min(1.5, (self.depth_target / max(D, 1e-6)) ** 0.25))

            # track events if debug
            if self.debug:
                self._events.setdefault(b, [])

            for src, dst, rule in self.rules:
                m_src = self._match_suffix(ids, src)
                m_dst = self._match_suffix(ids, dst)

                # if we just saw the full long form, start compressed path
                if m_src == len(src):
                    scores[b, dst[0]] += rule.bias * bias_scale
                    if rule.forbid_src_cont and len(src) >= 1:
                        scores[b, src[-1]] += rule.forbid_bias
                    # Log matches if debug enabled
                    if self.debug:
                        import sys, os
                        print(f"✅ RULE FIRED: {self.tok.decode(src)} → {self.tok.decode(dst)} (bias={rule.bias*bias_scale:.2f})", file=sys.stderr)
                        self._events[b].append({"event":"src_full","src":src,"dst":dst,"bias":rule.bias*bias_scale})

                # if we're partway through short form, gently favor next token
                if 0 < m_dst < len(dst):
                    scores[b, dst[m_dst]] += 1.2 * bias_scale
                    if self.debug:
                        self._events[b].append({"event":"dst_progress","at":m_dst,"nudge":1.2*bias_scale})

                # discourage drifting deeper into the long form mid-phrase
                if 0 < m_src < len(src) and rule.forbid_src_cont:
                    scores[b, src[m_src]] += rule.forbid_bias
                    if self.debug:
                        self._events[b].append({"event":"src_block","at":m_src,"forbid":rule.forbid_bias})

        return scores
