"""Decontamination heuristics: behaviour on hand-built pairs (renamed binders, reordered hypotheses, numeral
types, unrelated statements). The expectations here ARE the measured false-positive / false-negative profile."""
from __future__ import annotations

import numpy as np

from align.data.decontam import (
    ContaminationLedger, Decontaminator, EvalItem, Hit, LSHIndex, MinHasher, Params, lean_tokens, nl_tokens,
    ngram_set, shingles, text_hash,
)

H = "import Mathlib\nimport Aesop\n\nset_option maxHeartbeats 0\n\nopen BigOperators Real Nat Topology Rat\n\n"

BASE = H + """/-- Let x and y be reals with x + y = 10 and x - y = 2. Prove that the product x * y equals 24 given that both
are positive and their sum of squares is 52. -/
theorem base_problem (x y : ℝ) (h₀ : 0 < x ∧ 0 < y) (h₁ : x + y = 10) (h₂ : x - y = 2) (h₃ : x ^ 2 + y ^ 2 = 52) :
  x * y = 24 := by sorry"""
import re as _re


def _rename(src, mp):
    return _re.sub(r"(?<![\w'])(" + "|".join(_re.escape(k) for k in mp) + r")(?![\w'])", lambda m: mp[m.group(0)], src)


RENAMED = _rename(BASE, {"base_problem": "other_name", "x": "a", "y": "b", "h₀": "p", "h₁": "q", "h₂": "r", "h₃": "s"})
EVAL_OTHER = H + "theorem other_eval (n : ℕ) (hn : 3 < n) (hm : n < 100) (hp : n.Prime) : ¬ (n % 6 = 0) := by sorry"
REORDERED = H + """theorem base_problem (x y : ℝ) (h₃ : x ^ 2 + y ^ 2 = 52) (h₂ : x - y = 2) (h₁ : x + y = 10) (h₀ : 0 < x ∧ 0 < y) :
  x * y = 24 := by sorry"""
NUMERAL = BASE.replace("(x y : ℝ)", "(x y : ℕ)").replace("= 24", "= (24 : ℕ)")
CHANGED_CONST = BASE.replace("= 52", "= 58").replace("= 24", "= 21")
UNRELATED = H + """theorem unrelated (n : ℕ) (hn : 0 < n) (h : n % 7 = 3) : (n ^ 3 + 2 * n + 5) % 7 = 3 := by sorry"""
SHORT_A = H + "theorem a (x : ℝ) : x = x := by sorry"
SHORT_B = H + "theorem b (y : ℝ) : y = y := by sorry"


def dec(statement=BASE, informal=None, **kw):
    return Decontaminator([EvalItem("E", "evalset", statement, informal)], Params(**kw) if kw else None)


def stages(d, stmt, informal=None):
    return {h.stage for h in d.check(stmt, informal)}


# ---------------------------------------------------------------- minhash

def test_minhash_estimates_jaccard():
    mh = MinHasher(num_perm=256, seed=3)
    a = {f"s{i}" for i in range(200)}
    b = {f"s{i}" for i in range(100, 300)}  # |A∩B|=100, |A∪B|=300, J=1/3
    sa = mh.signature(np.array([hash_(x) for x in a], dtype=np.uint64))
    sb = mh.signature(np.array([hash_(x) for x in b], dtype=np.uint64))
    assert abs(MinHasher.jaccard(sa, sb) - 1 / 3) < 0.09
    assert MinHasher.jaccard(sa, sa) == 1.0


def hash_(s):
    import zlib
    return zlib.crc32(s.encode())


def test_lsh_recalls_similar_not_disjoint():
    mh = MinHasher()
    idx = LSHIndex(mh, bands=32, threshold=0.5)
    base = [f"t{i}" for i in range(60)]
    idx.add("base", mh.signature(shingles(base, 3)))
    near = base[:55] + ["x1", "x2", "x3", "x4", "x5"]
    far = [f"z{i}" for i in range(60)]
    assert [i for i, _ in idx.query(mh.signature(shingles(near, 3)))] == ["base"]
    assert idx.query(mh.signature(shingles(far, 3))) == []
    assert idx.query(None) == []


def test_signature_is_deterministic_across_instances():
    a = MinHasher(seed=1).signature(shingles(list("abcdefgh"), 3))
    b = MinHasher(seed=1).signature(shingles(list("abcdefgh"), 3))
    assert (a == b).all()


# ---------------------------------------------------------------- statement stages: measured profile

def test_identical_hits_every_exact_stage():
    s = stages(dec(), BASE)
    assert {"text_hash", "sig_sha", "alpha_sha", "minhash_lsh", "ngram_lean"} <= s


def test_renamed_theorem_name_and_comment_only_is_text_hash_hit():
    s = stages(dec(), BASE.replace("base_problem", "renamed_thm").replace("Let x and y", "Suppose x and y"))
    assert "text_hash" in s and "sig_sha" in s


def test_renamed_binders_caught_by_alpha_not_by_exact_hashes():
    s = stages(dec(), RENAMED)
    assert "alpha_sha" in s
    assert "sig_sha" not in s and "text_hash" not in s  # exact stages are blind to renaming (by design)
    assert "minhash_lsh" in s  # alpha-renamed tokens make the fuzzy stages invariant too


def test_reordered_hypotheses_missed_by_alpha_caught_by_fuzzy():
    s = stages(dec(), REORDERED)
    assert "alpha_sha" not in s and "text_hash" not in s  # known false negative of the hash stages
    assert "minhash_lsh" in s or "ngram_lean" in s       # the fuzzy stages recover it


def test_numeral_type_variant_is_flagged_by_fuzzy_stages_but_not_by_hashes():
    s = stages(dec(), NUMERAL)
    assert not ({"text_hash", "sig_sha", "alpha_sha"} & s)   # different statement: exact stages correctly differ
    assert "minhash_lsh" in s or "ngram_lean" in s           # over-flagged on purpose (documented false positive)


def test_changed_constants_not_flagged_exactly():
    s = stages(dec(), CHANGED_CONST)
    assert not ({"text_hash", "sig_sha", "alpha_sha"} & s)


def test_unrelated_statement_no_hit_at_all():
    assert dec().check(UNRELATED) == []


def test_short_generic_statements_collide_only_in_alpha_hash():
    """`x = x` vs `y = y` is the same problem up to renaming (alpha_sha says so, correctly); the fuzzy stages are
    disabled below `min_lean_tokens` precisely because tiny statements all look alike."""
    d = dec(SHORT_A)
    s = stages(d, SHORT_B)
    assert "alpha_sha" in s
    assert "minhash_lsh" not in s and "ngram_lean" not in s


def test_alpha_sha_conservative_false_positive_documented():
    a = H + "theorem a (x : ℝ) (h : Real.sin x = 1) : x ≠ 0 := by sorry"
    b = H + "theorem b (y : ℝ) (h : Real.cos y = 1) : y ≠ 0 := by sorry"
    # p2_common keeps `sin`/`cos` lowercase-dotted as distinct tokens: the dotted names are NOT renamed
    assert "alpha_sha" not in stages(dec(a), b)


def test_lean_tokens_alpha_invariance():
    assert lean_tokens(BASE) == lean_tokens(RENAMED)
    assert lean_tokens(BASE, alpha=False) != lean_tokens(RENAMED, alpha=False)


def test_text_hash_ignores_comments_header_name():
    assert text_hash(BASE) == text_hash(BASE.replace("base_problem", "zzz").replace("Let x", "Take x"))
    assert text_hash(BASE) != text_hash(CHANGED_CONST)


# ---------------------------------------------------------------- informal text

INFORMAL = ("Let x and y be real numbers such that x plus y equals ten and x minus y equals two, "
            "find the product of x and y when both are positive and their squares sum to fifty two.")


def test_informal_ngram_catches_embedded_copy_and_hash_catches_exact():
    d = dec(EVAL_OTHER, INFORMAL)
    cand_informal = "Problem 7. " + INFORMAL + " Show your work."
    s = stages(d, UNRELATED.replace("unrelated", "q"), cand_informal)
    assert "ngram_nl" in s
    assert "informal_hash" in stages(d, UNRELATED.replace("unrelated", "q"), INFORMAL)


def test_informal_paraphrase_below_ngram_is_missed_known_false_negative():
    d = dec(EVAL_OTHER, INFORMAL)
    para = "Two positive reals have sum 10 and difference 2; compute their product given the sum of squares is 52."
    assert "ngram_nl" not in stages(d, UNRELATED.replace("unrelated", "q"), para)


def test_informal_unrelated_no_hit():
    d = dec(EVAL_OTHER, INFORMAL)
    assert stages(d, UNRELATED.replace("unrelated", "q"), "How many ways can five people sit around a circular table?") == set()


def test_docstring_of_eval_statement_counts_as_informal():
    d = dec(BASE, None)  # BASE's docstring is the informal text
    cand = "Let x and y be reals with x + y = 10 and x - y = 2. Prove that the product x * y equals 24 given that both"
    assert "ngram_nl" in stages(d, UNRELATED.replace("unrelated", "q"), cand)


def test_ngram_nl_ignores_formula_streams_and_stock_phrases():
    d = dec(EVAL_OTHER, "Show that there exists a positive integer n such that a 2 b 2 b 2 c 2 c 2 d 2 holds for all values.")
    cand = UNRELATED.replace("unrelated", "q")
    assert "ngram_nl" not in stages(d, cand, "Prove: show that there exists a positive integer n such that the claim is trivial")
    assert "ngram_nl" not in stages(d, cand, "We expand a 2 b 2 b 2 c 2 c 2 d 2 by hand in this example")


def test_nl_helpers():
    assert nl_tokens("A, b-C 12!") == ["a", "b", "c", "12"]
    assert len(ngram_set(list("abcdef"), 5)) == 2 and ngram_set(list("abc"), 5) == set()


# ---------------------------------------------------------------- elab hash + ancestry

def test_elab_hash_stage():
    d = Decontaminator([EvalItem("E", "s", UNRELATED, elab_sha="deadbeef")])
    assert [h.stage for h in d.check("theorem z : True := by sorry", elab_sha="deadbeef")] == ["elab_hash"]
    assert d.check("theorem z : True := by sorry", elab_sha="other") == []


def test_ancestry_child_inherits_parent_contamination():
    led = ContaminationLedger()
    led.add("parent", [Hit("alpha_sha", "E")], {"root": "parent", "parents": []})
    led.add("child", [], {"root": "parent", "parents": ["parent"]})
    led.add("grandchild", [], {"root": "parent", "parents": ["child"]})
    led.add("other", [], {"root": "other", "parents": []})
    assert led.resolve("child") == (True, "inherited:parent")
    assert led.resolve("grandchild")[0] is True
    assert led.resolve("other") == (False, "clean")
    assert led.resolve("parent") == (True, "direct:alpha_sha")


def test_ancestry_unknown_parent_is_reported_not_assumed_clean():
    led = ContaminationLedger()
    led.add("child", [], {"root": "ghost", "parents": ["ghost"]})
    assert led.resolve("child") == (False, "unknown_parent")
    led2 = ContaminationLedger(known={"ghost": True})
    led2.add("child", [], {"root": "ghost", "parents": ["ghost"]})
    assert led2.resolve("child")[0] is True


def test_ancestry_cycle_terminates():
    led = ContaminationLedger()
    led.add("a", [], {"parents": ["b"]})
    led.add("b", [], {"parents": ["a"]})
    assert led.resolve("a")[0] is False