"""
babilong_eval.preprocessor — HCP-based context filtering for BABILong.

Three experimental conditions:
  1. full_context():    Pass-through — returns raw BABILong input unchanged
  2. hcp_filtered():    HCP seeds + queries → top chunks by holographic attention
  3. random_filtered(): Random chunks at same retain_ratio (scientific control)

The core challenge: BABILong contexts are 99%+ PG-19 book filler with
tiny bAbI facts ("Mary went to the kitchen.") injected at random positions.
HCP's content classifier (designed for instructions/constraints) classifies
ALL bAbI facts as NARRATIVE — the same type as filler.

Solution: Anomaly Scoring
  Instead of relying on content TYPE to differentiate facts from filler,
  we detect LEXICAL ANOMALIES — statistical outliers in the text that
  signal injected facts rather than natural book prose.

  bAbI facts have distinctive properties:
    1. SHORT sentences (5-12 words vs. book paragraphs of 20-100+ words)
    2. HIGH name density (person names per word)
    3. SPATIAL verbs (went, journeyed, travelled, moved)
    4. SIMPLE syntax (Subject + Verb + Preposition + Place)

  PG-19 filler (19th century literature) has:
    1. LONG sentences with complex clause structure
    2. LOW name density (relative to word count)
    3. LITERARY verbs (regarded, possessed, contemplated)
    4. COMPLEX syntax (nested clauses, semicolons)

  The anomaly score detects chunks that look "too simple and spatial"
  for their surrounding context — exactly what injected bAbI facts are.

Framework:
  KV₃ (Observer non-interference): anomaly detection FINDS injected facts
       without creating false ones
  M4  (Roughness = signal): the statistical roughness of injected facts
       amidst smooth prose IS the signal
  AX34 (Constitutive dependency): the content gap between fact and filler
       is the functional interface HCP operates on
"""

from __future__ import annotations

import math
import random
import re
from dataclasses import dataclass
from typing import Dict, FrozenSet, List, Optional, Tuple

from hcp.types import ContextChunk, HCPConfig
from hcp.protocol import HCPProtocol
from hcp.seed import extract_keywords


# ---------------------------------------------------------------------------
# Constants for anomaly scoring
# ---------------------------------------------------------------------------

# bAbI person names (complete set from BABILong tasks)
_BABI_NAMES: frozenset = frozenset({
    "mary", "john", "sandra", "daniel", "bill",
    "fred", "jeff", "bob", "julie", "jane",
})

# bAbI spatial/movement verbs
_SPATIAL_VERBS: frozenset = frozenset({
    "went", "journeyed", "travelled", "traveled", "moved",
    "got", "is", "was", "picked", "dropped", "left",
    "took", "put", "gave", "handed", "passed",
    "discarded", "grabbed",
})

# bAbI location names
_BABI_LOCATIONS: frozenset = frozenset({
    "bathroom", "bedroom", "garden", "hallway",
    "kitchen", "office", "park", "cinema",
})

# Directional words (for QA4)
_DIRECTIONS: frozenset = frozenset({
    "north", "south", "east", "west",
})


# ---------------------------------------------------------------------------
# Anomaly scoring — the KEY adaptation for narrative text
# ---------------------------------------------------------------------------

def _compute_anomaly_score(text: str) -> float:
    """Compute how anomalous a text chunk is relative to expected PG-19 prose.

    HIGH anomaly = likely a bAbI fact (short, named, spatial)
    LOW anomaly  = likely PG-19 filler (long, literary, complex)

    Factors (weighted sum, normalized to [0, 1]):
      1. Shortness: shorter sentences score higher (bAbI facts are 5-12 words)
      2. Name density: ratio of bAbI person names to total words
      3. Spatial verb density: ratio of spatial/movement verbs to total words
      4. Location density: ratio of bAbI location words to total words
      5. Sentence simplicity: low comma/semicolon density = simpler syntax

    Args:
        text: Raw text content of a chunk.

    Returns:
        Anomaly score in [0.0, 1.0]. Higher = more likely a bAbI fact.
    """
    if not text.strip():
        return 0.0

    words = text.lower().split()
    n_words = len(words)
    if n_words == 0:
        return 0.0

    # Split into sentences for per-sentence analysis
    sentences = [s.strip() for s in re.split(r"[.!?]+", text) if s.strip()]
    n_sentences = max(len(sentences), 1)

    # Factor 1: Shortness — average words per sentence
    # bAbI facts: ~6-10 words. PG-19: ~20-50 words.
    avg_words_per_sentence = n_words / n_sentences
    # Sigmoid-like: score approaches 1.0 for very short sentences
    shortness = 1.0 / (1.0 + math.exp(0.3 * (avg_words_per_sentence - 10)))

    # Factor 2: Name density — bAbI names per word
    name_count = sum(1 for w in words if w.strip(".,;:!?\"'()") in _BABI_NAMES)
    name_density = name_count / n_words

    # Factor 3: Spatial verb density
    spatial_count = sum(1 for w in words if w.strip(".,;:!?\"'()") in _SPATIAL_VERBS)
    spatial_density = spatial_count / n_words

    # Factor 4: Location density
    location_count = sum(1 for w in words if w.strip(".,;:!?\"'()") in _BABI_LOCATIONS)
    location_density = location_count / n_words

    # Factor 5: Sentence simplicity (fewer commas/semicolons = simpler)
    punctuation_count = text.count(",") + text.count(";") + text.count(":")
    punct_density = punctuation_count / n_words
    simplicity = 1.0 / (1.0 + 5.0 * punct_density)

    # Weighted combination
    score = (
        0.25 * shortness
        + 0.25 * min(name_density * 5, 1.0)     # Scale up: 1 name in 5 words → 1.0
        + 0.20 * min(spatial_density * 5, 1.0)   # Scale up: 1 verb in 5 words → 1.0
        + 0.15 * min(location_density * 5, 1.0)  # Scale up: 1 location in 5 words → 1.0
        + 0.15 * simplicity
    )

    return min(max(score, 0.0), 1.0)


# ---------------------------------------------------------------------------
# Domain-agnostic anomaly scoring — generalized for any benchmark
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class _CorpusStats:
    """Corpus-level statistics computed from all chunks in one context.

    Used by the general anomaly scorer to detect statistically unusual
    chunks WITHOUT relying on domain-specific vocabulary lists.
    """
    mean_sentence_len: float   # mean words per sentence across corpus
    std_sentence_len: float    # std deviation of words per sentence
    word_freq: Dict[str, int]  # word → count across entire corpus
    total_words: int           # total word count in corpus
    rare_threshold: int        # frequency count at 10th percentile


def _compute_corpus_stats(chunks: List[ContextChunk]) -> _CorpusStats:
    """Compute corpus-level statistics from all chunks for anomaly scoring.

    Analyzes the entire context to establish statistical baselines, then
    individual chunks are measured AGAINST these baselines. Chunks that
    deviate significantly (injected facts in filler prose) score higher.

    Args:
        chunks: All ContextChunks for one BABILong sample.

    Returns:
        _CorpusStats with mean/std sentence length, word frequencies, etc.
    """
    all_text = " ".join(c.content for c in chunks)
    words = all_text.lower().split()
    total_words = len(words) or 1

    # Word frequencies
    word_freq: Dict[str, int] = {}
    for w in words:
        clean = w.strip(".,;:!?\"'()")
        if clean:
            word_freq[clean] = word_freq.get(clean, 0) + 1

    # Sentence lengths across entire corpus
    sentences = [s.strip() for s in re.split(r"[.!?]+", all_text) if s.strip()]
    sent_lengths = [len(s.split()) for s in sentences] if sentences else [0]
    mean_len = sum(sent_lengths) / len(sent_lengths)
    variance = sum((sl - mean_len) ** 2 for sl in sent_lengths) / len(sent_lengths)
    std_len = math.sqrt(variance) if variance > 0 else 1.0

    # Rare threshold: 10th percentile of word frequencies
    freq_values = sorted(word_freq.values())
    p10_idx = max(0, int(len(freq_values) * 0.10))
    rare_threshold = freq_values[p10_idx] if freq_values else 1

    return _CorpusStats(
        mean_sentence_len=mean_len,
        std_sentence_len=std_len,
        word_freq=word_freq,
        total_words=total_words,
        rare_threshold=rare_threshold,
    )


def _compute_general_anomaly_score(
    text: str,
    stats: _CorpusStats,
    query_words: Optional[FrozenSet[str]] = None,
) -> float:
    """Domain-agnostic anomaly score: detects statistically unusual chunks.

    Unlike `_compute_anomaly_score` (bAbI-specific), this scorer uses
    corpus-level statistics to detect ANY injected/anomalous text —
    not just bAbI facts. This makes HCP applicable beyond BABILong.

    Four factors (weighted sum, normalized to [0, 1]):
      1. Length z-score (0.30): deviation from corpus mean sentence length
      2. Vocabulary rarity (0.25): ratio of corpus-rare words in chunk
      3. Query-term overlap (0.25): Jaccard with question keywords
      4. Syntax simplicity (0.20): low punctuation density

    KV₃ (Observer non-interference): scorer derives thresholds from the
         corpus itself — no hardcoded vocabulary that could create bias.

    Args:
        text: Raw text content of a chunk.
        stats: Corpus-level statistics from _compute_corpus_stats.
        query_words: Optional frozenset of lowercased question keywords.

    Returns:
        Anomaly score in [0.0, 1.0]. Higher = more statistically unusual.
    """
    if not text.strip():
        return 0.0

    words = text.lower().split()
    n_words = len(words)
    if n_words == 0:
        return 0.0

    # Split into sentences
    sentences = [s.strip() for s in re.split(r"[.!?]+", text) if s.strip()]
    n_sentences = max(len(sentences), 1)
    avg_wps = n_words / n_sentences

    # Factor 1: Length z-score — how much does avg sentence length deviate?
    z = abs(avg_wps - stats.mean_sentence_len) / max(stats.std_sentence_len, 1.0)
    length_score = 1.0 - 1.0 / (1.0 + 0.5 * z)  # z=0→0, z=2→0.5, z=4→0.67

    # Factor 2: Vocabulary rarity — ratio of words rare in corpus
    clean_words = [w.strip(".,;:!?\"'()") for w in words]
    clean_words = [w for w in clean_words if w]
    rare_count = sum(
        1 for w in clean_words
        if stats.word_freq.get(w, 0) <= stats.rare_threshold
    )
    rarity_score = rare_count / max(len(clean_words), 1)

    # Factor 3: Query-term overlap (Jaccard)
    if query_words:
        chunk_words = frozenset(clean_words)
        intersection = len(chunk_words & query_words)
        union = len(chunk_words | query_words)
        overlap_score = intersection / max(union, 1)
    else:
        overlap_score = 0.0

    # Factor 4: Syntax simplicity
    punct_count = text.count(",") + text.count(";") + text.count(":")
    punct_density = punct_count / n_words
    simplicity_score = 1.0 / (1.0 + 5.0 * punct_density)

    # Weighted combination
    score = (
        0.30 * length_score
        + 0.25 * min(rarity_score * 3, 1.0)   # Scale up: 33% rare → 1.0
        + 0.25 * min(overlap_score * 5, 1.0)   # Scale up: 20% overlap → 1.0
        + 0.20 * simplicity_score
    )

    return min(max(score, 0.0), 1.0)


def _score_chunks_with_anomaly(
    chunks: List[ContextChunk],
    question: str,
    task: str,
) -> List[Tuple[int, float]]:
    """Score each chunk using HCP attention + anomaly scoring.

    The combined score merges:
      1. HCP holographic attention (structural importance + keyword relevance)
      2. Anomaly score (statistical outlier detection for bAbI facts)

    The anomaly score AUGMENTS the HCP seed's position_importance,
    boosting chunks that look like injected facts.

    Args:
        chunks: List of ContextChunks from HCP ingestion.
        question: The BABILong question.
        task: Task name (qa1–qa5).

    Returns:
        List of (chunk_position, combined_score) sorted descending by score.
    """
    # Extract keywords from question for query-driven scoring
    question_keywords = extract_keywords(question, top_k=10)

    # Also extract bAbI-specific keywords from the question
    q_words = set(question.lower().split())
    babi_q_names = q_words & _BABI_NAMES
    babi_q_locations = q_words & _BABI_LOCATIONS
    extra_kw = tuple(babi_q_names | babi_q_locations)
    all_keywords = question_keywords + extra_kw

    # Build HCP protocol and ingest
    config = HCPConfig(
        seed_top_k=20,
        chunk_size=200,       # Smaller chunks → finer granularity for bAbI facts
        attention_decay=0.01,  # Less position decay — facts can be ANYWHERE
        seed_boost_factor=0.7,
        u_shape_boost=0.05,    # Minimize U-shape (we want content-driven attention)
    )
    protocol = HCPProtocol(config=config)
    protocol.ingest_chunks(chunks)

    if not protocol.is_ready:
        # Fallback: anomaly only
        scored = []
        for chunk in chunks:
            anomaly = _compute_anomaly_score(chunk.content)
            scored.append((chunk.position, anomaly))
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored

    # Query HCP for holographic attention weights
    results = protocol.query(question, keywords=all_keywords, top_k=len(chunks))

    # Build position → hcp_score map
    hcp_scores: Dict[int, float] = {}
    for seeded_chunk, weight in results:
        hcp_scores[seeded_chunk.chunk.position] = weight

    # Combine HCP attention with anomaly scoring
    combined: List[Tuple[int, float]] = []
    for chunk in chunks:
        hcp_score = hcp_scores.get(chunk.position, 0.0)
        anomaly = _compute_anomaly_score(chunk.content)

        # Weighted combination: 40% HCP + 60% anomaly
        # Anomaly gets more weight because HCP's content classifier
        # treats all bAbI facts as NARRATIVE (same as filler).
        # The anomaly score is what actually differentiates them.
        score = 0.4 * hcp_score + 0.6 * anomaly
        combined.append((chunk.position, score))

    combined.sort(key=lambda x: x[1], reverse=True)
    return combined


# ---------------------------------------------------------------------------
# Chunking helper
# ---------------------------------------------------------------------------

def _chunk_context(context: str, chunk_size: int = 200) -> List[ContextChunk]:
    """Split BABILong context into ContextChunks.

    Uses sentence-boundary splitting to avoid breaking bAbI facts mid-sentence.
    Small chunk_size (200 chars) to isolate individual facts.

    Args:
        context: The raw context string from BABILong.
        chunk_size: Maximum characters per chunk.

    Returns:
        List of ContextChunks with positions.
    """
    # Split by sentences
    sentences = re.split(r"(?<=[.!?])\s+", context.strip())

    chunks: List[ContextChunk] = []
    current = ""
    position = 0

    for sentence in sentences:
        if len(current) + len(sentence) + 1 > chunk_size and current:
            chunks.append(ContextChunk(
                content=current.strip(),
                position=position,
            ))
            position += 1
            current = sentence
        else:
            current = f"{current} {sentence}" if current else sentence

    if current.strip():
        chunks.append(ContextChunk(
            content=current.strip(),
            position=position,
        ))

    return chunks


# ---------------------------------------------------------------------------
# Public API: Three experimental conditions
# ---------------------------------------------------------------------------

def full_context(context: str, **kwargs) -> str:
    """Condition 1: Pass-through — return raw context unchanged.

    This is the baseline: the LLM sees the FULL BABILong context
    (filler + facts) exactly as provided.

    Args:
        context: Raw BABILong context string.

    Returns:
        The same context string, unchanged.
    """
    return context


def hcp_filtered(
    context: str,
    question: str,
    task: str,
    retain_ratio: float = 0.25,
) -> str:
    """Condition 2: HCP-filtered context.

    Uses HCP's holographic attention + anomaly scoring to identify
    the most likely relevant chunks, then returns only the top
    retain_ratio fraction.

    Args:
        context: Raw BABILong context string.
        question: The question to answer.
        task: Task name (qa1–qa5).
        retain_ratio: Fraction of chunks to keep (0.0 to 1.0).

    Returns:
        Filtered context containing only the highest-scoring chunks,
        in their ORIGINAL order (to preserve narrative flow).
    """
    if retain_ratio >= 1.0:
        return context

    chunks = _chunk_context(context)
    if not chunks:
        return context

    # Score all chunks
    scored = _score_chunks_with_anomaly(chunks, question, task)

    # Keep top retain_ratio
    n_keep = max(1, int(len(chunks) * retain_ratio))
    top_positions = set(pos for pos, _ in scored[:n_keep])

    # Reconstruct in original order
    kept = [c for c in chunks if c.position in top_positions]
    kept.sort(key=lambda c: c.position)

    return " ".join(c.content for c in kept)


def random_filtered(
    context: str,
    retain_ratio: float = 0.25,
    seed: Optional[int] = None,
    **kwargs,
) -> str:
    """Condition 3: Random-filtered context (scientific control).

    Keeps the same number of chunks as hcp_filtered but selected
    RANDOMLY. This isolates whether HCP's scoring is better than
    random chance.

    If hcp_filtered accuracy > random_filtered accuracy, then HCP
    genuinely identifies relevant content.

    Args:
        context: Raw BABILong context string.
        retain_ratio: Fraction of chunks to keep.
        seed: Random seed for reproducibility.

    Returns:
        Randomly filtered context with chunks in original order.
    """
    if retain_ratio >= 1.0:
        return context

    chunks = _chunk_context(context)
    if not chunks:
        return context

    n_keep = max(1, int(len(chunks) * retain_ratio))

    rng = random.Random(seed)
    selected = rng.sample(range(len(chunks)), min(n_keep, len(chunks)))
    selected_set = set(selected)

    # Reconstruct in original order
    kept = [c for i, c in enumerate(chunks) if i in selected_set]
    kept.sort(key=lambda c: c.position)

    return " ".join(c.content for c in kept)


# ---------------------------------------------------------------------------
# Condition 4: BM25 retrieval baseline
# ---------------------------------------------------------------------------

def bm25_filtered(
    context: str,
    question: str,
    retain_ratio: float = 0.25,
    **kwargs,
) -> str:
    """Condition 4: BM25 retrieval baseline — lexical chunk selection.

    Uses BM25 (Okapi variant) to score chunks by term-frequency-based
    relevance to the question. Pure information retrieval baseline —
    no anomaly scoring, no HCP, no embeddings.

    BM25 is the standard sparse-retrieval baseline in IR research.
    It captures term overlap between query and document but has NO
    semantic understanding.

    Args:
        context: Raw BABILong context string.
        question: The question to answer.
        retain_ratio: Fraction of chunks to keep (0.0 to 1.0).

    Returns:
        BM25-filtered context with top chunks in original order.
    """
    if retain_ratio >= 1.0:
        return context

    chunks = _chunk_context(context)
    if not chunks:
        return context

    from rank_bm25 import BM25Okapi

    # Tokenize corpus (simple whitespace tokenization, lowercased)
    tokenized_corpus = [c.content.lower().split() for c in chunks]
    bm25 = BM25Okapi(tokenized_corpus)

    # Score all chunks against question
    tokenized_query = question.lower().split()
    scores = bm25.get_scores(tokenized_query)

    # Keep top retain_ratio
    n_keep = max(1, int(len(chunks) * retain_ratio))
    top_indices = sorted(
        range(len(scores)), key=lambda i: scores[i], reverse=True
    )[:n_keep]
    top_set = set(top_indices)

    # Reconstruct in original order
    kept = [chunks[i] for i in sorted(top_set)]
    return " ".join(c.content for c in kept)


# ---------------------------------------------------------------------------
# Condition 5: Embedding retrieval baseline
# ---------------------------------------------------------------------------

_EMBEDDING_MODEL = None


def _get_embedding_model():
    """Lazy-load sentence-transformers model (~80MB download on first use).

    Uses all-MiniLM-L6-v2 — fast, lightweight, good quality for IR.
    Singleton pattern: loaded once and cached for the process lifetime.
    """
    global _EMBEDDING_MODEL
    if _EMBEDDING_MODEL is None:
        from sentence_transformers import SentenceTransformer
        _EMBEDDING_MODEL = SentenceTransformer("all-MiniLM-L6-v2")
    return _EMBEDDING_MODEL


def embedding_filtered(
    context: str,
    question: str,
    retain_ratio: float = 0.25,
    **kwargs,
) -> str:
    """Condition 5: Embedding retrieval baseline — semantic chunk selection.

    Uses sentence-transformers (all-MiniLM-L6-v2) to encode chunks and
    question into dense vectors, then selects top chunks by cosine
    similarity. Dense-retrieval baseline — captures semantic meaning
    but requires neural network inference.

    This is the strongest retrieval baseline: it understands meaning,
    not just keyword overlap. If HCP beats embedding retrieval, that's
    a strong result. If not, it tells us about HCP's relative power.

    Args:
        context: Raw BABILong context string.
        question: The question to answer.
        retain_ratio: Fraction of chunks to keep (0.0 to 1.0).

    Returns:
        Embedding-filtered context with top chunks in original order.
    """
    if retain_ratio >= 1.0:
        return context

    chunks = _chunk_context(context)
    if not chunks:
        return context

    import gc
    import numpy as np

    model = _get_embedding_model()

    # Encode chunks in small batches to avoid OOM on CPU
    chunk_texts = [c.content for c in chunks]
    batch_size = 32
    all_embeds = []
    for start in range(0, len(chunk_texts), batch_size):
        batch = chunk_texts[start : start + batch_size]
        emb = model.encode(batch, show_progress_bar=False)
        all_embeds.append(emb)
    chunk_embeddings = np.concatenate(all_embeds, axis=0)
    question_embedding = model.encode([question], show_progress_bar=False)[0]
    gc.collect()

    # Cosine similarity
    norms_chunks = np.linalg.norm(chunk_embeddings, axis=1)
    norm_question = np.linalg.norm(question_embedding)
    # Guard against zero norms
    norms_chunks = np.maximum(norms_chunks, 1e-10)
    norm_question = max(norm_question, 1e-10)

    similarities = np.dot(chunk_embeddings, question_embedding) / (
        norms_chunks * norm_question
    )

    # Keep top retain_ratio
    n_keep = max(1, int(len(chunks) * retain_ratio))
    top_indices = np.argsort(similarities)[-n_keep:]
    top_set = set(top_indices.tolist())

    # Reconstruct in original order
    kept = [chunks[i] for i in sorted(top_set)]
    return " ".join(c.content for c in kept)
