"""
Phase S — Scaling: large synthetic corpora + hippocampal consolidation gate
+ accuracy-at-scale harness.

• make_scaled_corpus(n)      : n fact entries (EN + AR + hard distractors)
• make_eval_pairs(...)       : held-out paraphrase queries (unseen views/wordings)
• ConsolidationGate          : duplicate/near-duplicate detection before write
                               (prevents memory pollution — real hippocampus
                               consolidates before permanent storage)
"""
import random
import re

import numpy as np

from phase_c_augment import augment

TOPICS_EN = ["geography", "history", "science", "music", "sports", "cooking",
             "astronomy", "biology", "physics", "cities", "rivers", "mountains",
             "oceans", "planets", "animals", "plants", "languages", "economies"]
TOPICS_AR = ["جغرافيا", "تاريخ", "علوم", "موسيقى", "رياضة", "طبخ",
             "فلك", "أحياء", "فيزياء", "مدن", "أنهار", "جبال",
             "محيطات", "كواكب", "حيوانات", "نباتات", "لغات", "اقتصاد"]

# novel-wording templates (NEVER stored — used only for evaluation)
TEMPLATES_NOVEL = [
    ("Tell me about fact {i} of {t}", "fact {i} about {t}"),
    ("I want to know fact {i} regarding {t}", "fact {i} about {t}"),
    ("Could you remind me of fact {i} in {t}?", "fact {i} about {t}"),
    ("اعرض لي الحقيقة {i} بخصوص {t}", "الحقيقة {i} حول {t}"),
    ("ممكن تقول الحقيقة رقم {i} في {t}؟", "القيقة {i} حول {t}"),
]


def make_scaled_corpus(n, seed=0, hard_ratio=0.3):
    """n entries: EN facts + AR facts + hard distractors (shared template words,
    different values) — the honest stress for lexical confusion."""
    rng = random.Random(seed)
    out = []
    for i in range(n):
        kind = i % 3
        if kind == 0:
            t = TOPICS_EN[i % len(TOPICS_EN)]
            v = rng.randint(1, 999)
            out.append({"prompt": f"What is fact number {i} about {t}?",
                        "response": f"Fact {i}: the {t} value is {v}.",
                        "domain": "scale-en"})
        elif kind == 1:
            t = TOPICS_AR[i % len(TOPICS_AR)]
            v = rng.randint(1, 999)
            out.append({"prompt": f"ما هي الحقيقة رقم {i} حول {t}؟",
                        "response": f"الحقيقة {i}: قيمة {t} هي {v}.",
                        "domain": "scale-ar"})
        else:
            # hard distractor pair: same template vocabulary, unique id slot
            t = TOPICS_EN[(i // 2) % len(TOPICS_EN)]
            v = rng.randint(1, 999)
            out.append({"prompt": f"Give me the fact {i} about the {t} value",
                        "response": f"Sure — fact {i} gives {t} = {v}.",
                        "domain": "scale-hard"})
    return out


def make_eval_pairs(corpus, n_q=2000, seed=1):
    """Held-out queries: (query, target_prompt).
    50% augmented views (unseen noise) + 50% novel wording templates."""
    rng = random.Random(seed)
    nrs = np.random.RandomState(seed)
    idx = rng.sample(range(len(corpus)), min(n_q, len(corpus)))
    pairs = []
    for k, i in enumerate(idx):
        p = corpus[i]["prompt"]
        if k % 2 == 0:
            pairs.append((augment(p, nrs), p))
        else:
            m = re.search(r"(\d+)", p)
            num = m.group(1) if m else str(i)
            m2 = (re.search(r"about (?:the )?(\w+)", p) or
                  re.search(r"حول (\w+)", p) or
                  re.search(r"the (\w+) value", p))
            topic = m2.group(1) if m2 else "topic"
            tpl_q, _ = TEMPLATES_NOVEL[k % len(TEMPLATES_NOVEL)]
            q = tpl_q.format(i=num, t=topic)
            pairs.append((q, p))
    return pairs


class ConsolidationGate:
    """Hippocampal gate with PATTERN SEPARATION (two-factor rule):
      duplicate  ⇔  semantic cosine ≥ θ  AND  slot-signature match
    (slot signature = the set of numbers/values in the prompt).

    Why: bag-of-ngrams embeddings cannot tell 'fact 42' from 'fact 43'
    (cosine ≈ 0.95+) — the signature factor keeps template siblings
    (genuinely different facts) while merging true paraphrases, exactly
    like hippocampal pattern separation. Signature index => O(1) lookup."""

    def __init__(self, encoder, threshold=0.7):
        self.enc = encoder
        self.threshold = threshold
        self._index = {}
        self._n = 0

    @staticmethod
    def slot_sig(text):
        return frozenset(re.findall(r"\d+", text))

    def _sync(self, memory):
        for j in range(self._n, len(memory.entries)):
            sig = self.slot_sig(memory.entries[j]["prompt"])
            self._index.setdefault(sig, []).append(j)
        self._n = len(memory.entries)

    def check(self, memory, prompt):
        self._sync(memory)
        sig = self.slot_sig(prompt)
        cands = self._index.get(sig, [])
        if not cands:
            return {"is_new": True, "score": 0.0, "nearest": None}
        e = self.enc.encode(prompt)
        M = np.stack([memory.codes[j] for j in cands])
        scores = M @ e
        k = int(np.argmax(scores))
        s = float(scores[k])
        return {"is_new": bool(s < self.threshold), "score": s,
                "nearest": memory.entries[cands[k]]["prompt"] if s >= self.threshold else None}

    def ingest(self, memory, entry):
        """Write only if the gate accepts; returns the gate decision."""
        d = self.check(memory, entry["prompt"])
        if d["is_new"]:
            memory.write(entry)
        return d