#!/usr/bin/env python3
"""
morphsat.py - MorphSAT Token Adjacency Constraints for Helix-CDC

Ported from echo-box/runtimes/flowtorch/morphsat.py
Uses numpy instead of torch for compatibility with helix_infer_glyph.py

MorphSAT enforces soft constraints on token adjacency:
- Which tokens can follow which
- Position-aware masking (pos_mod4 for base64 rhythm)
- Soft penalties rather than hard masks

Integration with vocab_256:
    Lane 0 (ENTITY)   → Nouns/objects
    Lane 1 (ACTION)   → Verbs
    Lane 2 (QUALITY)  → Adjectives/modifiers
    Lane 3 (RELATION) → Prepositions/connectors

The default adjacency rules enforce natural sentence flow:
    ENTITY → ACTION preferred (noun → verb)
    ACTION → QUALITY preferred (verb → adjective)
    QUALITY → RELATION preferred (adjective → preposition)
    RELATION → ENTITY preferred (preposition → noun)
"""

from __future__ import annotations

import json
import math
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Set, Any, Union


# Default adjacency rules based on 4-lane semantic structure
# Format: lane -> list of preferred next lanes
DEFAULT_LANE_ADJACENCY = {
    0: [1, 2],      # ENTITY → ACTION or QUALITY
    1: [2, 3, 0],   # ACTION → QUALITY, RELATION, or ENTITY
    2: [3, 0, 1],   # QUALITY → RELATION, ENTITY, or ACTION
    3: [0, 1],      # RELATION → ENTITY or ACTION
}

# Lane names for reference
LANE_NAMES = ["ENTITY", "ACTION", "QUALITY", "RELATION"]


class MorphSAT:
    """
    Token adjacency constraint system for Helix-CDC.

    Applies soft penalties to tokens that violate adjacency rules,
    respecting the 4-lane semantic structure of vocab_256.

    Args:
        morph_table: Optional dict with adjacency rules
            - mode: "simple" or "prev+pos" or "lane"
            - adj: adjacency mappings
            - default: "all" (allow all) or "deny" (block unlisted)
        sat_T0: Initial temperature for annealing
        sat_Tmin: Minimum temperature
        steps_anneal: Steps to reach minimum temperature
        soft_lambda: Soft penalty magnitude
    """

    def __init__(
        self,
        morph_table: Optional[Dict[str, Any]] = None,
        sat_T0: float = 1.0,
        sat_Tmin: float = 0.5,
        steps_anneal: int = 200,
        soft_lambda: float = 4.0,
    ):
        self.mt = morph_table
        self.T0 = sat_T0
        self.Tmin = sat_Tmin
        self.N = max(1, steps_anneal)
        self.soft_lambda = soft_lambda
        self.step = 0

    def _allowed_set(
        self, last_token: int, pos_mod4: int, vocab_size: int
    ) -> Optional[List[int]]:
        """
        Get allowed token indices for (prev_token, position_mod4).

        Returns None to allow all tokens, or list of allowed indices.
        """
        if self.mt is None:
            # Use default lane-based adjacency
            return self._lane_based_allowed(last_token, pos_mod4, vocab_size)

        mode = self.mt.get("mode", "simple")

        if mode == "prev+pos":
            key = f"{int(last_token)}:{int(pos_mod4)}"
            allowed = self.mt.get("adj", {}).get(key)
            if allowed is None:
                return None if self.mt.get("default", "all") == "all" else []
            return allowed

        if mode == "lane":
            # Lane-based mode: use lane adjacency rules
            return self._lane_based_allowed(last_token, pos_mod4, vocab_size)

        # Simple mode: just check prev token
        allowed = self.mt.get(str(int(last_token)))
        if allowed is None and self.mt.get("default", "all") == "all":
            return None
        return allowed or []

    def _lane_based_allowed(
        self, last_token: int, pos_mod4: int, vocab_size: int
    ) -> Optional[List[int]]:
        """
        Lane-based adjacency using vocab_256 structure.

        The 4-lane structure (ENTITY/ACTION/QUALITY/RELATION) creates
        natural sentence flow when adjacency is respected.
        """
        # Get lane of last token (token_idx mod 4 in the original b64 space)
        last_lane = last_token % 4

        # Get preferred next lanes from adjacency table
        preferred_lanes = DEFAULT_LANE_ADJACENCY.get(last_lane, [0, 1, 2, 3])

        # For vocab_256: tokens 0-63 are base indices
        # Position determines which lane (pos_mod4)
        # So we allow all base indices that, at the current position,
        # would produce tokens in preferred lanes

        # If current position is pos_mod4, the token will be in lane pos_mod4
        # We want pos_mod4 to be in preferred_lanes for the last token's lane
        current_lane = pos_mod4

        if current_lane in preferred_lanes:
            # Current position lane is preferred - allow all
            return None

        # Current lane not preferred - still allow but with lower priority
        # Return None to allow all, but score() will apply soft penalty
        return None

    def score(
        self,
        prev_token: Optional[str],
        next_token: str,
        prev_idx: Optional[int],
        next_idx: int,
        position: int,
    ) -> float:
        """
        Score a token transition.

        Args:
            prev_token: Previous token word (or None if first)
            next_token: Next token word
            prev_idx: Previous token base index (0-63) - NOT used for lane calculation
            next_idx: Next token base index (0-63) - NOT used for lane calculation
            position: Current position in sequence

        Returns:
            Score multiplier (1.0 = full score, <1.0 = penalized)

        Note: Lane is determined by POSITION (pos % 4), not by base index.
              Position 0,4,8,... = ENTITY
              Position 1,5,9,... = ACTION
              Position 2,6,10,... = QUALITY
              Position 3,7,11,... = RELATION
        """
        if position == 0:
            return 1.0  # No constraint on first token

        # Lane is determined by position, not by base index
        prev_lane = (position - 1) % 4
        curr_lane = position % 4

        # Get preferred lanes for prev position's lane
        preferred = DEFAULT_LANE_ADJACENCY.get(prev_lane, [0, 1, 2, 3])

        # Check if current lane is preferred
        if curr_lane in preferred:
            return 1.0

        # Apply soft penalty for non-preferred transition
        T = self.temperature()
        self.step += 1
        return math.exp(-self.soft_lambda / T)

    def mask_scores(
        self,
        scores: np.ndarray,
        last_token_idx: Optional[int],
        pos_mod4: int,
    ) -> np.ndarray:
        """
        Apply soft penalties to a score array.

        Args:
            scores: Array of shape [vocab_size] with token scores
            last_token_idx: Previous token index (0-63) or None
            pos_mod4: Current position mod 4

        Returns:
            Modified scores with penalties applied
        """
        if self.mt is None and last_token_idx is None:
            return scores

        vocab_size = len(scores)
        allowed = self._allowed_set(
            last_token_idx if last_token_idx is not None else 0,
            pos_mod4,
            vocab_size,
        )

        if allowed is None:
            # All allowed - apply lane-based soft penalties
            if last_token_idx is not None:
                prev_lane = last_token_idx % 4
                preferred = DEFAULT_LANE_ADJACENCY.get(prev_lane, [0, 1, 2, 3])

                # Current lane determined by position
                curr_lane = pos_mod4

                if curr_lane not in preferred:
                    # Apply soft penalty to all scores
                    T = self.temperature()
                    scores = scores - self.soft_lambda / T

            self.step += 1
            return scores

        # Hard constraints from morph_table
        mask = np.full_like(scores, -self.soft_lambda)
        mask[allowed] = 0.0
        T = self.temperature()
        self.step += 1
        return (scores + mask) / T

    def temperature(self) -> float:
        """Cosine annealing from T0 to Tmin."""
        x = min(1.0, self.step / self.N)
        return self.Tmin + 0.5 * (self.T0 - self.Tmin) * (1 + math.cos(math.pi * (1 - x)))

    def reset(self) -> None:
        """Reset step counter for new sequence."""
        self.step = 0


def load_morph_table(path: Optional[Union[str, Path]]) -> Optional[Dict[str, Any]]:
    """
    Load morph table from JSON file.

    Format:
        {
            "mode": "lane" | "simple" | "prev+pos",
            "default": "all" | "deny",
            "adj": {...}  # adjacency rules
        }
    """
    if not path:
        return None
    p = Path(path)
    if not p.exists():
        return None
    with open(p, "r") as f:
        return json.load(f)


def create_default_morph_table() -> Dict[str, Any]:
    """
    Create a default morph table using 4-lane adjacency.

    This encodes natural English sentence structure:
        - ENTITY (nouns) → ACTION (verbs) or QUALITY (adjectives)
        - ACTION (verbs) → QUALITY, RELATION, or back to ENTITY
        - QUALITY (adjectives) → RELATION (prepositions) or ENTITY
        - RELATION (prepositions) → ENTITY (noun phrases)
    """
    return {
        "mode": "lane",
        "default": "all",
        "lane_adjacency": DEFAULT_LANE_ADJACENCY,
        "lane_names": LANE_NAMES,
        "description": "Default 4-lane semantic adjacency for vocab_256",
    }


# Convenience function for integration with helix_infer_glyph
def score_token_sequence(
    tokens: List[str],
    indices: List[int],
    morphsat: Optional[MorphSAT] = None,
) -> List[Dict[str, Any]]:
    """
    Score a complete token sequence.

    Args:
        tokens: List of token words
        indices: List of base indices (0-63)
        morphsat: MorphSAT instance (creates default if None)

    Returns:
        List of dicts with scoring info for each token
    """
    if morphsat is None:
        morphsat = MorphSAT()
    morphsat.reset()

    results = []
    prev_idx = None

    for pos, (token, idx) in enumerate(zip(tokens, indices)):
        score = morphsat.score(
            prev_token=tokens[pos - 1] if pos > 0 else None,
            next_token=token,
            prev_idx=prev_idx,
            next_idx=idx,
            position=pos,
        )

        lane = pos % 4
        results.append({
            "pos": pos,
            "token": token,
            "base_idx": idx,
            "lane": lane,
            "lane_name": LANE_NAMES[lane],
            "score": score,
            "penalized": score < 1.0,
        })

        prev_idx = idx

    return results
