"""Exact typed-span evaluation and held-out pooled isotonic calibration.

Gold offsets are Unicode character indices; predicted offsets are UTF-8 bytes.
Predicted-span reliability excludes missed entities; recall is reported separately.
"""
import bisect
import hashlib
import json
from pathlib import Path

LABELS = ('first_name', 'last_name', 'email', 'phone_number', 'street_address', 'customer_id', 'account_number')


def read_jsonl(path):
    with open(path, encoding='utf-8') as f:
        return [json.loads(line) for line in f if line.strip()]


def write_json(path, obj):
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    Path(path).write_text(json.dumps(obj, indent=2, allow_nan=False) + '\n')


def apply_calibration(score, calibration):
    xs, ys = calibration['xs'], calibration['ys']
    if score <= xs[0]:
        return ys[0]
    if score >= xs[-1]:
        return ys[-1]
    i = bisect.bisect_right(xs, score)
    return ys[i-1] + (score-xs[i-1])/(xs[i]-xs[i-1])*(ys[i]-ys[i-1])


def _index(records):
    out = {}
    for r in records:
        if r['id'] in out:
            raise ValueError(f"Duplicate document ID: {r['id']}")
        out[r['id']] = r
    return out


def aligned(gold_records, prediction_records, calibration=None):
    gold, predictions = _index(gold_records), _index(prediction_records)
    if gold.keys() != predictions.keys():
        raise ValueError('Gold/prediction document IDs differ; missing predictions cannot be silently omitted')
    for doc_id, g in gold.items():
        p = predictions[doc_id]
        if g['text'] != p['text']:
            raise ValueError(f'Text mismatch: {doc_id}')
        offsets = [0]
        for ch in g['text']:
            offsets.append(offsets[-1] + len(ch.encode('utf-8')))
        boundaries = set(offsets)
        target = set()
        for span in g['spans']:
            a, b, label = span['start'], span['end'], span['label']
            if label not in LABELS or not (isinstance(a, int) and isinstance(b, int) and 0 <= a < b < len(offsets)):
                raise ValueError(f'Invalid gold span: {doc_id}: {span}')
            target.add((label, offsets[a], offsets[b]))
        found = {}
        for span in p['findings']:
            label, a, b = span['label'], span['start'], span['end']
            raw = float(span['raw_score'])
            if label not in LABELS or a not in boundaries or b not in boundaries or a >= b or not 0 <= raw <= 1:
                raise ValueError(f'Invalid prediction: {doc_id}: {span}')
            score = apply_calibration(raw, calibration) if calibration else span.get('confidence')
            if score is None:
                score = raw
            score = float(score)
            if not 0 <= score <= 1:
                raise ValueError(f'Invalid confidence: {score}')
            key = (label, a, b)
            if key in found:
                raise ValueError(f'Duplicate predicted span: {doc_id}: {key}')
            found[key] = (raw, score)
        yield doc_id, target, found


def _prf(tp, predicted, gold):
    precision = tp / predicted if predicted else 0.0
    recall = tp / gold if gold else 0.0
    return {'tp': tp, 'fp': predicted-tp, 'fn': gold-tp, 'predicted': predicted, 'support': gold,
            'precision': precision, 'recall': recall, 'f1': 2*tp/(predicted+gold) if predicted+gold else 0.0}


def _merge(spans):
    merged = []
    for a, b in sorted((s[1], s[2]) for s in spans):
        if merged and a <= merged[-1][1]:
            merged[-1][1] = max(b, merged[-1][1])
        else:
            merged.append([a, b])
    return merged


def _coverage(gold, pred):
    g, p = _merge(gold), _merge(pred)
    i = j = overlap = 0
    while i < len(g) and j < len(p):
        overlap += max(0, min(g[i][1], p[j][1])-max(g[i][0], p[j][0]))
        if g[i][1] <= p[j][1]:
            i += 1
        else:
            j += 1
    return overlap, sum(b-a for a,b in p), sum(b-a for a,b in g)


def reliability(rows, score_index):
    bins = [{'lower': i/10, 'upper': (i+1)/10, 'support': 0, 'mean_score': 0.0, 'accuracy': 0.0} for i in range(10)]
    squared = 0.0
    for row in rows:
        score, correct = row[score_index], row[2]
        b = bins[min(9, int(score*10))]
        b['support'] += 1
        b['mean_score'] += score
        b['accuracy'] += correct
        squared += (score-correct)**2
    ece = 0.0
    for b in bins:
        if b['support']:
            b['mean_score'] /= b['support']
            b['accuracy'] /= b['support']
            ece += b['support'] * abs(b['mean_score']-b['accuracy'])
        else:
            b['mean_score'] = b['accuracy'] = None
    n = len(rows)
    return {'support': n, 'brier': squared/n if n else None, 'ece': ece/n if n else None,
            'binning': '10 equal-width bins; left-inclusive, right-exclusive except final bin includes 1', 'bins': bins}


def evaluate(gold_records, prediction_records, threshold=None, calibration=None):
    if threshold is None and calibration:
        threshold = calibration['threshold']
    counts = {label: [0, 0, 0] for label in LABELS}
    coverage = [0, 0, 0]
    docs = dict(total=0, positive_support=0, clean_support=0, false_positive=0, false_negative=0,
                with_any_missed_entity=0, with_any_spurious_entity=0)
    rows = []
    for _, target, found in aligned(gold_records, prediction_records, calibration):
        accepted = {k for k, (_, s) in found.items() if threshold is None or s >= threshold}
        for k, (raw, score) in found.items():
            rows.append((raw, score, int(k in target)))
        for label in LABELS:
            gs, ps = {s for s in target if s[0] == label}, {s for s in accepted if s[0] == label}
            for i, value in enumerate((len(gs & ps), len(ps), len(gs))):
                counts[label][i] += value
        for i, value in enumerate(_coverage(target, accepted)):
            coverage[i] += value
        docs['total'] += 1
        docs['positive_support'] += bool(target)
        docs['clean_support'] += not target
        docs['false_positive'] += bool(accepted) and not target
        docs['false_negative'] += bool(target) and not accepted
        docs['with_any_missed_entity'] += bool(target-accepted)
        docs['with_any_spurious_entity'] += bool(accepted-target)
    return {'threshold': threshold, 'overall': _prf(*(sum(c[i] for c in counts.values()) for i in range(3))),
            'per_label': {label: _prf(*c) for label, c in counts.items()},
            'redaction_byte_coverage': _prf(*coverage), 'documents': docs,
            'raw_reliability': reliability(rows, 0), 'confidence_reliability': reliability(rows, 1),
            'reliability_population': 'All predicted spans before threshold, correctness = exact typed span. Missed gold entities excluded; see recall.',
            'confidence_semantics': 'Pooled calibrated exact-span score' if calibration else 'Provided confidence when non-null; otherwise uncalibrated raw span score',
            'coverage_semantics': 'Union of UTF-8 bytes; ignores labels. Zero-denominator P/R/F1 reported as 0.'}


def fit_calibration(fit_gold, fit_predictions, select_gold, select_predictions):
    fit_gold, select_gold = list(fit_gold), list(select_gold)
    if {r['id'] for r in fit_gold} & {r['id'] for r in select_gold}:
        raise ValueError('Calibration fit/select document IDs overlap')
    hashes = lambda records: {hashlib.sha256(r['text'].encode()).hexdigest() for r in records}
    if hashes(fit_gold) & hashes(select_gold):
        raise ValueError('Calibration fit/select texts overlap')
    fit_rows = [(raw, int(k in target)) for _, target, found in aligned(fit_gold, fit_predictions)
                for k, (raw, _) in found.items()]
    if not fit_rows:
        raise ValueError('No predicted spans available to fit calibration')
    from sklearn.isotonic import IsotonicRegression
    model = IsotonicRegression(out_of_bounds='clip').fit([r[0] for r in fit_rows], [r[1] for r in fit_rows])
    artifact = {'version': 1, 'method': 'isotonic', 'xs': model.X_thresholds_.tolist(), 'ys': model.y_thresholds_.tolist(),
                'interpolation': 'linear; clip outside fitted x range', 'score_semantics': 'Estimated exact typed-span correctness for predicted spans; pooled labels; raw score = minimum predicted-class token softmax',
                'experimental': True, 'fit_documents': len(fit_gold), 'fit_predicted_spans': len(fit_rows),
                'fit_correct_spans': sum(r[1] for r in fit_rows), 'select_documents': len(select_gold),
                'warnings': []}
    if len({r[1] for r in fit_rows}) < 2:
        artifact['warnings'].append('Fit targets contain only one outcome; constant map cannot demonstrate probability reliability.')
    select_rows, support = [], 0
    for _, target, found in aligned(select_gold, select_predictions, artifact):
        support += len(target)
        select_rows.extend((score, int(k in target)) for k, (_, score) in found.items())
    if not select_rows or not support:
        raise ValueError('Threshold selection requires predicted spans and gold entities')
    select_rows.sort(reverse=True)
    curve = []
    tp = predicted = 0
    i = 0
    while i < len(select_rows):
        threshold = select_rows[i][0]
        while i < len(select_rows) and select_rows[i][0] == threshold:
            tp += select_rows[i][1]
            predicted += 1
            i += 1
        curve.append({'threshold': threshold, **_prf(tp, predicted, support)})
    best = max(curve, key=lambda row: (row['f1'], row['recall'], row['threshold']))
    artifact.update(threshold=best['threshold'], threshold_objective='Max exact typed-span F1 on calibration-select; ties favor recall, then higher threshold',
                    select_operating_point=best, threshold_curve=curve)
    return artifact