"""Multi-modal encoding utilities for HIEROMEM.

This module provides a deterministic encoder that fuses heterogeneous
modalities (currently text, image, and audio placeholders) into a single
embedding that is compatible with the dense retrieval subsystem.  The
implementation mirrors the mathematical treatment in the HIEROMEM paper by
maintaining modality-specific weights that are normalised to enforce the
hierarchical bounded-growth constraint.

The encoder is intentionally dependency-light: images and audio inputs are
hashed into pseudo-random but deterministic vectors so the tests remain fully
offline.  When heavier encoders are available they can be swapped in by
subclassing :class:`MultiModalEncoder`.
"""
# flake8: noqa

from __future__ import annotations

from dataclasses import dataclass, field
import hashlib
import logging
import math
import os
from typing import Any, Iterable, Mapping

import numpy as np

from hieromem_core.encoder import Encoder, EncoderConfig


Vector = np.ndarray


def _ensure_iterable(value: Any) -> Iterable[Any]:
    if value is None:
        return []
    if isinstance(value, (list, tuple, set)):
        return value
    return [value]


def _hash_to_vector(payload: Any, dim: int, salt: str) -> Vector:
    """Create a deterministic unit vector from *payload* using SHA256."""

    text = f"{salt}:{payload!r}".encode("utf8")
    digest = hashlib.sha256(text).digest()
    # Expand digest to requested dimensionality by cycling the bytes.
    raw_uint = np.frombuffer(
        (digest * (math.ceil(dim / len(digest)))), dtype=np.uint8)
    raw = np.asarray(raw_uint[:dim], dtype=np.float32)
    centered = raw - float(np.mean(raw))
    norm = np.linalg.norm(centered) + 1e-9
    return centered / norm


@dataclass(slots=True)
class MultiModalOutput:
    """Container describing the fused embedding and modality weights."""

    vector: Vector
    modality_weights: Mapping[str, float] = field(default_factory=dict)


class MultiModalEncoder:
    """Fuse heterogeneous modalities into a single embedding.

    Parameters
    ----------
    text_encoder:
        Optional encoder used for text payloads.  Defaults to the repository's
        lightweight hash encoder.
    embedding_dim:
        Dimensionality of the fused vector.  If ``None`` the value from the
        text encoder configuration is used.
    temperature:
        Temperature parameter for the softmax that normalises modality
        contributions.
    """

    def __init__(
        self,
        text_encoder: Encoder | None = None,
        *,
        embedding_dim: int | None = None,
        temperature: float = 1.0,
    ) -> None:
        self.text_encoder = text_encoder or Encoder(EncoderConfig())
        if embedding_dim is None:
            embedding_dim = self.text_encoder.config.embedding_dim or 384
        self.embedding_dim = int(embedding_dim)
        self.temperature = float(max(temperature, 1e-6))

    # ------------------------------------------------------------------ #
    def encode(self, payload: Mapping[str, Any]) -> MultiModalOutput:
        """Encode ``payload`` and return a fused embedding.

        ``payload`` is a mapping from modality name to the corresponding raw
        data.  Supported keys are ``"text"``, ``"image"``, ``"audio"``, and any
        other user-defined modality.  Unknown modalities are hashed using the
        deterministic stub so downstream algorithms can still reason about
        them.
        """

        if not payload:
            raise ValueError(
                "MultiModalEncoder.encode requires a non-empty payload")

        modality_vectors: list[Vector] = []
        modality_strengths: list[float] = []
        modality_names: list[str] = []

        for modality, raw_value in payload.items():
            values = list(_ensure_iterable(raw_value))
            if not values:
                continue

            if modality == "text":
                vector = self._encode_text(values)
            else:
                vector = self._encode_generic(values, modality)

            modality_vectors.append(vector)
            modality_strengths.append(float(len(values)))
            modality_names.append(modality)

        if not modality_vectors:
            raise ValueError("No encodable modalities supplied in payload")

        weights = self._normalise_strengths(
            np.asarray(modality_strengths, dtype=np.float64))
        fused = np.zeros(self.embedding_dim, dtype=np.float32)

        for weight, vector in zip(weights, modality_vectors):
            fused += weight * vector

        norm = np.linalg.norm(fused)
        if norm > 0:
            fused /= norm

        modality_weights = {
    name: float(weight) for name,
    weight in zip(
        modality_names,
         weights)}
        return MultiModalOutput(
    vector=fused.astype(
        np.float32),
         modality_weights=modality_weights)

    # ------------------------------------------------------------------ #
    def _encode_text(self, texts: list[Any]) -> Vector:
        encodable = [str(item) for item in texts]
        encoded = np.asarray(
    self.text_encoder.encode(encodable),
     dtype=np.float32)
        mean_vector = np.mean(encoded, axis=0)
        norm = float(np.linalg.norm(mean_vector))
        if norm > 0:
            mean_vector /= norm
        return np.asarray(mean_vector, dtype=np.float32)

    def _encode_generic(self, values: list[Any], modality: str) -> Vector:
        # Honour explicit request to stay offline/lightweight.
        force_stub = os.getenv("HIEROMEM_FORCE_STUB_ENCODER", "").lower() in {
            "1",
            "true",
            "yes",
        }
        if force_stub:
            aggregate = np.zeros(self.embedding_dim, dtype=np.float32)
            for item in values:
                aggregate += _hash_to_vector(item, self.embedding_dim, modality)
            norm = np.linalg.norm(aggregate)
            if norm > 0:
                aggregate /= norm
            return aggregate

        # Try to get a real encoder for this modality
        try:
            from hieromem_core.encoders.factory import EncoderFactory
            real_encoder = EncoderFactory.get_encoder(modality)

            # If we have a real encoder, use it
            aggregate = np.zeros(real_encoder.dimension, dtype=np.float32)
            count = 0
            for item in values:
                try:
                    # Real encoders expect file paths or raw data
                    # If item is a string (filepath) or appropriate object, this works
                    vec = real_encoder.encode(item)

                    # Handle dimension mismatch if necessary (though we expect consistency)
                    if vec.shape[0] != self.embedding_dim:
                        # Simple projection or padding if dimensions don't match
                        # For now, we assume the user configured dimensions to match or we just warn
                        # But to be safe and keep it working with the fixed embedding_dim of the system:
                        if vec.shape[0] > self.embedding_dim:
                            vec = vec[:self.embedding_dim]
                        else:
                            padded = np.zeros(self.embedding_dim, dtype=np.float32)
                            padded[:vec.shape[0]] = vec
                            vec = padded

                    aggregate += vec
                    count += 1
                except Exception:
                    # Fallback to hash for this specific item if real encoder fails
                    aggregate += _hash_to_vector(item, self.embedding_dim, modality)
                    count += 1

            if count > 0:
                aggregate /= count  # Average pooling

            norm = np.linalg.norm(aggregate)
            if norm > 0:
                aggregate /= norm
            return aggregate.astype(np.float32)

        except (ValueError, ImportError):
            # No real encoder found for this modality, use deterministic stub
            aggregate = np.zeros(self.embedding_dim, dtype=np.float32)
            for item in values:
                aggregate += _hash_to_vector(item, self.embedding_dim, modality)
            norm = np.linalg.norm(aggregate)
            if norm > 0:
                aggregate /= norm
            return aggregate
        allow_external = os.getenv("HIEROMEM_ENABLE_MODALITY_ENCODERS", "").lower() in {
            "1",
            "true",
            "yes",
        }

        if allow_external:
            try:
                from hieromem_core.encoders.factory import EncoderFactory

                real_encoder = EncoderFactory.get_encoder(modality)
                aggregate = np.zeros(real_encoder.dimension, dtype=np.float32)
                count = 0
                for item in values:
                    try:
                        vec = real_encoder.encode(item)
                        if vec.shape[0] != self.embedding_dim:
                            if vec.shape[0] > self.embedding_dim:
                                vec = vec[: self.embedding_dim]
                            else:
                                padded = np.zeros(self.embedding_dim, dtype=np.float32)
                                padded[: vec.shape[0]] = vec
                                vec = padded
                        aggregate += vec
                        count += 1
                    except Exception:
                        aggregate += _hash_to_vector(item, self.embedding_dim, modality)
                        count += 1

                if count > 0:
                    aggregate /= count

                norm = np.linalg.norm(aggregate)
                if norm > 0:
                    aggregate /= norm
                return aggregate.astype(np.float32)
            except Exception as exc:
                # Any failure (missing deps, offline downloads, etc.) falls back
                # to the deterministic stub below.
                logging.getLogger(__name__).warning(
                    "Falling back to hashed multimodal encoding: %s", exc
                )

        aggregate = np.zeros(self.embedding_dim, dtype=np.float32)
        for item in values:
            aggregate += _hash_to_vector(item, self.embedding_dim, modality)
        norm = np.linalg.norm(aggregate)
        if norm > 0:
            aggregate /= norm
        return aggregate

    def _normalise_strengths(self, strengths: np.ndarray) -> np.ndarray:
        strengths = strengths.astype(np.float64)
        max_strength = float(np.max(strengths))
        if max_strength > 0:
            strengths = strengths / max_strength
        scaled = strengths / self.temperature
        max_scaled = float(np.max(scaled))
        scaled = scaled - max_scaled
        exp = np.exp(scaled)
        total = float(np.sum(exp))
        if total == 0:
            return np.ones_like(exp) / len(exp)
        return exp / total
