"""Inference-only Morton and spatial BPE codec."""
from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import json
import re
import numpy as np

_DIR_TO_VEC = {0: (1, 0, 0), 1: (0, 1, 0), 2: (0, 0, 1)}
MORTON_GRID_SIZE = 32
MORTON_BITS = 5
MORTON_MAX_CODE = (1 << 15) - 1
MORTON_TOKEN_RE = re.compile(r"<morton_(\d+)>")
MESH_TOKEN_RE = re.compile(r"<mesh_(\d+)>")
MORTON_MESH_PAIR_RE = re.compile(r"<morton_(\d+)>\s*<mesh_(\d+)>")
PairKey = Tuple[int, int, int, int, int]

@dataclass
class MergeEntry:
    new_id: int
    pair: Tuple[int, int]
    direction: int
    rel_offset: Tuple[int, int, int]
    leaf_tokens: List[int]
    leaf_offsets: List[Tuple[int, int, int]]

@dataclass
class MortonMeshParseResult:
    ids: np.ndarray
    anchors: np.ndarray
    dropped_count: int = 0

def _as_int64_array(x: Any, name: str, ndim: int) -> np.ndarray:
    arr = np.asarray(x, dtype=np.int64)
    if arr.ndim != ndim:
        raise ValueError(f'{name} must be {ndim}D, got shape={arr.shape}')
    return arr

def morton3d_encode_xyz32(coords: Any) -> Any:
    """Encode block coordinates ``[0, 31]^3`` with fixed x/y/z bit interleave.

    Bit order is fixed for the whole project:
    ``x_i -> bit 3*i``, ``y_i -> bit 3*i+1``, ``z_i -> bit 3*i+2``.
    """
    arr = np.asarray(coords, dtype=np.int64)
    single = arr.ndim == 1
    if single:
        if arr.shape[0] != 3:
            raise ValueError(f'coords must be [3] or [N,3], got shape={arr.shape}')
        arr2 = arr.reshape(1, 3)
    elif arr.ndim == 2 and arr.shape[1] == 3:
        arr2 = arr
    else:
        raise ValueError(f'coords must be [3] or [N,3], got shape={arr.shape}')
    if arr2.size and ((arr2 < 0).any() or (arr2 >= MORTON_GRID_SIZE).any()):
        raise ValueError(f'Morton coordinates must be in [0, {MORTON_GRID_SIZE - 1}], got min={arr2.min()} max={arr2.max()}')
    x = arr2[:, 0]
    y = arr2[:, 1]
    z = arr2[:, 2]
    code = np.zeros(arr2.shape[0], dtype=np.int64)
    for bit in range(MORTON_BITS):
        code |= (x >> bit & 1) << 3 * bit
        code |= (y >> bit & 1) << 3 * bit + 1
        code |= (z >> bit & 1) << 3 * bit + 2
    return int(code[0]) if single else code

def morton3d_decode_xyz32(codes: Any) -> Any:
    """Decode fixed 15-bit Morton code(s) back to block coordinates."""
    arr = np.asarray(codes, dtype=np.int64)
    single = arr.ndim == 0
    arr1 = arr.reshape(-1)
    if arr1.size and ((arr1 < 0).any() or (arr1 > MORTON_MAX_CODE).any()):
        raise ValueError(f'Morton codes must be in [0, {MORTON_MAX_CODE}], got min={arr1.min()} max={arr1.max()}')
    coords = np.zeros((arr1.shape[0], 3), dtype=np.int64)
    for bit in range(MORTON_BITS):
        coords[:, 0] |= (arr1 >> 3 * bit & 1) << bit
        coords[:, 1] |= (arr1 >> 3 * bit + 1 & 1) << bit
        coords[:, 2] |= (arr1 >> 3 * bit + 2 & 1) << bit
    return coords[0] if single else coords

def serialize_morton_mesh_pairs(ids: Sequence[int], anchors: Any) -> str:
    """Serialize BPE/VQ ids with explicit block-anchor Morton tokens."""
    ids_np = np.asarray(ids, dtype=np.int64).reshape(-1)
    anchors_np = _as_int64_array(anchors, 'anchors', 2)
    if anchors_np.shape[1] != 3:
        raise ValueError(f'anchors must be [K,3], got {anchors_np.shape}')
    if ids_np.shape[0] != anchors_np.shape[0]:
        raise ValueError('ids and anchors length mismatch')
    if ids_np.size == 0:
        return '<mesh_start><mesh_end>'
    morton_codes = morton3d_encode_xyz32(anchors_np)
    parts = [f'<morton_{int(m)}><mesh_{int(t)}>' for (m, t) in zip(morton_codes.tolist(), ids_np.tolist())]
    return '<mesh_start>' + ''.join(parts) + '<mesh_end>'

def parse_morton_mesh_pairs(text: str, max_mesh_id: Optional[int]=None) -> MortonMeshParseResult:
    """Parse valid ``<morton_*><mesh_*>`` pairs, dropping damaged fragments.

    Only adjacent Morton/Mesh pairs are accepted. Lone Morton tokens, lone mesh
    tokens, out-of-range Morton codes, and mesh ids >= ``max_mesh_id`` are
    skipped and counted in ``dropped_count``.
    """
    ids: List[int] = []
    morton_codes: List[int] = []
    used_morton_spans = set()
    used_mesh_spans = set()
    dropped = 0
    for match in MORTON_MESH_PAIR_RE.finditer(text):
        morton = int(match.group(1))
        mesh_id = int(match.group(2))
        used_morton_spans.add((match.start(1), match.end(1)))
        used_mesh_spans.add((match.start(2), match.end(2)))
        if morton < 0 or morton > MORTON_MAX_CODE:
            dropped += 1
            continue
        if max_mesh_id is not None and mesh_id >= int(max_mesh_id):
            dropped += 1
            continue
        morton_codes.append(morton)
        ids.append(mesh_id)
    total_morton = len(list(MORTON_TOKEN_RE.finditer(text)))
    total_mesh = len(list(MESH_TOKEN_RE.finditer(text)))
    dropped += max(0, total_morton - len(used_morton_spans))
    dropped += max(0, total_mesh - len(used_mesh_spans))
    if ids:
        anchors = morton3d_decode_xyz32(np.asarray(morton_codes, dtype=np.int64))
        ids_np = np.asarray(ids, dtype=np.int64)
    else:
        anchors = np.zeros((0, 3), dtype=np.int64)
        ids_np = np.zeros((0,), dtype=np.int64)
    return MortonMeshParseResult(ids=ids_np, anchors=anchors, dropped_count=dropped)

def _morton_order(coords: np.ndarray) -> np.ndarray:
    if coords.size == 0:
        return np.array([], dtype=np.int64)
    cmax = int(np.max(coords))
    return _morton_sort_indices(coords.astype(np.int64), coord_max=max(cmax + 1, 1))

def _morton_sort_indices(coords_xyz: np.ndarray, coord_max: int) -> np.ndarray:
    n = coords_xyz.shape[0]
    if n == 0:
        return np.array([], dtype=np.int64)
    bits = max(1, int(np.ceil(np.log2(coord_max + 1))))
    x = coords_xyz[:, 0].astype(np.int64)
    y = coords_xyz[:, 1].astype(np.int64)
    z = coords_xyz[:, 2].astype(np.int64)
    bit_idx = np.arange(bits, dtype=np.int64)
    x_bits = (x[:, None] >> bit_idx & 1) << 3 * bit_idx
    y_bits = (y[:, None] >> bit_idx & 1) << 3 * bit_idx + 1
    z_bits = (z[:, None] >> bit_idx & 1) << 3 * bit_idx + 2
    code = x_bits.sum(axis=1) + y_bits.sum(axis=1) + z_bits.sum(axis=1)
    return np.argsort(code, kind='stable')

def _canonical_min_offset(offsets: Sequence[Tuple[int, int, int]]) -> Tuple[int, int, int]:
    return min(((int(c[0]), int(c[1]), int(c[2])) for c in offsets))

def _add3(a: Tuple[int, int, int], b: Tuple[int, int, int]) -> Tuple[int, int, int]:
    return (a[0] + b[0], a[1] + b[1], a[2] + b[2])

def _sub3(a: Tuple[int, int, int], b: Tuple[int, int, int]) -> Tuple[int, int, int]:
    return (a[0] - b[0], a[1] - b[1], a[2] - b[2])

def _make_pair_key(a_tok: int, b_tok: int, a_anchor: Sequence[int], b_anchor: Sequence[int]) -> PairKey:
    return (int(a_tok), int(b_tok), int(b_anchor[0]) - int(a_anchor[0]), int(b_anchor[1]) - int(a_anchor[1]), int(b_anchor[2]) - int(a_anchor[2]))

class _IncEncodeState:
    """Per-sample slot graph for fast BPE encode (replays ``merge_table`` in order)."""
    __slots__ = ('tok', 'M', 'cell_xyz', 'directions', 'coord_to_idx', 'slot_of_cell', 'token_of_slot', 'slot_anchor', 'cells_by_slot', 'pair_slot_pairs', 'incidence')

    def __init__(self, tok: 'BPE3DTokenizer', tokens: np.ndarray, coords: np.ndarray) -> None:
        self.tok = tok
        self.M = int(tokens.shape[0])
        self.cell_xyz = coords.astype(np.int64, copy=False)
        self.directions = tok.directions
        self.coord_to_idx: Dict[Tuple[int, int, int], int] = {}
        for i in range(self.M):
            t = (int(self.cell_xyz[i, 0]), int(self.cell_xyz[i, 1]), int(self.cell_xyz[i, 2]))
            self.coord_to_idx[t] = i
        self.slot_of_cell = np.arange(self.M, dtype=np.int64)
        self.token_of_slot = tokens.astype(np.int64, copy=False).copy()
        self.slot_anchor = self.cell_xyz.copy()
        self.cells_by_slot: Dict[int, np.ndarray] = {i: np.array([i], dtype=np.int64) for i in range(self.M)}
        self.pair_slot_pairs: Dict[PairKey, set] = {}
        self.incidence: Dict[int, Dict[PairKey, Dict[int, int]]] = defaultdict(dict)
        if self.M > 0:
            cmax = int(np.max(self.cell_xyz))
            K = int(cmax) + 2
            xyz = self.cell_xyz
            codes = ((xyz[:, 0] * K + xyz[:, 1]) * K + xyz[:, 2]).astype(np.int64)
            order = np.argsort(codes, kind='stable')
            sorted_codes = codes[order]
            tokv = self.token_of_slot
            for d in self.directions:
                dv = np.asarray(_DIR_TO_VEC[int(d)], dtype=np.int64)
                nbr = xyz + dv
                nc = ((nbr[:, 0] * K + nbr[:, 1]) * K + nbr[:, 2]).astype(np.int64)
                idx = np.searchsorted(sorted_codes, nc)
                hit = (idx < self.M) & (sorted_codes[np.minimum(idx, self.M - 1)] == nc)
                a_cell = np.nonzero(hit)[0]
                if a_cell.size == 0:
                    continue
                b_cell = order[idx[hit]]
                mask = b_cell != a_cell
                a_cell = a_cell[mask]
                b_cell = b_cell[mask]
                for k in range(int(a_cell.shape[0])):
                    i = int(a_cell[k])
                    j = int(b_cell[k])
                    self.register_directed(_make_pair_key(int(tokv[i]), int(tokv[j]), self.slot_anchor[i], self.slot_anchor[j]), int(i), int(j))

    def register_directed(self, key: PairKey, a_slot: int, b_slot: int) -> None:
        bucket = self.pair_slot_pairs.setdefault(key, set())
        if (a_slot, b_slot) in bucket:
            return
        bucket.add((a_slot, b_slot))
        self.incidence[a_slot].setdefault(key, {})[b_slot] = 0
        self.incidence[b_slot].setdefault(key, {})[a_slot] = 1

    def remove_directed(self, key: PairKey, a_slot: int, b_slot: int) -> None:
        bucket = self.pair_slot_pairs.get(key)
        if bucket is None or (a_slot, b_slot) not in bucket:
            return
        bucket.discard((a_slot, b_slot))
        if not bucket:
            del self.pair_slot_pairs[key]
        if key in self.incidence[a_slot]:
            self.incidence[a_slot][key].pop(b_slot, None)
            if not self.incidence[a_slot][key]:
                del self.incidence[a_slot][key]
        if key in self.incidence[b_slot]:
            self.incidence[b_slot][key].pop(a_slot, None)
            if not self.incidence[b_slot][key]:
                del self.incidence[b_slot][key]

    def remove_all_incident(self, slot: int) -> None:
        for (key, others) in list(self.incidence[slot].items()):
            for (other, role) in list(others.items()):
                if int(role) == 0:
                    self.remove_directed(key, int(slot), int(other))
                else:
                    self.remove_directed(key, int(other), int(slot))

    def merge_slots(self, sa: int, sb: int, new_id: int) -> None:
        if sa == sb:
            return
        ca = self.cells_by_slot[sa]
        cb = self.cells_by_slot[sb]
        merged = np.concatenate([ca, cb], axis=0)
        self.cells_by_slot[sa] = merged
        del self.cells_by_slot[sb]
        for idx in cb.tolist():
            self.slot_of_cell[int(idx)] = sa
        self.token_of_slot[sa] = int(new_id)
        self.token_of_slot[sb] = -1
        if tuple((int(v) for v in self.slot_anchor[sb].tolist())) < tuple((int(v) for v in self.slot_anchor[sa].tolist())):
            self.slot_anchor[sa] = self.slot_anchor[sb]

    def add_edges_for_slot(self, slot: int) -> None:
        if int(self.token_of_slot[slot]) < 0:
            return
        ta = int(self.token_of_slot[slot])
        cells = self.cells_by_slot[slot]
        for ci in cells.tolist():
            c = (int(self.cell_xyz[int(ci), 0]), int(self.cell_xyz[int(ci), 1]), int(self.cell_xyz[int(ci), 2]))
            for d in self.directions:
                dv = _DIR_TO_VEC[int(d)]
                n = (c[0] + dv[0], c[1] + dv[1], c[2] + dv[2])
                j = self.coord_to_idx.get(n)
                if j is not None:
                    sj = int(self.slot_of_cell[int(j)])
                    if sj != int(slot):
                        tb = int(self.token_of_slot[sj])
                        if tb >= 0:
                            self.register_directed(_make_pair_key(ta, tb, self.slot_anchor[slot], self.slot_anchor[sj]), int(slot), sj)
                n_sub = (c[0] - dv[0], c[1] - dv[1], c[2] - dv[2])
                k = self.coord_to_idx.get(n_sub)
                if k is not None:
                    sk = int(self.slot_of_cell[int(k)])
                    if sk != int(slot):
                        t_left = int(self.token_of_slot[sk])
                        if t_left >= 0:
                            self.register_directed(_make_pair_key(t_left, ta, self.slot_anchor[sk], self.slot_anchor[slot]), sk, int(slot))

    def apply_merge_table(self) -> None:
        for entry in self.tok.merge_table:
            (a_tok, b_tok) = (int(entry.pair[0]), int(entry.pair[1]))
            rel_offset = tuple((int(v) for v in entry.rel_offset))
            new_id = int(entry.new_id)
            key = (a_tok, b_tok, rel_offset[0], rel_offset[1], rel_offset[2])
            if key not in self.pair_slot_pairs:
                continue
            snapshot = list(self.pair_slot_pairs[key])
            pairs = [(int(x), int(y)) for (x, y) in snapshot if int(self.token_of_slot[int(x)]) == a_tok and int(self.token_of_slot[int(y)]) == b_tok]
            if not pairs:
                continue
            if len(pairs) > 1:
                anchors = np.stack([self.slot_anchor[pa] for (pa, _) in pairs])
                order = _morton_order(anchors)
                pairs = [pairs[int(i)] for i in order.tolist()]
            used: set = set()
            for (sa, sb) in pairs:
                if sa in used or sb in used:
                    continue
                if int(self.token_of_slot[int(sa)]) < 0 or int(self.token_of_slot[int(sb)]) < 0:
                    continue
                if int(self.token_of_slot[int(sa)]) != a_tok or int(self.token_of_slot[int(sb)]) != b_tok:
                    continue
                self.remove_all_incident(int(sa))
                self.remove_all_incident(int(sb))
                self.merge_slots(int(sa), int(sb), new_id)
                self.add_edges_for_slot(int(sa))
                used.add(sa)
                used.add(sb)

    def final_ids_anchors(self) -> Tuple[np.ndarray, np.ndarray]:
        if self.M == 0:
            return (np.zeros((0,), dtype=np.int64), np.zeros((0, 3), dtype=np.int64))
        alive_slots = [s for s in self.cells_by_slot if int(self.token_of_slot[s]) >= 0]
        if not alive_slots:
            return (np.zeros((0,), dtype=np.int64), np.zeros((0, 3), dtype=np.int64))
        ids = np.asarray([int(self.token_of_slot[s]) for s in alive_slots], dtype=np.int64)
        anchors = np.stack([self.slot_anchor[s] for s in alive_slots]).astype(np.int64, copy=False)
        order = _morton_order(anchors)
        return (ids[order], anchors[order])

class BPE3DTokenizer:

    def __init__(self, base_vocab_size: int=8192, directions: Tuple[int, ...]=(0, 1, 2)) -> None:
        self.base_vocab_size = int(base_vocab_size)
        self.directions = tuple((int(d) for d in directions))
        for d in self.directions:
            if d not in _DIR_TO_VEC:
                raise ValueError(f'unsupported direction={d}, expected one of {tuple(_DIR_TO_VEC.keys())}')
        self.merge_table: List[MergeEntry] = []

    @property
    def vocab_size(self) -> int:
        return self.base_vocab_size + len(self.merge_table)

    def _build_flat_leaf_templates(self) -> None:
        flat: Dict[int, Tuple[List[int], List[Tuple[int, int, int]]]] = {i: ([i], [(0, 0, 0)]) for i in range(self.base_vocab_size)}
        for entry in self.merge_table:
            (a, b) = entry.pair
            shift = tuple((int(v) for v in entry.rel_offset))
            (a_leaf_ids, a_leaf_off) = flat[a]
            (b_leaf_ids, b_leaf_off) = flat[b]
            leaf_ids = list(a_leaf_ids)
            leaf_offs = list(a_leaf_off)
            for (t, off) in zip(b_leaf_ids, b_leaf_off):
                leaf_ids.append(int(t))
                leaf_offs.append(_add3(off, shift))
            anchor_shift = _canonical_min_offset(leaf_offs)
            leaf_offs = [_sub3(off, anchor_shift) for off in leaf_offs]
            entry.leaf_tokens = leaf_ids
            entry.leaf_offsets = leaf_offs
            flat[entry.new_id] = (leaf_ids, leaf_offs)

    def _encode_single_incremental(self, tokens: np.ndarray, coords: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Replay ``merge_table`` with local edge updates (fast encode path)."""
        tokens = _as_int64_array(tokens, 'tokens', 1)
        coords = _as_int64_array(coords, 'coords', 2)
        if tokens.shape[0] == 0:
            return (np.zeros((0,), dtype=np.int64), np.zeros((0, 3), dtype=np.int64))
        st = _IncEncodeState(self, tokens, coords)
        st.apply_merge_table()
        return st.final_ids_anchors()

    def _encode_single(self, tokens: np.ndarray, coords: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        return self._encode_single_incremental(tokens, coords)

    def encode_sparse_numpy(self, tokens: np.ndarray, coords: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """Encode one mesh sample on CPU only (no per-point ``torch.tensor``).

        Returns ``(macro_ids [K], macro_anchors [K,3])`` in Morton order of anchors.
        """
        return self._encode_single(tokens, coords)

    def decode_to_leaf_batches(self, batches: Sequence[Dict[str, np.ndarray]], *, coordinate_min: int=0, coordinate_max: int=MORTON_GRID_SIZE - 1, duplicate_policy: str='first') -> List[Dict[str, Any]]:
        """Expand BPE macros and sanitize leaf VQ tokens before sparse decode.

        ``duplicate_policy='first'`` is deliberately deterministic: the first
        generated leaf at a coordinate is retained and later leaves at that
        coordinate are counted and discarded. Out-of-range coordinates and VQ
        ids never reach ``SparseTensor``.
        """
        if duplicate_policy != 'first':
            raise ValueError(f"unsupported duplicate_policy={duplicate_policy!r}; expected 'first'")
        lo = int(coordinate_min)
        hi = int(coordinate_max)
        if lo > hi:
            raise ValueError(f'invalid coordinate range [{lo}, {hi}]')
        outputs: List[Dict[str, Any]] = []
        for rec in batches:
            ids = _as_int64_array(rec['ids'], 'ids', 1)
            anchors = _as_int64_array(rec['anchors'], 'anchors', 2)
            if anchors.shape[1] != 3:
                raise ValueError(f'anchors must be [K,3], got {anchors.shape}')
            if ids.shape[0] != anchors.shape[0]:
                raise ValueError('ids and anchors length mismatch')
            leaf_ids: List[int] = []
            leaf_coords: List[Tuple[int, int, int]] = []
            seen_coords: set[Tuple[int, int, int]] = set()
            seen_anchors: set[Tuple[int, int, int]] = set()
            stats = {'macro_pair_count': int(ids.shape[0]), 'duplicate_macro_anchors': 0, 'invalid_bpe_ids': 0, 'expanded_leaf_token_count': 0, 'invalid_vq_ids': 0, 'out_of_range_leaf_coords': 0, 'duplicate_leaf_coords': 0, 'valid_leaf_token_count': 0, 'duplicate_policy': duplicate_policy, 'coordinate_min': lo, 'coordinate_max': hi}
            for (tid, anc_raw) in zip(ids.tolist(), anchors.tolist()):
                aid = int(tid)
                anchor = tuple((int(v) for v in anc_raw))
                if anchor in seen_anchors:
                    stats['duplicate_macro_anchors'] += 1
                else:
                    seen_anchors.add(anchor)
                if aid < 0 or aid >= self.vocab_size:
                    stats['invalid_bpe_ids'] += 1
                    continue
                if aid < self.base_vocab_size:
                    expanded = [(aid, (0, 0, 0))]
                else:
                    entry = self.merge_table[aid - self.base_vocab_size]
                    expanded = list(zip(entry.leaf_tokens, entry.leaf_offsets))
                for (leaf_id_raw, offset_raw) in expanded:
                    stats['expanded_leaf_token_count'] += 1
                    leaf_id = int(leaf_id_raw)
                    if leaf_id < 0 or leaf_id >= self.base_vocab_size:
                        stats['invalid_vq_ids'] += 1
                        continue
                    coord = _add3(anchor, tuple((int(v) for v in offset_raw)))
                    if any((v < lo or v > hi for v in coord)):
                        stats['out_of_range_leaf_coords'] += 1
                        continue
                    if coord in seen_coords:
                        stats['duplicate_leaf_coords'] += 1
                        continue
                    seen_coords.add(coord)
                    leaf_ids.append(leaf_id)
                    leaf_coords.append(coord)
            stats['valid_leaf_token_count'] = len(leaf_ids)
            outputs.append({'ids': np.asarray(leaf_ids, dtype=np.int64), 'coords': np.asarray(leaf_coords, dtype=np.int64).reshape(-1, 3), 'stats': stats})
        return outputs

    @classmethod
    def load(cls, path: str) -> 'BPE3DTokenizer':
        p = Path(path)
        data = json.loads(p.read_text(encoding='utf-8'))
        tok = cls(base_vocab_size=int(data['base_vocab_size']), directions=tuple((int(x) for x in data.get('directions', [0, 1, 2]))))
        table: List[MergeEntry] = []
        for r in data['merge_table']:
            if 'rel_offset' not in r:
                raise ValueError('This merge_table.json was created by an older non-reversible 3D BPE format without rel_offset. Please retrain the BPE merge table before using decode_to_sparse.')
            table.append(MergeEntry(new_id=int(r['new_id']), pair=(int(r['pair'][0]), int(r['pair'][1])), direction=int(r['direction']), rel_offset=tuple((int(v) for v in r['rel_offset'])), leaf_tokens=[int(x) for x in r.get('leaf_tokens', [])], leaf_offsets=[tuple((int(v) for v in t)) for t in r.get('leaf_offsets', [])]))
        tok.merge_table = table
        if len(tok.merge_table) > 0 and len(tok.merge_table[0].leaf_tokens) == 0:
            tok._build_flat_leaf_templates()
        return tok