#!/usr/bin/env python3
"""
Helix64 Grammar Module

Encodes the valid operation sequences for Helix64 programs.
Discovered empirically: the VM has strict sequential dependencies:

    A: INIT         - Initialize state
    B: LOAD_W1      - Load weights (required before matmul)
    C: LOAD_B1      - Load biases (required before add_bias)
    D: MATMUL_1     - Matrix multiply (requires W1)
    E: ADD_BIAS_1   - Add bias + ReLU (requires b1)
    F: LOAD_W2      - Load second layer weights
    G: MATMUL_2     - Second matrix multiply
    H: OUTPUT       - Finalize output

Current discovery: Only ABCDEFGH is valid in the strict VM.
Future: This grammar can be relaxed as the VM becomes more flexible.

Usage:
    from helix_cdc.helix_grammar import (
        is_valid_sequence,
        validate_sequence,
        nearest_valid_sequence,
        VALID_OPS_BY_POSITION,
    )

    # Check if a sequence is valid
    is_valid_sequence("ABCDEFGH")  # True
    is_valid_sequence("BACDFEGH")  # False

    # Get validation errors
    errors = validate_sequence("BACDFEGH")
    # ["Position 0: 'B' not in allowed set {'A'}",
    #  "Position 1: 'A' not in allowed set {'B'}", ...]

    # Snap invalid sequence to nearest valid
    nearest_valid_sequence("HGFEDCBA")  # "ABCDEFGH"
"""

from __future__ import annotations

from typing import Dict, List, Optional, Set, Tuple

# Canonical codons
CODONS = list("ABCDEFGH")
PROGRAM_LENGTH = 8

# Semantic meaning of each codon (from CDR programs.py canonical mapping)
# These define what each operation DOES in the Helix8 execution model
CODON_SEMANTICS = {
    "A": {
        "name": "CONTEXT_GATHER",
        "description": "Gather context from input stream",
        "vm_op": "INIT",
        "cdr_role": "Symbol ingestion and context preparation",
    },
    "B": {
        "name": "EMBEDDING",
        "description": "Project input into embedding space",
        "vm_op": "LOAD_W1",
        "cdr_role": "Map symbols to dense representations",
    },
    "C": {
        "name": "PROJECTION_UP",
        "description": "Project embeddings to higher dimension",
        "vm_op": "LOAD_B1",
        "cdr_role": "Expand representational capacity",
    },
    "D": {
        "name": "LAYER_NORM",
        "description": "Normalize activations for stability",
        "vm_op": "MATMUL_1",
        "cdr_role": "Stabilize information flow",
    },
    "E": {
        "name": "BIAS_RELU_GATE",
        "description": "Add bias + ReLU activation (ANCHOR GATE)",
        "vm_op": "ADD_BIAS_1",
        "cdr_role": "Gate position - controls information flow",
        "is_gate": True,
    },
    "F": {
        "name": "W1_MATMUL",
        "description": "First weight matrix multiply",
        "vm_op": "LOAD_W2",
        "cdr_role": "Transform through learned weights",
    },
    "G": {
        "name": "SOFTMAX",
        "description": "Apply softmax for probability distribution",
        "vm_op": "MATMUL_2",
        "cdr_role": "Normalize to probability space",
    },
    "H": {
        "name": "EMIT_LOGITS",
        "description": "Emit final logits/output",
        "vm_op": "OUTPUT",
        "cdr_role": "Return result to caller",
    },
}

# Position semantics (what role each position plays in the Helix8 pipeline)
# From CDR programs.py canonical definition
POSITION_SEMANTICS = {
    0: {
        "codon": "A",
        "role": "context_gather",
        "description": "Input processing and context preparation",
        "phase": "input",
    },
    1: {
        "codon": "B",
        "role": "embedding",
        "description": "Symbol-to-vector embedding",
        "phase": "encode",
    },
    2: {
        "codon": "C",
        "role": "projection_up",
        "description": "Dimension expansion",
        "phase": "encode",
    },
    3: {
        "codon": "D",
        "role": "layer_norm",
        "description": "Activation normalization",
        "phase": "transform",
    },
    4: {
        "codon": "E",
        "role": "anchor_gate",
        "description": "Information gating (tunable via gate_policy)",
        "phase": "gate",
        "is_tunable": True,
    },
    5: {
        "codon": "F",
        "role": "w1_matmul",
        "description": "Weight transformation",
        "phase": "transform",
    },
    6: {
        "codon": "G",
        "role": "softmax",
        "description": "Probability distribution",
        "phase": "output",
    },
    7: {
        "codon": "H",
        "role": "emit_logits",
        "description": "Final output emission",
        "phase": "output",
    },
}

# Valid codons per position (empirically discovered from VM errors)
# Current: strict ordering required
# Future: can be relaxed as VM becomes more flexible
VALID_OPS_BY_POSITION: Dict[int, Set[str]] = {
    0: {"A"},  # Must initialize first
    1: {"B"},  # Must load weights before matmul
    2: {"C"},  # Must load biases before add_bias
    3: {"D"},  # Matmul after weights loaded
    4: {"E"},  # Add bias after biases loaded (the "gate" position)
    5: {"F"},  # Second layer weights
    6: {"G"},  # Second matmul
    7: {"H"},  # Output
}

# Dependencies: which ops must come before others
CODON_DEPENDENCIES: Dict[str, List[str]] = {
    "A": [],           # A has no dependencies
    "B": ["A"],        # B needs A (init)
    "C": ["A"],        # C needs A (init)
    "D": ["B"],        # D needs B (weights)
    "E": ["C", "D"],   # E needs C (biases) and D (matmul result)
    "F": ["E"],        # F needs E (previous layer output)
    "G": ["F"],        # G needs F (second weights)
    "H": ["G"],        # H needs G (second matmul result)
}

# The canonical valid program (the one that works)
CANONICAL_PROGRAM = "ABCDEFGH"


def is_valid_codon(codon: str) -> bool:
    """Check if a single character is a valid codon."""
    return codon in CODONS


def is_valid_at_position(codon: str, position: int) -> bool:
    """Check if a codon is valid at a specific position."""
    if position < 0 or position >= PROGRAM_LENGTH:
        return False
    if position not in VALID_OPS_BY_POSITION:
        return True  # No constraint defined
    return codon in VALID_OPS_BY_POSITION[position]


def is_valid_sequence(seq: str) -> bool:
    """
    Check if a program sequence is valid according to the grammar.

    Args:
        seq: 8-character codon sequence

    Returns:
        True if valid, False otherwise
    """
    if not isinstance(seq, str) or len(seq) != PROGRAM_LENGTH:
        return False

    for i, codon in enumerate(seq):
        if not is_valid_codon(codon):
            return False
        if not is_valid_at_position(codon, i):
            return False

    return True


def validate_sequence(seq: str) -> List[str]:
    """
    Validate a sequence and return list of errors.

    Args:
        seq: Sequence to validate

    Returns:
        List of error strings (empty if valid)
    """
    errors: List[str] = []

    if not isinstance(seq, str):
        errors.append(f"Sequence must be a string, got {type(seq).__name__}")
        return errors

    if len(seq) != PROGRAM_LENGTH:
        errors.append(f"Sequence must be {PROGRAM_LENGTH} characters, got {len(seq)}")
        return errors

    for i, codon in enumerate(seq):
        if not is_valid_codon(codon):
            errors.append(f"Position {i}: '{codon}' is not a valid codon (must be A-H)")
        elif not is_valid_at_position(codon, i):
            allowed = VALID_OPS_BY_POSITION.get(i, set(CODONS))
            errors.append(f"Position {i}: '{codon}' not in allowed set {allowed}")

    return errors


def nearest_valid_sequence(seq: str) -> str:
    """
    Convert an invalid sequence to the nearest valid one.

    Currently: Since only ABCDEFGH is valid, this always returns that.
    Future: Could implement actual nearest-neighbor search.

    Args:
        seq: Input sequence (may be invalid)

    Returns:
        Valid sequence
    """
    if is_valid_sequence(seq):
        return seq

    # For now, snap to canonical program
    # Future: could compute edit distance and find nearest valid
    return CANONICAL_PROGRAM


def get_valid_mutations(seq: str, position: int) -> List[str]:
    """
    Get all valid single-position mutations at a given position.

    Args:
        seq: Current program sequence
        position: Position to mutate (0-7)

    Returns:
        List of valid mutated sequences
    """
    if position < 0 or position >= PROGRAM_LENGTH:
        return []

    allowed = VALID_OPS_BY_POSITION.get(position, set(CODONS))
    current = seq[position] if position < len(seq) else ""

    results = []
    for codon in allowed:
        if codon != current:
            mutated = seq[:position] + codon + seq[position + 1:]
            if is_valid_sequence(mutated):
                results.append(mutated)

    return results


def get_mutable_positions(strict: bool = True) -> List[int]:
    """
    Get positions that can be safely mutated.

    Args:
        strict: If True, only return positions with multiple valid options.
                If False, return all positions.

    Returns:
        List of position indices
    """
    if strict:
        # Only positions with more than one valid codon
        return [
            pos for pos, allowed in VALID_OPS_BY_POSITION.items()
            if len(allowed) > 1
        ]
    else:
        return list(range(PROGRAM_LENGTH))


def describe_program(seq: str) -> str:
    """
    Generate a human-readable description of a program.

    Args:
        seq: Program sequence

    Returns:
        Multi-line description string
    """
    lines = [f"Program: {seq}", ""]

    for i, codon in enumerate(seq):
        pos_info = POSITION_SEMANTICS.get(i, {})
        codon_info = CODON_SEMANTICS.get(codon, {})

        role = pos_info.get("role", "unknown")
        codon_name = codon_info.get("name", codon)
        codon_desc = codon_info.get("description", "")

        valid = is_valid_at_position(codon, i)
        status = "OK" if valid else "INVALID"

        lines.append(f"  [{i}] {codon} ({codon_name}): {codon_desc} [{status}]")

    errors = validate_sequence(seq)
    if errors:
        lines.append("")
        lines.append("Errors:")
        for e in errors:
            lines.append(f"  - {e}")

    return "\n".join(lines)


def get_grammar_summary() -> Dict:
    """
    Get a summary of the current grammar constraints.

    Returns:
        Dict with grammar information
    """
    total_programs = 1
    for pos in range(PROGRAM_LENGTH):
        allowed = len(VALID_OPS_BY_POSITION.get(pos, set(CODONS)))
        total_programs *= allowed

    return {
        "program_length": PROGRAM_LENGTH,
        "codons": CODONS,
        "valid_ops_by_position": {
            pos: list(allowed) for pos, allowed in VALID_OPS_BY_POSITION.items()
        },
        "canonical_program": CANONICAL_PROGRAM,
        "total_valid_programs": total_programs,
        "mutable_positions": get_mutable_positions(strict=True),
        "note": "Current grammar is strict - only ABCDEFGH is valid. "
                "Future versions may relax constraints.",
    }


# Quick self-test
if __name__ == "__main__":
    print("=== Helix64 Grammar Module ===\n")

    print("Grammar Summary:")
    summary = get_grammar_summary()
    for k, v in summary.items():
        print(f"  {k}: {v}")
    print()

    print("Canonical program description:")
    print(describe_program(CANONICAL_PROGRAM))
    print()

    print("Testing validation:")
    test_seqs = ["ABCDEFGH", "HGFEDCBA", "AAAAAAAA", "ABCD"]
    for seq in test_seqs:
        valid = is_valid_sequence(seq)
        errors = validate_sequence(seq)
        print(f"  {seq!r}: valid={valid}, errors={len(errors)}")
    print()

    print("Nearest valid for 'HGFEDCBA':", nearest_valid_sequence("HGFEDCBA"))
