#!/usr/bin/env python3
"""CPU-only controlled filtering experiment; no LLM or semantic encoder used."""
import hashlib
import json
import math
import platform
import random
import re
from pathlib import Path

SEED = 65
ROOT = Path(__file__).resolve().parent


def tokens(s):
    return re.findall(r"\w+", s.lower())


def features(row):
    ts = tokens(row['response'])
    return [1.0, float(bool(ts)), min(len(ts), 40) / 40,
            len(set(ts)) / max(1, len(ts))]


def probability(row, weights):
    z = sum(a * b for a, b in zip(features(row), weights))
    return 1 / (1 + math.exp(-max(-40, min(40, z))))


def fit(rows):
    weights = [0.0] * 4
    losses = []
    for step in range(1000):
        grad = [0.0] * 4
        loss = 0.0
        for row in rows:
            p = probability(row, weights)
            y = row['label']
            loss -= y * math.log(max(p, 1e-12)) + (1-y)*math.log(max(1-p, 1e-12))
            for j, x in enumerate(features(row)):
                grad[j] += (p-y)*x
        if step in (0, 999):
            losses.append(loss / len(rows))
        weights = [w - 0.3 * (g/len(rows) + (0.001*w if j else 0))
                   for j, (w, g) in enumerate(zip(weights, grad))]
    return weights, losses


def labeled_family(n, split):
    instruction = 'Explain how to add %d and %d.' % (n, n+2)
    answers = [
        ('To add these two integers, start with %d and increase it by %d. '
         'The resulting total is %d; direct counting gives the same answer.' % (n, n+2, 2*n+2), 1),
        ('', 0), ('answer ' * 28, 0)]
    return [dict(id='%s-%d-%d' % (split, n, k), family='seed-%d' % n,
                 instruction=instruction, response=a, label=y, split=split)
            for k, (a, y) in enumerate(answers)]


def expected(row):
    op, a = row['operation'], row['args']
    if op == 'add':
        return sum(a)
    if op == 'sort':
        return sorted(a)
    if op == 'reverse':
        return a[::-1]
    if op == 'sumsq':
        return sum(x*x for x in a)
    raise ValueError(op)


def fixture_pool():
    rows = []
    for op, arguments in [
        ('add', [[i, i+3] for i in range(20, 26)]),
        ('sort', [[9, 2, 7], [8, 1, 4]]),
        ('reverse', ['laboratory', 'research']),
        ('sumsq', [[2, 3, 4], [3, 4, 5]])]:
        for i, args in enumerate(arguments):
            row = dict(id='%s-%d' % (op, i), family=op, operation=op, args=args,
                       instruction='%s: %s' % (op, json.dumps(args)), split='pool',
                       difficulty=3 if op == 'sumsq' else 1)
            value = expected(row)
            if op == 'add':
                response = ('To solve the addition task carefully, combine the first quantity with '
                            'the second quantity and count their total. The requested result is %s, '
                            'which follows directly from the integer addition operation.' % value)
            elif op == 'sort':
                response = 'Ascending order of the supplied integers is %s.' % value
            elif op == 'reverse':
                response = 'Reading from right to left gives %s.' % value
            else:
                response = 'Squared terms sum to %s.' % value
            row.update(response=response, answer=value)
            rows.append(row)
    forged = dict(rows[0], id='forged', answer=444,
                  response=rows[0]['response'].replace(str(rows[0]['answer']), '444'))
    empty = dict(rows[1], id='empty', response='', answer=None)
    return rows + [forged, empty]


def validity(row):
    if not row['response'].strip():
        return 'empty_response'
    if row['answer'] != expected(row):
        return 'wrong_structured_answer'
    return 'ok'


def cosine(a, b):
    na, nb = math.sqrt(sum(x*x for x in a)), math.sqrt(sum(x*x for x in b))
    if not na or not nb:
        raise ValueError('zero vector')
    return sum(x*y for x, y in zip(a, b))/(na*nb)


def diverse(ranked, vectors, threshold):
    chosen = []
    audit = []
    for r in ranked:
        sim = max((cosine(vectors[r['id']], vectors[s['id']]) for s in chosen), default=None)
        keep = sim is None or sim <= threshold
        audit.append(dict(id=r['id'], max_similarity=sim, kept=keep))
        if keep:
            chosen.append(r)
    return chosen, audit


def metrics(rows, pool, budget):
    groups = {r['family'] for r in pool if validity(r) == 'ok'}
    return dict(ids=[r['id'] for r in rows], count=len(rows), budget=budget,
                mean_quality=round(sum(r['quality'] for r in rows)/max(1, len(rows)), 6),
                valid_fraction=sum(validity(r)=='ok' for r in rows)/max(1, len(rows)),
                family_coverage=len({r['family'] for r in rows})/len(groups),
                hard_count=sum(r['difficulty']>=3 for r in rows),
                response_words=sum(len(tokens(r['response'])) for r in rows))


def main():
    rng = random.Random(SEED)
    train = [r for n in range(2, 10) for r in labeled_family(n, 'train')]
    dev = [r for n in range(12, 14) for r in labeled_family(n, 'dev')]
    assert {r['family'] for r in train}.isdisjoint(r['family'] for r in dev)
    weights, losses = fit(train)
    grid = []
    for t in (0.3, 0.4, 0.5, 0.6, 0.7):
        tp = sum(probability(r, weights)>=t and r['label']==1 for r in dev)
        fp = sum(probability(r, weights)>=t and r['label']==0 for r in dev)
        fn = sum(probability(r, weights)<t and r['label']==1 for r in dev)
        grid.append(dict(threshold=t, tp=tp, fp=fp, fn=fn,
                         f1=2*tp/max(1, 2*tp+fp+fn)))
    threshold = max(grid, key=lambda z: (z['f1'], -abs(z['threshold']-0.5)))['threshold']
    pool = fixture_pool()
    for r in pool:
        r['quality'] = probability(r, weights)
        r['validity'] = validity(r)
    # Pool gold/validity is never used to fit the classifier or choose its threshold.
    families = ['add', 'sort', 'reverse', 'sumsq']
    vectors = {r['id']: [float(r['family']==g) for g in families] for r in pool}
    eligible = [r for r in pool if r['validity']=='ok' and r['quality']>=threshold]
    ranked = sorted(eligible, key=lambda r: (-r['quality'], r['id']))
    representatives, trace = diverse(ranked, vectors, 0.9)
    budget = 3
    # Reserve one already-verified hard representative, then fill in score order.
    hard = [r for r in representatives if r['difficulty']>=3][:1]
    quota = (hard + [r for r in representatives if r not in hard])[:budget]
    strategies = {
        'random_valid': rng.sample(eligible, min(budget, len(eligible))),
        'quality_only': ranked[:budget],
        'quality_diversity': representatives[:budget],
        'quality_diversity_hard': quota,
        'unsafe_quality_no_validity': sorted(pool, key=lambda r: (-r['quality'], r['id']))[:budget]}
    report = {k: metrics(v, pool, budget) for k, v in strategies.items()}
    a, b = next(r for r in pool if r['id']=='add-0'), next(r for r in pool if r['id']=='forged')
    assert features(a) == features(b) and a['quality'] == b['quality']
    assert validity(a)=='ok' and validity(b)=='wrong_structured_answer'
    assert diverse([], {}, 0.9)[0] == []
    assert len(diverse([a, b], vectors, 0.9)[0]) == 1
    assert len(diverse([a, b], vectors, 1.0)[0]) == 2  # strict > rejection
    try:
        cosine([0.0]*4, [1.0]*4)
    except ValueError:
        pass
    else:
        raise AssertionError('zero vector must be rejected')
    assert losses[-1] < losses[0]
    assert all(validity(r)=='ok' for r in quota)
    payload = dict(seed=SEED, python=platform.python_version(),
                   feature_names=['bias', 'nonempty', 'word_count_capped_40', 'unique_ratio'],
                   train_shape=[len(train), 4], dev_shape=[len(dev), 4],
                   pool_shape=[len(pool), 4], vector_type='handwritten task one-hot, NOT neural embeddings',
                   weights=weights, training_bce=losses, threshold_grid=grid,
                   selected_threshold=threshold, cosine_threshold=0.9, budget=budget,
                   eligible_count=len(eligible), train=train, dev=dev, pool=pool,
                   diversity_trace=trace, results=report,
                   checks=['split family disjoint', 'loss decreases', 'wrong answer feature collision',
                           'empty pool', 'strict similarity boundary', 'zero vector rejected', 'quota validity'],
                   limitations=['Controlled shared templates; dev is not independent-domain validation',
                                'Only structured answers checked; prose correctness not verified',
                                'Difficulty is handwritten, not measured student loss',
                                'No LLM generation, neural scorer, embedding or downstream SFT was run'])
    payload['script_sha256'] = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
    (ROOT/'results.json').write_text(json.dumps(payload, ensure_ascii=False, indent=2)+'\n')
    print('Python', platform.python_version(), 'seed', SEED)
    print('train/dev/pool shapes', payload['train_shape'], payload['dev_shape'], payload['pool_shape'])
    print('BCE checkpoints', losses, 'threshold', threshold, 'eligible', len(eligible))
    print(json.dumps(report, indent=2))
    print('SMOKE PASS:', ', '.join(payload['checks']))


if __name__ == '__main__':
    main()