"""Holographic Reduced Representation (HRR) primitives on Apple MLX.

Binding/unbinding are circular (de)convolution in the FFT domain (Plate, 1995).
These are the building blocks the morpheme tokenizer uses to compose
``prefix (x) root (x) suffix`` into a single fixed-width "holistic" token vector.
"""
from __future__ import annotations

import mlx.core as mx
import numpy as np


def normalize(x: mx.array, eps: float = 1e-6) -> mx.array:
    """L2-normalize along the last axis; upcast to float32 for stable math."""
    x = x.astype(mx.float32)
    return x / mx.sqrt(mx.sum(x * x, axis=-1, keepdims=True) + eps)


def _bind(a: mx.array, b: mx.array) -> mx.array:
    """Circular convolution: real(IFFT(FFT(a) * FFT(b)))."""
    return mx.real(mx.fft.ifft(mx.fft.fft(a.astype(mx.float32)) * mx.fft.fft(b.astype(mx.float32))))


def _unbind(bound: mx.array, key: mx.array) -> mx.array:
    """Circular correlation (approx inverse of bind): real(IFFT(FFT(bound) * conj(FFT(key))))."""
    return mx.real(
        mx.fft.ifft(mx.fft.fft(bound.astype(mx.float32)) * mx.conj(mx.fft.fft(key.astype(mx.float32))))
    )


# Compiled hot paths.
bind = mx.compile(_bind)
unbind = mx.compile(_unbind)


def bundle(*vectors: mx.array) -> mx.array:
    """Superpose vectors by normalized sum (the VSA "add")."""
    if not vectors:
        raise ValueError("bundle requires at least one vector")
    return normalize(sum(v.astype(mx.float32) for v in vectors))


def cosine_similarity(a: mx.array, b: mx.array) -> mx.array:
    """Cosine similarity along the last axis (vectors are normalized first)."""
    return mx.sum(normalize(a) * normalize(b), axis=-1)


def make_unitary(dim: int, seed: int = 0) -> mx.array:
    """Deterministic unitary vector: its FFT has all-ones magnitude, so binding with
    it is exactly invertible (unbind perfectly recovers the bound value)."""
    rng = np.random.default_rng(seed)
    spectrum = np.ones(dim, dtype=np.complex64)
    half = dim // 2
    phases = rng.uniform(0, 2 * np.pi, max(0, half - 1))
    spectrum[1:half] = np.exp(1j * phases)
    spectrum[half + 1 :] = np.conj(spectrum[1:half][::-1])
    if dim % 2 == 0:
        spectrum[half] = 1.0
    return mx.array(np.fft.ifft(spectrum).real.astype(np.float32))


def update_context(context: mx.array, token: mx.array, position_key: mx.array) -> mx.array:
    """Incrementally fold a token into a running context: normalize(context + token (x) pos)."""
    return normalize(context.astype(mx.float32) + bind(token, position_key))