# file: token_encoder.py
"""Bidirectional tokenizer for math specification trees.

Encodes nested math spec dicts into flat token sequences with
structural boundary markers so that decode deterministically
reconstructs the original structure.  Supports a dynamic name
registry for variable/function/class identifiers not in the
static vocabulary.
"""

from __future__ import annotations

import json
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union


class MathSpecTokenizer:
    """Converts mathematical specifications to token sequences and back.

    Encoding uses structural separator tokens (KEY_SEP, LIST_START,
    LIST_END, DICT_START, DICT_END) embedded in the flattened token
    stream so the decoder can reconstruct dicts, lists, and nested
    structures deterministically.

    Dynamic names (variables, function names, class names, arbitrary
    numeric constants) are stored in an on-the-fly registry that maps
    Python objects to numeric IDs in the range 500–4999.  The static
    vocabulary occupies IDs 0–199.
    """

    # ---------- static vocabulary layout ----------
    #   0         PAD
    #   1         START
    #   2         END
    #   3         UNK
    #   4-15      math ops      (ADD, SUB, MUL, DIV, MOD, POW, EQ, NEQ, LT, LTE, GT, GTE)
    #  20-29      type / value  (function, return, variable, constant, list, dict, string, number, boolean, comp)
    #  30-35      control flow  (if, else, for, while, loop)
    #  40-43      common        (input, output, condition, body)
    #  46-49      value markers (VALUE_INT, VALUE_FLOAT, VALUE_STR, VALUE_BOOL)
    #  50         KEY_SEP
    #  52-55      structure     (LIST_START, LIST_END, DICT_START, DICT_END)
    #  56-125     structural keys + extra node types
    # 500-4999    dynamic value registry

    # -- special tokens --------------------------------------------------
    PAD: int = 0
    START: int = 1
    END: int = 2
    UNK: int = 3

    # -- value-type markers (prefix tokens for dynamic values) -----------
    VALUE_INT: int = 46
    VALUE_FLOAT: int = 47
    VALUE_STR: int = 48
    VALUE_BOOL: int = 49

    # -- structural boundary markers -------------------------------------
    KEY_SEP: int = 50          # separates dict key from its value
    LIST_START: int = 52
    LIST_END: int = 53
    DICT_START: int = 54
    DICT_END: int = 55

    # -- dynamic registry range ------------------------------------------
    VALUE_ID_MIN: int = 500
    VALUE_ID_MAX: int = 4999

    def __init__(self, vocab_size: int = 5000) -> None:
        """Build static vocabulary and initialise dynamic registries.

        Parameters
        ----------
        vocab_size:
            Upper bound on total token count (static + dynamic).
            Currently unused externally but preserved for API compat.
        """
        self.vocab_size: int = vocab_size
        self.vocab: Dict[str, int] = self._build_vocab()

        # Reverse lookup: int → str (static vocab only)
        self._id_to_token: Dict[int, str] = {v: k for k, v in self.vocab.items()}

        # Dynamic value registry (handles strings, ints, floats, bools, None)
        self.value_registry: Dict[Any, int] = {}   # value → id
        self._id_to_value: Dict[int, Any] = {}     # id → value
        self._next_value_id: int = self.VALUE_ID_MIN

    # ------------------------------------------------------------------
    #  Vocabulary construction
    # ------------------------------------------------------------------

    @staticmethod
    def _build_vocab() -> Dict[str, int]:
        """Return the static base vocabulary dict."""
        return {
            # special
            "<PAD>": 0,
            "<START>": 1,
            "<END>": 2,
            "<UNK>": 3,
            # math operations  (4-15)
            "ADD": 4,
            "SUB": 5,
            "MUL": 6,
            "DIV": 7,
            "MOD": 8,
            "POW": 9,
            "EQ": 10,
            "NEQ": 11,
            "LT": 12,
            "LTE": 13,
            "GT": 14,
            "GTE": 15,
            # type / value tokens  (20-29)
            "function": 20,
            "return": 21,
            "variable": 22,
            "constant": 23,
            "list": 24,
            "dict": 25,
            "string": 26,
            "number": 27,
            "boolean": 28,
            "comp": 29,
            # control flow  (30-35)
            "if": 30,
            "else": 31,
            "for": 32,
            "while": 33,
            "loop": 34,
            # common structural keys  (40-43)
            "input": 40,
            "output": 41,
            "condition": 42,
            "body": 43,
            # structural keys — AST node parts  (56-100)
            "binary_op": 56,
            "comparison": 57,
            "bool_op": 58,
            "unary_op": 59,
            "if_exp": 60,
            "call": 61,
            "attribute": 62,
            "subscript": 63,
            "slice": 64,
            "starred": 65,
            "elements": 66,
            "pairs": 67,
            "key": 68,
            "generators": 69,
            "elt": 70,
            "target": 71,
            "iter": 72,
            "ifs": 73,
            "assign": 74,
            "augassign": 75,
            "annassign": 76,
            "annotation": 77,
            "targets": 78,
            "test": 79,
            "orelse": 80,
            "finalbody": 81,
            "handlers": 82,
            "exc": 83,
            "module": 84,
            "class": 85,
            "bases": 86,
            "import": 87,
            "import_from": 88,
            "names": 89,
            "args": 90,
            "keywords": 91,
            "op": 92,
            "left": 93,
            "right": 94,
            "comparator": 95,
            "operand": 96,
            "attr": 97,
            "upper": 98,
            "lower": 99,
            "step": 100,
            # extra node-type identifiers  (101-129)
            "tuple": 101,
            "set": 102,
            "list_comprehension": 103,
            "lambda": 104,
            "yield": 105,
            "break": 106,
            "continue": 107,
            "pass": 108,
            "raise": 109,
            "try": 110,
            "except": 111,
            "with": 112,
            "withitem": 113,
            "assert": 114,
            "delete": 115,
            "fstring": 116,
            "global": 117,
            "nonlocal": 118,
            "unknown_node": 119,
            # extra dict keys commonly found in math specs
            "type": 120,
            "name": 121,
            "value": 122,
            "returns": 123,
            "NoneType": 124,
        }

    # ------------------------------------------------------------------
    #  Registry helpers
    # ------------------------------------------------------------------

    def _register_value(self, value: Any) -> int:
        """Register a dynamic value and return its numeric ID.

        If *value* already exists in ``self.value_registry`` the
        existing ID is returned.
        """
        if value in self.value_registry:
            return self.value_registry[value]

        if self._next_value_id > self.VALUE_ID_MAX:
            raise RuntimeError(
                f"Value registry exhausted ({self.VALUE_ID_MAX=}). "
                "Increase VALUE_ID_MAX or reduce registry size."
            )

        assigned = self._next_value_id
        self.value_registry[value] = assigned
        self._id_to_value[assigned] = value
        self._next_value_id += 1
        return assigned

    def _lookup_value(self, token_id: int) -> Any:
        """Look up a dynamic value from its numeric ID.

        Raises ``KeyError`` if *token_id* is not registered.
        """
        return self._id_to_value[token_id]

    # ------------------------------------------------------------------
    #  Encoding
    # ------------------------------------------------------------------

    def encode(self, math_spec: Dict[str, Any]) -> List[int]:
        """Encode a math-spec dict to a flat token sequence.

        Parameters
        ----------
        math_spec:
            A nested dict produced by ``CodeToMathConverter.convert()``.

        Returns
        -------
        List[int]
            Flat list of integer token IDs.
        """
        tokens: List[int] = [self.START]
        tokens.extend(self._flatten(math_spec))
        tokens.append(self.END)
        return tokens

    def _flatten(self, obj: Any) -> List[int]:
        """Recursively flatten *obj* into token IDs."""
        if isinstance(obj, dict):
            return self._flatten_dict(obj)
        if isinstance(obj, list):
            return self._flatten_list(obj)
        if isinstance(obj, bool):
            return [self.VALUE_BOOL, self._register_value(obj)]
        if isinstance(obj, int):
            return [self.VALUE_INT, self._register_value(obj)]
        if isinstance(obj, float):
            return [self.VALUE_FLOAT, self._register_value(obj)]
        if isinstance(obj, str):
            return self._flatten_string(obj)
        if obj is None:
            return [self.VALUE_STR, self._register_value("__NONE__")]
        # Fallback: treat as unknown string
        return [self.UNK, self.VALUE_STR, self._register_value(str(obj))]

    def _flatten_dict(self, dct: Dict[str, Any]) -> List[int]:
        """Flatten a dict: DICT_START, {key, KEY_SEP, value}*, DICT_END."""
        tokens: List[int] = [self.DICT_START]
        for key, val in dct.items():
            # encode the key
            if key in self.vocab:
                tokens.append(self.vocab[key])
            else:
                tokens.append(self.UNK)
                tokens.append(self._register_value(key))
            tokens.append(self.KEY_SEP)
            tokens.extend(self._flatten(val))
        tokens.append(self.DICT_END)
        return tokens

    def _flatten_list(self, lst: List[Any]) -> List[int]:
        """Flatten a list: LIST_START, items*, LIST_END."""
        tokens: List[int] = [self.LIST_START]
        for item in lst:
            tokens.extend(self._flatten(item))
        tokens.append(self.LIST_END)
        return tokens

    def _flatten_string(self, value: str) -> List[int]:
        """Encode a string value.

        If the string is a known vocabulary word (and not an
        integer/float literal coincidentally matching a key) encode
        it as a plain vocab token.  Otherwise use VALUE_STR +
        registry lookup.
        """
        if value in self.vocab:
            return [self.vocab[value]]
        return [self.VALUE_STR, self._register_value(value)]

    # ------------------------------------------------------------------
    #  Decoding
    # ------------------------------------------------------------------

    def decode(self, tokens: List[int]) -> Dict[str, Any]:
        """Decode a token sequence back to a math-spec dict.

        The reconstructed dict is structurally identical to the
        original (modulo Python float/int distinction for small
        whole-number floats).

        Parameters
        ----------
        tokens:
            Flat token list produced by :meth:`encode`.

        Returns
        -------
        Dict[str, Any]
            Reconstructed math spec.
        """
        if not tokens:
            raise ValueError("Empty token sequence")

        pos = 0
        if tokens[pos] == self.START:
            pos += 1

        result, pos = self._parse(tokens, pos)

        # optionally consume trailing END
        if pos < len(tokens) and tokens[pos] == self.END:
            pos += 1

        # result must be a dict (math specs are always dicts)
        if not isinstance(result, dict):
            raise TypeError(
                f"Decoded root value is {type(result).__name__}, expected dict"
            )
        return result

    def _parse(self, tokens: List[int], pos: int) -> Tuple[Any, int]:
        """Parse a single value starting at *pos*.

        Returns ``(value, next_pos)``.
        """
        if pos >= len(tokens):
            raise ValueError("Unexpected end of token sequence")

        token = tokens[pos]

        # --- structural containers ---
        if token == self.DICT_START:
            return self._parse_dict(tokens, pos + 1)

        if token == self.LIST_START:
            return self._parse_list(tokens, pos + 1)

        # --- value markers ---
        if token == self.VALUE_STR:
            return self._parse_registry_value(tokens, pos + 1)

        if token == self.VALUE_INT:
            return self._parse_registry_value(tokens, pos + 1)

        if token == self.VALUE_FLOAT:
            return self._parse_registry_value(tokens, pos + 1)

        if token == self.VALUE_BOOL:
            return self._parse_registry_value(tokens, pos + 1)

        # --- UNK (unknown key or value) ---
        if token == self.UNK:
            # next token is a registry ID for the string
            dynamic_val, pos2 = self._parse_registry_value(tokens, pos + 1)
            return dynamic_val, pos2

        # --- plain vocab token (string value) ---
        if token in self._id_to_token:
            return self._id_to_token[token], pos + 1

        # --- dynamic registry ID (should not normally appear bare) ---
        if token in self._id_to_value:
            return self._id_to_value[token], pos + 1

        raise ValueError(
            f"Unknown token {token} at position {pos}. "
            "Corrupted sequence or missing vocab entry."
        )

    def _parse_dict(self, tokens: List[int], pos: int) -> Tuple[Dict[str, Any], int]:
        """Parse a dict body (between DICT_START and DICT_END)."""
        result: Dict[str, Any] = {}
        while pos < len(tokens) and tokens[pos] != self.DICT_END:
            # --- read key ---
            if tokens[pos] == self.UNK:
                # unknown key → next token is registry ID
                key_val, pos = self._parse(tokens, pos)
                # key_val is expected to be a string
                key = str(key_val)
            else:
                key_raw, pos = self._parse(tokens, pos)
                key = str(key_raw)

            # expect KEY_SEP
            if pos >= len(tokens) or tokens[pos] != self.KEY_SEP:
                raise ValueError(
                    f"Expected KEY_SEP ({self.KEY_SEP}) at position {pos}, "
                    f"got {tokens[pos] if pos < len(tokens) else 'EOF'}"
                )
            pos += 1  # consume KEY_SEP

            # --- read value ---
            val, pos = self._parse(tokens, pos)
            result[key] = val

        if pos >= len(tokens) or tokens[pos] != self.DICT_END:
            raise ValueError(
                f"Expected DICT_END ({self.DICT_END}) at position {pos}"
            )
        pos += 1  # consume DICT_END
        return result, pos

    def _parse_list(self, tokens: List[int], pos: int) -> Tuple[List[Any], int]:
        """Parse a list body (between LIST_START and LIST_END)."""
        result: List[Any] = []
        while pos < len(tokens) and tokens[pos] != self.LIST_END:
            item, pos = self._parse(tokens, pos)
            result.append(item)
        if pos >= len(tokens) or tokens[pos] != self.LIST_END:
            raise ValueError(
                f"Expected LIST_END ({self.LIST_END}) at position {pos}"
            )
        pos += 1  # consume LIST_END
        return result, pos

    def _parse_registry_value(self, tokens: List[int], pos: int) -> Tuple[Any, int]:
        """Read one registry-ID token and return the stored value."""
        if pos >= len(tokens):
            raise ValueError("Expected registry ID but reached end of tokens")
        token = tokens[pos]
        if token in self._id_to_value:
            raw = self._id_to_value[token]
            # Convert sentinel back to Python None.
            if raw == "__NONE__":
                return None, pos + 1
            return raw, pos + 1
        raise ValueError(f"No registered value for token {token} at position {pos}")

    # ------------------------------------------------------------------
    #  Utility methods
    # ------------------------------------------------------------------

    def encode_batch(self, specs: List[Dict[str, Any]]) -> List[List[int]]:
        """Encode a batch of math specs.

        Parameters
        ----------
        specs:
            List of math-spec dicts.

        Returns
        -------
        List[List[int]]
            One token list per spec.
        """
        return [self.encode(spec) for spec in specs]

    def pad_sequences(
        self,
        sequences: List[List[int]],
        max_len: int,
    ) -> List[List[int]]:
        """Pad or truncate every sequence to *max_len* with PAD.

        Returns a plain ``List[List[int]]`` so the module stays
        dependency-free.  Callers that need ``torch.Tensor`` can
        wrap the result with ``torch.tensor(padded)``.

        Parameters
        ----------
        sequences:
            Variable-length token sequences.
        max_len:
            Target length.  Longer sequences are **truncated** from
            the right; shorter ones are padded with ``PAD`` (0).

        Returns
        -------
        List[List[int]]
            Uniform-length token lists.
        """
        padded: List[List[int]] = []
        for seq in sequences:
            if len(seq) > max_len:
                padded.append(seq[:max_len])
            else:
                padded.append(seq + [self.PAD] * (max_len - len(seq)))
        return padded

    def get_vocab_size(self) -> int:
        """Return the effective vocabulary size.

        This is ``len(self.vocab) + len(self.value_registry)``, i.e.
        every token ID that has been assigned so far.
        """
        return len(self.vocab) + len(self.value_registry)

    def decode_to_string(self, tokens: List[int]) -> str:
        """Decode a token sequence into a human-readable string.

        Useful for debugging and inspection during training.

        Parameters
        ----------
        tokens:
            Token list from :meth:`encode`.

        Returns
        -------
        str
            Space-separated token names.
        """
        parts: List[str] = []
        for t in tokens:
            if t in self._id_to_token:
                parts.append(self._id_to_token[t])
            elif t in self._id_to_value:
                raw = self._id_to_value[t]
                parts.append(repr(raw))
            elif t == self.PAD:
                parts.append("<PAD>")
            elif t == self.START:
                parts.append("<START>")
            elif t == self.END:
                parts.append("<END>")
            elif t == self.UNK:
                parts.append("<UNK>")
            elif t == self.KEY_SEP:
                parts.append("KEY_SEP")
            elif t == self.LIST_START:
                parts.append("LIST_START")
            elif t == self.LIST_END:
                parts.append("LIST_END")
            elif t == self.DICT_START:
                parts.append("DICT_START")
            elif t == self.DICT_END:
                parts.append("DICT_END")
            elif t == self.VALUE_INT:
                parts.append("VALUE_INT")
            elif t == self.VALUE_FLOAT:
                parts.append("VALUE_FLOAT")
            elif t == self.VALUE_STR:
                parts.append("VALUE_STR")
            elif t == self.VALUE_BOOL:
                parts.append("VALUE_BOOL")
            else:
                parts.append(f"<{t}>")
        return " ".join(parts)

    def save_vocabulary(self, path: str) -> None:
        """Persist vocabulary and registry state to a JSON file.

        The saved file can be reloaded with :meth:`load_vocabulary`
        to restore an identical tokenizer state (useful for
        inference after training).

        Parameters
        ----------
        path:
            Filesystem path for the JSON output.
        """
        # Convert value_registry keys to JSON-safe types.
        json_registry: Dict[str, int] = {}
        for val, vid in self.value_registry.items():
            if val is None:
                key = "__NONE__"
            elif isinstance(val, bool):
                key = f"__BOOL__{val}"
            elif isinstance(val, (int, float)):
                key = f"__NUM__{val}"
            elif isinstance(val, str):
                key = val
            else:
                key = str(val)
            json_registry[key] = vid

        payload = {
            "base_vocab": self.vocab,
            "value_registry": json_registry,
            "next_value_id": self._next_value_id,
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2, ensure_ascii=False)

    def load_vocabulary(self, path: str) -> None:
        """Restore vocabulary and registry state from a JSON file.

        Parameters
        ----------
        path:
            Path to a file previously written by :meth:`save_vocabulary`.
        """
        with open(path, encoding="utf-8") as f:
            payload = json.load(f)

        # Restore base vocab (overwrite any existing)
        self.vocab = payload["base_vocab"]
        self._id_to_token = {v: k for k, v in self.vocab.items()}

        # Rebuild value registries from JSON
        self.value_registry = {}
        self._id_to_value = {}
        for key_str, vid in payload["value_registry"].items():
            # Reverse the key encoding from save_vocabulary
            if key_str == "__NONE__":
                val: Any = None
            elif key_str.startswith("__BOOL__"):
                val = key_str[len("__BOOL__"):] == "True"
            elif key_str.startswith("__NUM__"):
                num_str = key_str[len("__NUM__"):]
                try:
                    val = int(num_str)
                except ValueError:
                    val = float(num_str)
            else:
                val = key_str
            self.value_registry[val] = vid
            self._id_to_value[vid] = val

        self._next_value_id = payload.get("next_value_id", self.VALUE_ID_MIN)

    # ------------------------------------------------------------------
    #  Name registry (convenience alias — kept for backward compat)
    # ------------------------------------------------------------------

    @property
    def name_registry(self) -> Dict[str, int]:
        """Convenience accessor for string-only entries in the value registry."""
        return {
            str(k): v
            for k, v in self.value_registry.items()
            if isinstance(k, str) and k != "__NONE__"
        }