"""
Drift computation utilities for MRS.

Drift measures local coherence deviation between reasoning steps.
Multiple strategies available:
    - Hash-based (fast, deterministic)
    - Length-based (simple heuristic)
    - Semantic (requires embeddings)
"""

from typing import Any, Optional
import hashlib


def hash_drift(text: str) -> float:
    """
    Compute drift using hash-based method (fast, deterministic).

    Maps text to a value in [0, 1] based on its hash.

    Args:
        text: Input text to compute drift for

    Returns:
        Drift value between 0.0 and 1.0
    """
    if not text:
        return 0.0
    text_bytes = str(text).encode("utf-8")
    hash_value = int(hashlib.md5(text_bytes).hexdigest(), 16)
    return (hash_value % 1000) / 1000.0


def length_drift(current: str, previous: Optional[str] = None) -> float:
    """
    Compute drift based on output length change.

    Args:
        current: Current output text
        previous: Previous output text (optional)

    Returns:
        Drift value between 0.0 and 1.0
    """
    if not previous:
        return hash_drift(current)

    curr_len = len(current)
    prev_len = len(previous)

    if prev_len == 0:
        return 0.5

    ratio = abs(curr_len - prev_len) / max(curr_len, prev_len)
    return min(1.0, ratio)


def semantic_drift(
    current: str, previous: str, embedding_fn: Optional[callable] = None
) -> float:
    """
    Compute semantic drift using embeddings (if available).

    Falls back to hash_drift if no embedding function provided.

    Args:
        current: Current output text
        previous: Previous output text
        embedding_fn: Function that returns embeddings for text

    Returns:
        Drift value between 0.0 and 1.0
    """
    if embedding_fn is None:
        combined = f"{previous}|||{current}"
        return hash_drift(combined)

    try:
        emb_current = embedding_fn(current)
        emb_previous = embedding_fn(previous)

        dot_product = sum(a * b for a, b in zip(emb_current, emb_previous))
        norm_current = sum(a * a for a in emb_current) ** 0.5
        norm_previous = sum(b * b for b in emb_previous) ** 0.5

        if norm_current == 0 or norm_previous == 0:
            return 0.5

        cosine_sim = dot_product / (norm_current * norm_previous)

        return (1.0 - cosine_sim) / 2.0
    except Exception:
        return hash_drift(f"{previous}|||{current}")


def compute_drift(
    output: Any, previous_output: Optional[Any] = None, method: str = "hash"
) -> float:
    """
    Unified drift computation interface.

    Args:
        output: Current reasoning output
        previous_output: Previous reasoning output (optional)
        method: Drift computation method ("hash", "length", "semantic")

    Returns:
        Drift value between 0.0 and 1.0
    """
    current_str = str(output)
    previous_str = str(previous_output) if previous_output else None

    if method == "hash":
        return hash_drift(current_str)
    elif method == "length":
        return length_drift(current_str, previous_str)
    elif method == "semantic":
        return semantic_drift(current_str, previous_str or "")
    else:
        return hash_drift(current_str)
