"""randomization_implementation_finder.py – precision/recall ladder for *randomization implementation* (who generated sequence, who enrolled, who assigned).
Five variants (v1–v5):
    • v1 – high recall: any implementation cue (sequence generated by statistician, independent pharmacist, clinicians enrolled, central system assigned, assignments generated by, etc.).
    • v2 – role cue + action verb (generated, enrolled, assigned) within ±4 tokens.
    • v3 – only inside a *Randomisation Implementation* or *Participants / Assignment* heading block (first ~500 chars).
    • v4 – v2 plus presence of at least two distinct roles (e.g., statistician + investigator + system) in same sentence.
    • v5 – tight template: “Statistician generated sequence; investigators enrolled participants; central web system assigned groups.”
Each finder returns tuples: (start_word_idx, end_word_idx, snippet).
"""
from __future__ import annotations
import re
from typing import List, Tuple, Sequence, Dict, Callable

TOKEN_RE = re.compile(r"\S+")

def _token_spans(text: str) -> List[Tuple[int, int]]:
    return [(m.start(), m.end()) for m in TOKEN_RE.finditer(text)]

def _char_to_word(span: Tuple[int, int], spans: Sequence[Tuple[int, int]]):
    s, e = span
    w_s = next(i for i, (a, b) in enumerate(spans) if a <= s < b)
    w_e = next(i for i, (a, b) in reversed(list(enumerate(spans))) if a < e <= b)
    return w_s, w_e

ROLE_RE = re.compile(
    r"\b(?:statisticians?|data\s+managers?|pharmacists?|investigators?|clinicians?|nurses?|research\s+assistants?|independent|central\s+system|web[- ]?based\s+system|interactive\s+voice\s+response|ivr|iwrs)\b",
    re.I,
)
ACTION_RE = re.compile(r"\b(?:generated|prepared|created|enrolled|screened|assigned|allocated|registered|entered)\b", re.I)
IMPLEMENT_CUE_RE = re.compile(r"\b(?:sequence\s+generated\s+by|generated\s+(?:the\s+)?sequence|(?:investigators?|clinicians?|nurses?)\s+enrolled|enrolled\s+(?:participants?|patients?)|(?:central\s+)?(?:system|web|ivr|iwrs)\s+assigned|assigned\s+(?:groups?|interventions?))\b", re.I)
HEAD_RE = re.compile(r"(?m)^(?:randomi[sz]ation\s+implementation|implementation|assignment|enrollment)\s*[:\-]?\s*$", re.I)
TRAP_RE = re.compile(r"\bimplemented\s+the\s+treatment|implemented\s+protocol\b", re.I)
TIGHT_TEMPLATE_RE = re.compile(r"statistician\s+generated\s+[^\.\n]{0,80}?(?:investigators?|clinicians?|nurses?)\s+enrolled[^\.\n]{0,80}?(?:central\s+)?(?:web|system|ivr|iwrs)[^\.\n]{0,80}?assigned", re.I)
OBJECT_RE = re.compile(r"\b(?:sequence|list|allocation|randomi[sz]ation|participants?|groups?|interventions?)\b", re.I)

def _collect(patterns: Sequence[re.Pattern[str]], text: str):
    spans = _token_spans(text)
    out: List[Tuple[int, int, str]] = []
    for patt in patterns:
        for m in patt.finditer(text):
            if TRAP_RE.search(text[max(0, m.start()-30):m.end()+30]):
                continue
            w_s, w_e = _char_to_word((m.start(), m.end()), spans)
            out.append((w_s, w_e, m.group(0)))
    return out

def find_randomization_implementation_v1(text: str):
    return _collect([IMPLEMENT_CUE_RE], text)

def find_randomization_implementation_v2(text: str, window: int = 4):
    spans = _token_spans(text)
    tokens = [text[s:e] for s, e in spans]
    role_idx = {i for i, t in enumerate(tokens) if ROLE_RE.search(t)}
    act_idx = {i for i, t in enumerate(tokens) if ACTION_RE.search(t)}
    obj_idx = {i for i, t in enumerate(tokens) if OBJECT_RE.search(t)}
    out = []
    for r in role_idx:
        for a in act_idx:
            if abs(a - r) <= window:
                if any(abs(o - a) <= window for o in obj_idx):
                    w_s, w_e = _char_to_word(spans[r], spans)
                    out.append((w_s, w_e, tokens[r]))
    return out

def find_randomization_implementation_v3(text: str, block_chars: int = 500):
    spans = _token_spans(text)
    blocks = []
    for h in HEAD_RE.finditer(text):
        s = h.end(); e = min(len(text), s + block_chars)
        blocks.append((s, e))
    inside = lambda p: any(s <= p < e for s, e in blocks)
    out = []
    for m in ROLE_RE.finditer(text):
        if inside(m.start()):
            w_s, w_e = _char_to_word((m.start(), m.end()), spans)
            out.append((w_s, w_e, m.group(0)))
    return out

def find_randomization_implementation_v4(text: str, window: int = 8):
    spans = _token_spans(text)
    tokens = [text[s:e] for s, e in spans]
    impl_matches = find_randomization_implementation_v2(text, window=window)
    out = []
    for w_s, w_e, snip in impl_matches:
        roles = {tokens[i].lower() for i in range(max(0, w_s-window), min(len(tokens), w_e+window)) if ROLE_RE.fullmatch(tokens[i])}
        if len(roles) >= 2:
            out.append((w_s, w_e, snip))
    return out

def find_randomization_implementation_v5(text: str):
    return _collect([TIGHT_TEMPLATE_RE], text)

RANDOMIZATION_IMPLEMENTATION_FINDERS: Dict[str, Callable[[str], List[Tuple[int,int,str]]]] = {
    "v1": find_randomization_implementation_v1,
    "v2": find_randomization_implementation_v2,
    "v3": find_randomization_implementation_v3,
    "v4": find_randomization_implementation_v4,
    "v5": find_randomization_implementation_v5,
}

__all__ = ["find_randomization_implementation_v1","find_randomization_implementation_v2","find_randomization_implementation_v3","find_randomization_implementation_v4","find_randomization_implementation_v5","RANDOMIZATION_IMPLEMENTATION_FINDERS"]

find_randomization_implementation_high_recall = find_randomization_implementation_v1
find_randomization_implementation_high_precision = find_randomization_implementation_v5
