#!/usr/bin/env python3
"""
Morph table for token compression and SAT guidance.

Morph table maps common token sequences to compressed representations,
enabling token savings at iso-quality.
"""
import json
from pathlib import Path
from typing import Dict, List, Optional

class MorphTable:
    """
    Morph table for token compression.

    Maps common sequences to shorter representations, reducing token count
    while preserving semantic quality (validated with KL divergence).
    """

    def __init__(self, table_path: Optional[str] = None):
        """
        Initialize morph table.

        Args:
            table_path: Path to morph table JSON (optional)
        """
        self.table = {}
        self.inverse_table = {}

        if table_path:
            self.load(table_path)
        else:
            # Default: common patterns that can be compressed
            # Format: {original_sequence: compressed_token}
            self.table = {
                # Common technical phrases (examples)
                "in order to": "to",
                "due to the fact that": "because",
                "at this point in time": "now",
                # Add more from actual usage patterns
            }
            self.inverse_table = {v: k for k, v in self.table.items()}

    def load(self, table_path: str):
        """Load morph table from JSON file."""
        path = Path(table_path)
        if path.exists():
            data = json.loads(path.read_text())
            self.table = data.get("table", {})
            self.inverse_table = {v: k for k, v in self.table.items()}

    def save(self, table_path: str):
        """Save morph table to JSON file."""
        data = {
            "version": "1.0",
            "table": self.table,
            "count": len(self.table)
        }
        Path(table_path).write_text(json.dumps(data, indent=2))

    def compress_tokens(self, tokens: List[str]) -> List[str]:
        """
        Compress token sequence using morph table.

        Args:
            tokens: Input token sequence

        Returns:
            Compressed token sequence
        """
        compressed = []
        i = 0
        while i < len(tokens):
            # Try to match longest sequence first
            matched = False
            for length in range(min(5, len(tokens) - i), 0, -1):
                sequence = " ".join(tokens[i:i+length])
                if sequence in self.table:
                    compressed.append(self.table[sequence])
                    i += length
                    matched = True
                    break

            if not matched:
                compressed.append(tokens[i])
                i += 1

        return compressed

    def decompress_tokens(self, compressed: List[str]) -> List[str]:
        """
        Decompress token sequence using inverse morph table.

        Args:
            compressed: Compressed token sequence

        Returns:
            Original token sequence
        """
        decompressed = []
        for token in compressed:
            if token in self.inverse_table:
                decompressed.extend(self.inverse_table[token].split())
            else:
                decompressed.append(token)

        return decompressed

    def compute_savings(self, original: List[str], compressed: List[str]) -> dict:
        """
        Compute compression statistics.

        Args:
            original: Original token sequence
            compressed: Compressed token sequence

        Returns:
            Stats dict with savings percentage, etc.
        """
        orig_len = len(original)
        comp_len = len(compressed)
        saved = orig_len - comp_len
        saved_pct = (saved / orig_len * 100) if orig_len > 0 else 0.0

        return {
            "original_tokens": orig_len,
            "compressed_tokens": comp_len,
            "tokens_saved": saved,
            "saved_pct": saved_pct
        }

def test_morph_table():
    """Test morph table."""
    table = MorphTable()

    # Test sequence
    tokens = ["in", "order", "to", "reduce", "energy", "consumption"]

    # Compress
    compressed = table.compress_tokens(tokens)
    stats = table.compute_savings(tokens, compressed)

    print(f"Original: {len(tokens)} tokens: {' '.join(tokens)}")
    print(f"Compressed: {len(compressed)} tokens: {' '.join(compressed)}")
    print(f"Saved: {stats['saved_pct']:.1f}%")

if __name__ == "__main__":
    test_morph_table()
