"""Describe mechanically split own-trace targets; never read native outcomes.

No hidden features or training target files are emitted. The split routine is
development code pending a new experiment contract.
"""
import hashlib
import json
from pathlib import Path

import numpy as np

from run_baseline import Qwen3Tokenizer, extract_final_candidate

ROOT = Path(__file__).resolve().parents[1]


def split_trace(row, tokenizer):
    tokens = row['completion']
    assert tokens[-1] == tokenizer.eos_token_id and tokenizer.eos_token_id not in tokens[:-1]
    text = tokenizer.decode(tokens[:-1])
    box = text.rfind('\\boxed')
    assert box >= 0
    cut = text.rfind('\n\n', 0, box)
    if cut < 1:
        raise ValueError('no nonempty prefix before final boxed paragraph')
    prefix, separator, suffix = text[:cut], text[cut:cut+2], text[cut+2:]
    assert prefix + separator + suffix == text
    assert separator == '\n\n' and prefix and suffix
    prediction = extract_final_candidate(text)
    assert prediction and extract_final_candidate(suffix) == prediction
    # This is a declared retokenization of the exact own-policy decoded text.
    # Original full-trace CE keeps the original completion IDs.
    prefix_ids = tokenizer.encode(prefix)
    separator_ids = tokenizer.encode(separator)
    suffix_ids = tokenizer.encode(suffix) + [tokenizer.eos_token_id]
    canonical = '\\boxed{' + prediction + '}'
    canonical_ids = tokenizer.encode(canonical) + [tokenizer.eos_token_id]
    assert extract_final_candidate(canonical) == prediction
    assert tokenizer.decode(canonical_ids[:-1]) == canonical
    encoding = tokenizer._tok.encode(canonical)
    assert encoding.ids == canonical_ids[:-1]
    body_start, body_end = len('\\boxed{'), len('\\boxed{') + len(prediction)
    content_mask = [start < body_end and end > body_start for start, end in encoding.offsets]
    mixed_content_tokens = sum(keep and (start < body_start or end > body_end)
                               for keep, (start, end) in zip(content_mask, encoding.offsets))
    content_mask.append(False)  # EOS is formatting/termination, not answer body.
    assert any(content_mask) and len(content_mask) == len(canonical_ids)
    assert tokenizer.decode(prefix_ids) == prefix
    assert tokenizer.decode(separator_ids) == separator
    assert tokenizer.decode(suffix_ids[:-1]) == suffix
    assert tokenizer.decode(prefix_ids + separator_ids + suffix_ids[:-1]) == text
    return {'prefix': prefix_ids, 'separator': separator_ids, 'suffix': suffix_ids,
            'canonical_answer': canonical_ids,
            'canonical_content_mask': content_mask,
            'mixed_content_tokens': mixed_content_tokens,
            'prefix_boxes': prefix.count('\\boxed'),
            'retokenized_ids_differ': prefix_ids + separator_ids + suffix_ids != tokens}


def describe(values):
    return {'n': len(values), 'mean': float(np.mean(values)),
            'min_median_max': np.quantile(values, [0, .5, 1]).tolist()}


def main():
    tokenizer = Qwen3Tokenizer(tokenizer_file_path=ROOT / 'code/data/tokenizer-base.json')
    contract = json.loads((ROOT / 'notes/trainable_recurrence_contract.json').read_text())
    result = {'status': 'PASS', 'scope': 'source-only boundary audit; no native readout or target files',
              'groups': {}}
    for field in ('source_targets', 'holdout_targets'):
        path = ROOT / contract[field]
        rows = [json.loads(line) for line in path.read_text().splitlines()]
        splits = [split_trace(row, tokenizer) for row in rows]
        result['groups'][field] = {
            'source_sha256': hashlib.sha256(path.read_bytes()).hexdigest(),
            'questions': len(rows), 'all_reconstruct_exact_text': True,
            'all_preserve_extracted_answer': True,
            'prefix_tokens': describe([len(s['prefix']) for s in splits]),
            'suffix_tokens_with_eos': describe([len(s['suffix']) for s in splits]),
            'canonical_answer_tokens_with_eos': describe([len(s['canonical_answer']) for s in splits]),
            'answer_body_overlapping_tokens': describe([sum(s['canonical_content_mask']) for s in splits]),
            'mixed_content_and_format_tokens': sum(s['mixed_content_tokens'] for s in splits),
            'separator_ids': sorted({tuple(s['separator']) for s in splits}),
            'earlier_box_rows': sum(s['prefix_boxes'] > 0 for s in splits),
            'retokenized_ids_differ': sum(s['retokenized_ids_differ'] for s in splits),
        }
    result['limits'] = ['Final-paragraph boundary is not a reasoning-step annotation.',
                        'Earlier answer mentions and intermediate correctness were not assessed.',
                        'Source terminal verification is inherited from the prior independent audit.',
                        'No source rows filtered or hand-selected; no native performance inferred.']
    (ROOT / 'notes/latent_source_boundary_audit.json').write_text(json.dumps(result, indent=2)+'\n')
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()
