"""
hcp.attention — Attention simulation: flat (baseline) vs holographic (HCP).

Simulates LLM attention weight distribution over context positions.
Does NOT require an actual LLM — models the DOCUMENTED attention pathology
and demonstrates how the holographic seed corrects it.

BASELINE (flat_attention):
  Implements the U-shaped attention bias documented by:
    - Liu et al. TACL 2024 ("Lost in the Middle"): LLMs degrade when
      relevant information is in the middle of long contexts
    - MIT (2025): position bias is ARCHITECTURAL — caused by causal
      attention masking and rotary positional encoding, not capacity
    - U-shape: high attention at beginning/end, degraded in middle

HOLOGRAPHIC (holographic_attention):
  Seed-modulated attention that uses the HolographicSeed's structural
  pointers to boost attention at structurally important positions,
  REGARDLESS of their location.  This inverts the positional bias:
  attention is driven by CONTENT STRUCTURE, not by position.

Framework Axioms:
  AX37: Holographic Structure — seed carries global structure to every point
  AX13: Ehadiyet — local reflection of global unity
  AX34: Constitutive dependency — the attention gap is the functional interface
        that the seed fills
  T15:  Structural > detail — structural importance weights > positional bias

KV₇ compliance: This module imports ONLY from hcp.types and standard library.
"""

from __future__ import annotations

import math
from typing import List, Optional, Tuple

from hcp.types import (
    AttentionProfile,
    HCPConfig,
    HolographicSeed,
    clamp_score,
)


# ---------------------------------------------------------------------------
# Flat Attention (Baseline) — U-shaped positional bias
# ---------------------------------------------------------------------------

def flat_attention(
    n_positions: int,
    query_position: int = -1,
    config: Optional[HCPConfig] = None,
) -> AttentionProfile:
    """Simulate the flat (standard) attention pattern over context positions.

    Models the U-shaped attention bias documented in empirical research:
    - High attention at the START (primacy effect)
    - High attention at the END (recency effect)
    - DEGRADED attention in the MIDDLE (the "lost in the middle" problem)

    The mathematical model:
      w(i) = base + u_shape(i) + distance_decay(i)

    where:
      - base: uniform floor (all positions get some attention)
      - u_shape: parabolic boost at edges, minimum at center
      - distance_decay: exponential decay from query position

    This accurately models the documented pathology:
      - MIT 2025: attention masking creates positional bias
      - Liu et al. 2024: 20+ percentage point drop for middle content

    Args:
        n_positions: Total number of context positions.
        query_position: Position of the query (-1 = end, default).
        config: HCPConfig for decay rate and U-shape parameters.

    Returns:
        AttentionProfile with U-shaped bias.
    """
    if n_positions <= 0:
        return AttentionProfile(weights=(), mode="flat")

    if config is None:
        config = HCPConfig()

    if query_position < 0:
        query_position = n_positions - 1

    weights: List[float] = []
    for i in range(n_positions):
        # Base: uniform floor
        base = 0.3

        # U-shape: parabola with minimum at center
        # Normalized so edge = +u_boost, center = 0
        center = (n_positions - 1) / 2.0
        if n_positions > 1:
            distance_from_center = abs(i - center) / center  # [0, 1]
        else:
            distance_from_center = 0.0
        u_shape = config.u_shape_boost * (distance_from_center ** 2)

        # Distance decay from query position
        distance_from_query = abs(i - query_position)
        decay = math.exp(-config.attention_decay * distance_from_query)

        weight = base + u_shape + 0.4 * decay
        weights.append(clamp_score(weight))

    # Normalize to sum to 1 (probability distribution)
    total = sum(weights)
    if total > 0:
        weights = [clamp_score(w / total) for w in weights]

    return AttentionProfile(weights=tuple(weights), mode="flat")


# ---------------------------------------------------------------------------
# Holographic Attention — Seed-modulated, structure-driven
# ---------------------------------------------------------------------------

def holographic_attention(
    n_positions: int,
    seed: HolographicSeed,
    query_position: int = -1,
    query_keywords: Tuple[str, ...] = (),
    per_position_keywords: Optional[List[Tuple[str, ...]]] = None,
    config: Optional[HCPConfig] = None,
) -> AttentionProfile:
    """Simulate holographic (HCP) attention over context positions.

    The KEY INNOVATION: attention is modulated by the HolographicSeed,
    which encodes GLOBAL context structure.  This means:

    1. Structurally important positions get boosted regardless of location
    2. Instruction/constraint positions always receive high attention
    3. Positions sharing entities with the query get boosted
    4. The U-shaped bias is COUNTERACTED by structural importance

    Mathematical model:
      w_hcp(i) = (1 - α) · w_flat(i) + α · importance(i) + β · relevance(i)

    where:
      - w_flat(i): baseline flat attention (the bias we're correcting)
      - importance(i): seed's position_importance (content-driven)
      - relevance(i): keyword overlap between query and position
      - α: seed_boost_factor (how much the seed overrides position bias)
      - β: query-relevance factor

    This implements AX37 (Holographic Architecture): the seed ensures
    that every position "knows about" the global structure, so attention
    can flow to the RIGHT positions, not just nearby positions.

    Args:
        n_positions: Total number of context positions.
        seed: The HolographicSeed computed from the full context.
        query_position: Position of the query (-1 = end).
        query_keywords: Keywords from the query (for relevance matching).
        config: HCPConfig for boost factors.

    Returns:
        AttentionProfile with seed-modulated weights.
    """
    if n_positions <= 0:
        return AttentionProfile(weights=(), mode="holographic")

    if config is None:
        config = HCPConfig()

    if query_position < 0:
        query_position = n_positions - 1

    # Start with flat attention as baseline
    flat = flat_attention(n_positions, query_position, config)

    alpha = config.seed_boost_factor  # How much seed overrides position

    weights: List[float] = []
    for i in range(n_positions):
        # Baseline (with its U-shaped bias)
        w_flat = flat.weights[i] if i < len(flat.weights) else 0.3

        # Structural importance from seed (position-independent!)
        if i < len(seed.position_importance):
            importance = seed.position_importance[i]
        else:
            importance = 0.0

        # Instruction/constraint boost (these are ALWAYS important)
        structural_boost = 0.0
        if i in seed.instruction_positions:
            structural_boost += 0.3
        if i in seed.constraint_positions:
            structural_boost += 0.25

        # Entity relevance: if query keywords overlap with entities at this position
        entity_boost = 0.0
        if query_keywords:
            query_set = set(kw.lower() for kw in query_keywords)
            for entity, positions in seed.entity_positions.items():
                if i in positions and entity in query_set:
                    entity_boost += 0.2

        # Keyword relevance: global keywords present at this position
        keyword_boost = 0.0
        if query_keywords and seed.global_keywords:
            query_set = set(kw.lower() for kw in query_keywords)
            global_set = set(seed.global_keywords)
            overlap = query_set & global_set
            if overlap:
                keyword_boost = 0.1 * len(overlap) / max(len(query_set), 1)

        # Per-position keyword relevance: direct content-query matching
        # This is the KEY retrieval mechanism — positions whose CONTENT
        # matches the query get boosted regardless of location.
        content_relevance = 0.0
        if query_keywords and per_position_keywords and i < len(per_position_keywords):
            query_set = set(kw.lower() for kw in query_keywords)
            pos_set = set(kw.lower() for kw in per_position_keywords[i])
            if pos_set:
                overlap = query_set & pos_set
                # Jaccard-like relevance: overlap / union
                union_size = len(query_set | pos_set)
                if union_size > 0:
                    content_relevance = 0.5 * len(overlap) / union_size

        # Combined: interpolate flat with structural
        relevance = structural_boost + entity_boost + keyword_boost + content_relevance
        w_hcp = (1 - alpha) * w_flat + alpha * importance + relevance

        weights.append(clamp_score(w_hcp))

    # Normalize
    total = sum(weights)
    if total > 0:
        weights = [clamp_score(w / total) for w in weights]

    return AttentionProfile(weights=tuple(weights), mode="holographic")


# ---------------------------------------------------------------------------
# Comparison utility
# ---------------------------------------------------------------------------

def attention_improvement(
    flat: AttentionProfile,
    holographic: AttentionProfile,
    target_position: int,
) -> float:
    """Compute the attention improvement at a target position.

    Returns the ratio of holographic / flat attention at the target.
    Values > 1.0 indicate the holographic approach puts MORE attention
    on the target.  This is most significant when the target is in
    the middle third (the "danger zone").

    Args:
        flat: Flat attention profile.
        holographic: Holographic attention profile.
        target_position: The position to compare.

    Returns:
        Improvement ratio.  Capped at 0.9999 per KV₄.
    """
    if target_position >= len(flat.weights) or target_position >= len(holographic.weights):
        return 0.0

    flat_w = flat.weights[target_position]
    holo_w = holographic.weights[target_position]

    if flat_w <= 0:
        return clamp_score(holo_w * 100)  # Infinite improvement capped

    return clamp_score(holo_w / flat_w)
