"""Vector composition helpers."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .ri import dense_l2_normalize
from .types import FeatureMap, WeightedDenseVectors

if TYPE_CHECKING:
    from .encoder import IdentifierEncoder


def compose_vectors(
    vectors_with_weights: WeightedDenseVectors,
    dim: int | None = None,
    normalize: bool = True,
    dtype: str | np.dtype = "float32",
) -> np.ndarray:
    """Compose weighted dense vectors into one output vector.

    This is intended for record-level or multi-signal retrieval, where several
    already-encoded identifiers should be collapsed into a single ANN query.
    """
    if vectors_with_weights:
        first_arr = np.asarray(vectors_with_weights[0][0], dtype=np.float32)
        if first_arr.ndim != 1:
            raise ValueError(f"Expected 1D vectors, got shape {first_arr.shape}")
        inferred_dim = int(first_arr.shape[0])
        dim = inferred_dim if dim is None else int(dim)
    elif dim is None:
        raise ValueError("dim is required when composing an empty vector list")
    else:
        dim = int(dim)

    acc = np.zeros(dim, dtype=np.float32)
    for vec, weight in vectors_with_weights:
        arr = np.asarray(vec, dtype=np.float32)
        if arr.ndim != 1:
            raise ValueError(f"Expected 1D vectors, got shape {arr.shape}")
        if arr.shape != (dim,):
            raise ValueError(f"Expected vector shape ({dim},), got {arr.shape}")
        if weight == 0.0:
            continue
        acc += np.float32(weight) * arr

    if normalize:
        acc = dense_l2_normalize(acc)

    return acc.astype(np.dtype(dtype), copy=False)


compose_dense_vectors = compose_vectors


class VectorComposer:
    """Incrementally build a composed vector from texts and/or dense vectors."""

    def __init__(self, encoder: IdentifierEncoder, normalize_output: bool = True) -> None:
        """Create a composer bound to an existing ``IdentifierEncoder``."""
        self._encoder = encoder
        self.normalize_output = normalize_output
        self._accumulator = np.zeros(self._encoder.dim, dtype=np.float32)
        self._feature_accumulator: FeatureMap = {}
        self._text_items = 0
        self._vector_items = 0

    def add_text(self, text: str, weight: float = 1.0, family: str | None = None) -> None:
        """Encode a text item and accumulate it with the provided weight."""
        vec, feats, _ = self._encoder._encode_with_features(text, family)
        if weight != 0.0:
            self._accumulator += np.float32(weight) * vec
            for k, v in feats.items():
                self._feature_accumulator[k] = self._feature_accumulator.get(k, 0.0) + (weight * v)
        self._text_items += 1

    def add_vector(self, vec: np.ndarray, weight: float = 1.0) -> None:
        """Accumulate a precomputed dense vector with the provided weight."""
        arr = np.asarray(vec, dtype=np.float32)
        if arr.shape != (self._encoder.dim,):
            raise ValueError(f"Expected vector shape ({self._encoder.dim},), got {arr.shape}")
        if weight != 0.0:
            self._accumulator += np.float32(weight) * arr
        self._vector_items += 1

    def build(self, normalize: bool | None = None) -> np.ndarray:
        """Return the current composed vector, optionally overriding normalization."""
        do_normalize = self.normalize_output if normalize is None else normalize
        out = self._accumulator.astype(np.float32, copy=True)
        if do_normalize:
            out = dense_l2_normalize(out)
        return out

    def explain(self, top_k: int = 10) -> dict:
        """Return a compact summary of the accumulated text-side features."""
        ranked = sorted(
            ((k, abs(v)) for k, v in self._feature_accumulator.items()),
            key=lambda item: item[1],
            reverse=True,
        )
        compact = [(k, float(v)) for k, v in ranked[: max(0, int(top_k))]]
        return {
            "num_text_items": self._text_items,
            "num_vector_items": self._vector_items,
            "normalize_output": self.normalize_output,
            "top_text_features": compact,
        }

    def reset(self) -> None:
        """Clear all accumulated vectors, features, and counters."""
        self._accumulator.fill(0.0)
        self._feature_accumulator.clear()
        self._text_items = 0
        self._vector_items = 0
