"""Morton Encoding - O(1) Z-order Space-Filling Curve

Pure Python implementation with optional numpy acceleration.
"""
from __future__ import annotations

from typing import Tuple, List


def _spread_bits(v: int) -> int:
    """Spread bits for 3D Morton encoding. O(1)."""
    v = v & 0x1FFFFF
    v = (v | (v << 32)) & 0x1F00000000FFFF
    v = (v | (v << 16)) & 0x1F0000FF0000FF
    v = (v | (v << 8)) & 0x100F00F00F00F00F
    v = (v | (v << 4)) & 0x10C30C30C30C30C3
    v = (v | (v << 2)) & 0x1249249249249249
    return v


def _compact_bits(v: int) -> int:
    """Compact bits for Morton decoding. O(1)."""
    v = v & 0x1249249249249249
    v = (v | (v >> 2)) & 0x10C30C30C30C30C3
    v = (v | (v >> 4)) & 0x100F00F00F00F00F
    v = (v | (v >> 8)) & 0x1F0000FF0000FF
    v = (v | (v >> 16)) & 0x1F00000000FFFF
    v = (v | (v >> 32)) & 0x1FFFFF
    return v


def encode(x: int, y: int, z: int = 0) -> int:
    """Encode 3D coordinates to Morton code. O(1), 21 operations."""
    return _spread_bits(x) | (_spread_bits(y) << 1) | (_spread_bits(z) << 2)


def decode(morton: int) -> Tuple[int, int, int]:
    """Decode Morton code to 3D coordinates. O(1)."""
    return _compact_bits(morton), _compact_bits(morton >> 1), _compact_bits(morton >> 2)


def encode_2d(x: int, y: int) -> int:
    """Encode 2D coordinates to Morton code."""
    return encode(x, y, 0)


def decode_2d(morton: int) -> Tuple[int, int]:
    """Decode Morton code to 2D coordinates."""
    x, y, _ = decode(morton)
    return x, y


def batch_encode(coords: List[Tuple[int, int, int]]) -> List[int]:
    """Batch encode coordinates. O(N)."""
    return [encode(x, y, z) for x, y, z in coords]


def batch_decode(mortons: List[int]) -> List[Tuple[int, int, int]]:
    """Batch decode Morton codes. O(N)."""
    return [decode(m) for m in mortons]


def distance(m1: int, m2: int) -> float:
    """Calculate Euclidean distance between two Morton codes."""
    x1, y1, z1 = decode(m1)
    x2, y2, z2 = decode(m2)
    return ((x2 - x1) ** 2 + (y2 - y1) ** 2 + (z2 - z1) ** 2) ** 0.5


def neighbors(morton: int, radius: int = 1) -> List[int]:
    """Get neighboring Morton codes within radius."""
    x, y, z = decode(morton)
    result = []
    for dx in range(-radius, radius + 1):
        for dy in range(-radius, radius + 1):
            for dz in range(-radius, radius + 1):
                if dx == 0 and dy == 0 and dz == 0:
                    continue
                nx, ny, nz = x + dx, y + dy, z + dz
                if nx >= 0 and ny >= 0 and nz >= 0:
                    result.append(encode(nx, ny, nz))
    return result
