"""Descriptive integrated report, available only after the frozen full audit.

Does not alter the contract, recompute the primary gate, or select a successor.
"""
import collections
import hashlib
import json
import pathlib
import statistics
import subprocess
import sys

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


def read_json(path):
    return json.loads(path.read_text())


def rows(path):
    return [json.loads(line) for line in path.read_text().splitlines()]


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main():
    contract_path = ROOT / 'notes/stable_contextual_contract.json'
    contract = read_json(contract_path)
    folder = ROOT / contract['output']
    result_path = folder / 'analysis/result.json'
    queue = read_json(ROOT / contract['queue_state'])
    if queue['status'] != 'COMPLETE' or not result_path.exists():
        raise SystemExit('Full frozen audit is not complete; descriptive reporting deferred.')
    result = read_json(result_path)
    assert result['status'] == 'COMPLETE_AUDITED'
    split = read_json(ROOT / contract['split'])
    manifest = {str(path.relative_to(ROOT)): digest(path) for path in (
        result_path, contract_path, pathlib.Path(__file__).resolve(),
        ROOT / 'code/account_stable_contextual_work.py',
        ROOT / 'notes/stable_contextual_reporting_plan.md')}
    records = {}
    before_reference = None
    for arm in contract['arms']:
        arm_folder = folder / arm
        updates, rollouts = [], []
        for chunk in range(contract['chunks'][arm]):
            summary = read_json(arm_folder / f'chunk{chunk}_summary.json')
            assert summary['contract_sha256'] == digest(contract_path)
            for kind, key, target in [('updates', 'updates_sha256', updates),
                                      ('rollouts', 'raw_sha256', rollouts)]:
                path = arm_folder / f'chunk{chunk}_{kind}.jsonl'
                assert digest(path) == summary[key]
                manifest[str(path.relative_to(ROOT))] = summary[key]
                target.extend(rows(path))
        assert len(updates) == contract['updates'][arm] and len(rollouts) == contract['updates'][arm]*32
        grouped = collections.defaultdict(list)
        for row in rollouts:
            grouped[row['step'], row['position']].append(row)
        exposure = collections.Counter()
        dense_divergence = 0.0
        diagnostics = []
        for update in updates:
            exposure['zero_signal_batches'] += sum(g['backward_calls'] for g in update['groups']) == 0
            exposure['clipped_batches'] += update['preclip_norm'] > contract['clip_norm']
            for position, group in enumerate(update['groups']):
                samples = grouped[update['step'], position]
                assert len(samples) == 8
                successes = sum(row['reward'] for row in samples)
                exposure['all_wrong_groups'] += successes == 0
                exposure['all_correct_groups'] += successes == 8
                exposure['mixed_reward_groups'] += 0 < successes < 8
                exposure['correct_capped_draws'] += sum(row['reward'] == 1 and not row['eos'] for row in samples)
                exposure['wrong_capped_draws'] += sum(row['reward'] == 0 and not row['eos'] for row in samples)
                exposure['donor_groups'] += group['donor'] is not None
                exposure['eligible_receiver_groups'] += bool(group['receivers'])
                exposure['eligible_receiver_draws'] += len(group['receivers'])
                exposure['rl_active_draws'] += sum(a != 0 for a in group['advantages'])
                exposure['dense_target_tokens'] += group['dense_tokens']
                dense_divergence += group['dense_jsd_sum']
        assert all(summary['teacher_exact_unchanged'] for summary in [read_json(arm_folder / f'chunk{j}_summary.json') for j in range(contract['chunks'][arm])])
        retention = {}
        for phase in ('before', 'after'):
            path = arm_folder / f'retention_{phase}.json'
            manifest[str(path.relative_to(ROOT))] = digest(path)
            data = read_json(path)
            retention[phase] = {r['index']: r['nll'] for r in data['rows']}
            assert len(retention[phase]) == 64
            assert not set(retention[phase]) & set(split['train_indices'])
        assert retention['before'].keys() == retention['after'].keys()
        if before_reference is None:
            before_reference = retention['before']
        # Identical initial weights/interface; tolerate bounded GPU execution.
        assert max(abs(before_reference[i] - retention['before'][i]) for i in before_reference) < 1e-5
        deltas = [retention['after'][i] - v for i, v in retention['before'].items()]
        records[arm] = {
            'exposure': dict(exposure), 'scheduled_batches': contract['updates'][arm],
            'questions': contract['updates'][arm]*4, 'draws': contract['updates'][arm]*32,
            'mean_jsd_per_dense_token': dense_divergence / exposure['dense_target_tokens'] if exposure['dense_target_tokens'] else None,
            'median_preclip_norm': statistics.median(u['preclip_norm'] for u in updates),
            'retention': {
                'before_mean_nll': statistics.mean(retention['before'].values()),
                'after_mean_nll': statistics.mean(retention['after'].values()),
                'mean_paired_delta_nll': statistics.mean(deltas),
                'questions_with_increased_nll': sum(d > 0 for d in deltas),
                'questions': 64,
            },
        }
    subprocess.run([sys.executable, str(ROOT / 'code/account_stable_contextual_work.py')], check=True, stdout=subprocess.DEVNULL)
    work = read_json(ROOT / 'notes/stable_contextual_work_accounting.json')
    output = {
        'scope': 'Descriptive adjunct prepared while G was queued, before native generation. No additional gates or inferential claims. Combined gradients do not measure component dose or Adam displacement. Retention is historical source likelihood, not native capability. Exposure is treatment-dependent, not matched causal strata.',
        'contract_sha256': digest(contract_path), 'primary': result,
        'training_descriptive': records, 'work': work, 'input_sha256': manifest,
    }
    destination = folder / 'analysis/descriptive_report.json'
    destination.write_text(json.dumps(output, indent=2))
    lines = ['# Frozen contextual teacher: integrated development readout', '', output['scope'], '',
             '| Policy | Correct / 512 |', '|---|---:|']
    lines += [f'| {a} | {n} |' for a, n in result['counts'].items()]
    lines += ['', f"Frozen development spending gate: **{result['development_gate']}**.", '',
              '| Contrast | Net correct | Difference, pp | Paired 95% interval, pp |', '|---|---:|---:|---|']
    for label, contrast in result['contrasts'].items():
        lower, upper = contrast['ci95']
        lines.append(f"| {label} | {contrast['net_correct']} | {100*contrast['difference']:.2f} | [{100*lower:.2f}, {100*upper:.2f}] |")
    lines += ['', '| Arm | Optimizer updates / scheduled | Donor groups / questions | Receiver draws / draws | Dense tokens | Clipped batches / scheduled | Source NLL change |', '|---|---:|---:|---:|---:|---:|---:|']
    for arm, record in records.items():
        e, r = record['exposure'], record['retention']
        lines.append(f"| {arm} | {result['training'][arm]['optimizer_updates']}/{record['scheduled_batches']} | {e['donor_groups']}/{record['questions']} | {e['eligible_receiver_draws']}/{record['draws']} | {e['dense_target_tokens']} | {e['clipped_batches']}/{record['scheduled_batches']} | {r['mean_paired_delta_nll']:+.6f} |")
    lines += ['', '| Arm | Training allocated GPU seconds | Matrix + attention proxy, FLOPs |', '|---|---:|---:|']
    for arm, accounting in work['training'].items():
        proxy = accounting['proxy']
        flops = proxy['matrix_flops_proxy'] + proxy['dense_attention_flops_proxy']
        lines.append(f"| {arm} | {accounting['allocated_gpu_seconds']:.1f} | {flops:.4e} |")
    lines += ['', '| Policy | Evaluation GPU seconds | Generated tokens |', '|---|---:|---:|']
    for arm, accounting in work['evaluation'].items():
        lines.append(f"| {arm} | {accounting['allocated_gpu_seconds']:.1f} | {accounting['work']['generated_tokens']} |")
    lower, upper = work['total_released_plus_interrupted_gpu_seconds_bounds']
    lines += ['', f'Total training/evaluation allocation including interrupted work: [{lower:.2f}, {upper:.2f}] GPU seconds. Any nonzero interval width reflects missing interruption exit timing, not statistical uncertainty. Readiness57.0s and freshWarm460.1s remain additional, as do older research ancestry costs.',
              '', 'Interrupted attempts:', '']
    for interruption in work['interrupted_attempts']:
        lines.append(f"- {interruption['record']}: {interruption['retained_partial_draws']} retained draws, {interruption['completed_but_uncheckpointed_updates']} uncheckpointed updates, {interruption['retained_generated_tokens_lower_bound']} retained generated tokens. This work was repeated and remains charged.")
    lines += ['', work['scope'], '', result['scope'], '',
              'Primary claims follow the unchanged frozen contrasts. F−E is the fixed contextual-versus-empty contrast. Frozen teacher identity does not make visited prefixes or coarsened target distributions fixed. No current/stable interaction or semantic-content-isolation claim follows. Shared ancestry and source costs are additional, as recorded in the JSON ledger.']
    (folder / 'analysis/descriptive_report.md').write_text('\n'.join(lines) + '\n')
    print(destination.relative_to(ROOT))


if __name__ == '__main__':
    main()
