"""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/contextual_rl_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_contextual_work.py',
        ROOT / 'notes/contextual_reporting_plan_2239.md')}
    records = {}
    before_reference = None
    for arm in contract['arms']:
        arm_folder = folder / arm
        updates, rollouts = [], []
        for chunk in range(contract['chunks']):
            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) == 64 and len(rollouts) == 2048
        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']
            diagnostic = update['diagnostic']
            if diagnostic is not None:
                assert diagnostic['actual_gradient_restoration_exact']
                components = diagnostic['components']
                rl_norm = components['rl']['weighted_gradient_norm']
                auxiliary = 'dense' if arm in ('C', 'E') else 'replay' if arm == 'H' else None
                auxiliary_norm = components[auxiliary]['weighted_gradient_norm'] if auxiliary else None
                diagnostics.append({
                    'step': update['step'], 'components': components,
                    'auxiliary_to_rl_norm_ratio': auxiliary_norm / rl_norm if auxiliary_norm is not None and rl_norm > 0 else None,
                    'combined_preclip_norm': update['preclip_norm'],
                })
        assert [d['step'] for d in diagnostics] == contract['gradient_diagnostic_updates']
        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), 'gradient_diagnostics': diagnostics,
            '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_contextual_work.py')], check=True, stdout=subprocess.DEVNULL)
    work = read_json(ROOT / 'notes/contextual_rl_work_accounting.json')
    output = {
        'scope': 'Descriptive adjunct specified during G training, before native generation. No additional gates or inferential claims. Two gradient snapshots cannot characterize the entire trajectory. 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 = ['# Contextual dense RL: 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 /64 | Donor groups /256 | Receiver draws /2048 | Dense tokens | Clipped batches /64 | Source NLL change |', '|---|---:|---:|---:|---:|---:|---:|']
    for arm, record in records.items():
        e, r = record['exposure'], record['retention']
        lines.append(f"| {arm} | {result['training'][arm]['optimizer_updates']} | {e['donor_groups']} | {e['eligible_receiver_draws']} | {e['dense_target_tokens']} | {e['clipped_batches']} | {r['mean_paired_delta_nll']:+.6f} |")
    lines += ['', '| Arm | Snapshot | RL norm | Auxiliary / RL norm | Combined norm |', '|---|---:|---:|---:|---:|']
    for arm, record in records.items():
        for d in record['gradient_diagnostics']:
            ratio = d['auxiliary_to_rl_norm_ratio']
            ratio_text = 'n/a' if ratio is None else f'{ratio:.6f}'
            lines.append(f"| {arm} | {d['step']} | {d['components']['rl']['weighted_gradient_norm']:.4f} | {ratio_text} | {d['combined_preclip_norm']:.4f} |")
    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. The interval reflects missing interruption exit timing, not statistical uncertainty. Readiness, probes and shared ancestry remain additional.',
              '', '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. C−E/C−H determine contextual/dense specificity; this report does not infer such specificity from the spending gate alone. 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()
