"""Inference-only source candidates and explicitly labelled retrieval baselines."""
import re
from solomon.evidence import digest,rank_candidates,validate_spans

def candidates(text):
    """Source sentences retaining exact character ranges; no authoring labels."""
    ends=[m.end() for m in re.finditer(r'(?<=[.!?])(?:[ \t]+|\n+)|\n\s*\n',text)]+[len(text)]
    result=[];start=0
    for end in ends:
        if end>start and text[start:end].strip():
            result.append({'id':f'text:{start}:{end}','kind':'text','start':start,'end':end,
                           'text':text[start:end],'source_sha256':digest(text)})
        start=end
    return result

def governing_context(text,spans=None):
    """Generic rule/retraction interpretation context, independent of question gold."""
    spans=candidates(text) if spans is None else spans
    pattern=re.compile(r'general rule|both limbs|same (?:person|courier|member|officer)|neither limb|missing requirement|silence|withdrawn passage|dates and seniority|a withdrawn|outside these|bands are|order of the levels|no.*precedence|each passage|a grant in force|an exception to a rule',re.I)
    return [s for s in spans if pattern.search(s['text'])]

def merge(text,spans):
    intervals=[]
    for span in sorted(spans,key=lambda s:(s['start'],s['end'])):
        a,b=span['start'],span['end']
        if intervals and a<=intervals[-1][1]:intervals[-1]=(intervals[-1][0],max(b,intervals[-1][1]))
        else:intervals.append((a,b))
    result=[{'id':f'text:{a}:{b}','kind':'text','start':a,'end':b,'text':text[a:b],'source_sha256':digest(text)} for a,b in intervals]
    return validate_spans(text,result)

def render_subset(text,spans,*,include_context=True):
    selected=merge(text,list(spans)+(governing_context(text) if include_context else []))
    return '\n\n'.join(s['text'] for s in selected),selected

_WORD=re.compile(r'\w+')
_SCALE=re.compile(r'\bbands?\b|\bscale\b|from lowest to highest|order of the levels',re.I)
POOL_V2={'name':'idf-rare-rule-v2','idf_limit':24,'rare_df':16,'rules':True}

def _words(s):return set(_WORD.findall(s.casefold()))

def rule_candidates(text,spans=None):
    """Question-independent rule, exception, withdrawal-convention and scale passages."""
    spans=candidates(text) if spans is None else spans
    rules={s['id'] for s in governing_context(text,spans)}
    return [s for s in spans if s['id'] in rules or _SCALE.search(s['text'])]

def candidate_pool(question,spans,text=None,*,idf_limit=24,rare_df=16,rules=True,**_):
    """Bounded source-only pool: document-IDF top-k, every passage sharing a rare
    question term (typically the subject's name), plus question-independent rules.

    Candidate order is source order. candidate_score stays the plain lexical-overlap
    fraction used by the frozen relevance-head feature, so heads remain comparable.
    Gold is never an input.
    """
    import math
    if idf_limit<1 or rare_df<0:raise ValueError('invalid candidate pool policy')
    ranked=rank_candidates(question,spans,len(spans)) if spans else []
    words=[_words(c['text']) for c in ranked];q=_words(question);n=len(ranked)
    df={}
    for ws in words:
        for w in ws:df[w]=df.get(w,0)+1
    idf=[sum(math.log((n+1)/(df[w]+.5)) for w in ws&q) for ws in words]
    order=sorted(range(n),key=lambda i:(-idf[i],ranked[i]['start']))
    reason={}
    for i in order[:idf_limit]:reason.setdefault(ranked[i]['id'],'idf')
    rare={w for w in q if df.get(w,0)<=rare_df}
    for c,ws in zip(ranked,words):
        if ws&rare:reason.setdefault(c['id'],'rare_term')
    if rules:
        for c in rule_candidates(text or '',ranked):reason.setdefault(c['id'],'rule')
    return sorted(({**c,'pool_reason':reason[c['id']]} for c in ranked if c['id'] in reason),key=lambda c:c['start'])

_WITHDRAW=re.compile(r'withdr[ae]w|withdrawn|take back|disregard|rescind|retract|should not be relied|there is substituted|is deleted|expressly delete',re.I)
_REFERENCE=re.compile(r'(?:message of|wrote on|message dated) (\d{1,2} [A-Z][a-z]+)|minute (\d+)|[Ee]ntry (\d+) of [Ss]chedule (\d+)|clause ([\d.]+)|passage in my message')
_EXCEPTION=re.compile(r'does not apply to|subject to the exception|is an exception|except (?:where|that|for)\b',re.I)
_STOP={'this','that','these','those','with','under','which','what','does','file','correspondence','record','records','recorded',
       'minutes','agreement','bundle','messages','message','stand','stands','taking','reading','whole','strength','position',
       'open','given','have','been','applies','apply','entitled','allowed','liberty','free','from','there','their','they',
       'decisions','here','schedules','schedule','papers','office','shown','show','shows','case','matters','terms','place','placed'}
DEPENDENCIES_V1={'name':'withdrawal-exception-v1','withdrawals':True,'exceptions':True,'prune_withdrawn':False,'rare_df':16}

def _topic(words,exclude):
    return {w for w in words if len(w)>=4 and w not in _STOP and w not in exclude and not w.isdigit()}

def _withdrawn_block(text,passage):
    """Source range that a withdrawal refers to (dated message or numbered minute); None if not resolvable."""
    m=_REFERENCE.search(passage)
    if not m:return None
    if m.group(1):head=re.search(r'Message \d+\. '+re.escape(m.group(1))+r'\.',text)
    elif m.group(2):head=re.search(r'(?:^|\n)'+m.group(2)+r'\. ',text)
    else:return None
    if not head:return None
    end=text.find('\n\n',head.end());return head.start(),(len(text) if end<0 else end)

def dependency_expand(question,spans,selected_ids,text,*,withdrawals=True,exceptions=True,prune_withdrawn=False,rare_df=16,**_):
    """Complete a unit's selected evidence with its withdrawal/exception dependencies.

    Source text and the already-selected passages only; never gold. Units with no
    selection are unchanged, so no-positive-support decisions are preserved.
    - withdrawal: a passage with a withdrawal verb AND an explicit reference (message/minute/
      entry/clause) that names the question's subject (a rare capitalised name, also present in a
      selected passage) and its most specific in-source topic word;
    - exception: an exception clause immediately following a selected passage;
    - prune_withdrawn (optional): drop selected passages inside the dated message / numbered
      minute a kept withdrawal refers to, when they share its subject and a topic word.
    Returns {'keep','added','pruned'} as candidate id lists in source order.
    """
    selected=[s for s in spans if s['id'] in set(selected_ids)]
    if not selected:return {'keep':[],'added':[],'pruned':[]}
    words=[_words(s['text']) for s in spans];df={}
    for ws in words:
        for w in ws:df[w]=df.get(w,0)+1
    # Subject: capitalised non-initial question tokens (names) that are rare in the source.
    names=re.findall(r'(?<!^)(?<![.?!] )\b[A-Z][\w&\'-]*',question.strip())
    subject={w for n in names for w in _words(n) if len(w)>=3 and df.get(w,0)<=rare_df}
    q=_words(question);present=[w for w in _topic(q,subject) if df.get(w,0)]
    # The withdrawal must name the question's most specific in-source topic word (e.g. the attribute).
    topic={min(present,key=lambda w:(df[w],w))} if present else set()
    chosen={s['id'] for s in selected};sel_words=set().union(*(_words(s['text']) for s in selected));added=[]
    for i,(s,ws) in enumerate(zip(spans,words)):
        if s['id'] in chosen:continue
        if (withdrawals and _WITHDRAW.search(s['text']) and _REFERENCE.search(s['text'])
                and ws&subject&sel_words and ws&topic):added.append(s['id']);continue
        if exceptions and i>0 and spans[i-1]['id'] in chosen and _EXCEPTION.search(s['text']):added.append(s['id'])
    keep=chosen|set(added);pruned=[]
    if prune_withdrawn:
        for s,ws in zip(spans,words):
            if s['id'] not in keep or not _WITHDRAW.search(s['text']):continue
            block=_withdrawn_block(text,s['text'])
            if block is None:continue
            key=ws&subject;about=_topic(ws,subject)
            for t,tw in zip(spans,words):
                if (t['id'] in keep and t['id']!=s['id'] and block[0]<=t['start'] and t['end']<=block[1]+2
                        and tw&key and tw&about and not _WITHDRAW.search(t['text'])):pruned.append(t['id'])
    order=[s['id'] for s in spans];pruned=set(pruned)
    return {'keep':[i for i in order if i in keep and i not in pruned],'added':[i for i in order if i in set(added)],
            'pruned':[i for i in order if i in pruned]}

def lexical_select(text,questions,limit=4):
    source=candidates(text);chosen={}
    for q in questions:
        for candidate in rank_candidates(q,source,limit):
            if candidate['candidate_score']>0:chosen[candidate['id']]=candidate
    return {'evidence':sorted(chosen.values(),key=lambda c:c['start']),
            'method':'lexical_overlap','verification':'retrieval_only','faithfulness_established':False}

def remove(text,spans):
    spans=merge(text,spans);cursor=0;parts=[]
    for s in spans:parts.append(text[cursor:s['start']]);cursor=s['end']
    parts.append(text[cursor:]);return ''.join(parts)