"""Frozen same-family feature teacher on audited own-policy source contexts.

Development tool, not yet launched. Inputs must be fixed before execution.
Every GPU invocation must use gpu-run. No model update or native readout.
"""
import argparse
import hashlib
import json
import time
from pathlib import Path

import torch

from audit_latent_source_boundary import split_trace
from latent_deliberation import LatentDeliberation
from run_baseline import FastModel, QWEN_CONFIG_06_B, Qwen3Tokenizer

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


def sha(path):
    digest = hashlib.sha256()
    with path.open('rb') as handle:
        for chunk in iter(lambda: handle.read(8*1024*1024), b''):
            digest.update(chunk)
    return digest.hexdigest()


@torch.no_grad()
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', action='append', required=True)
    parser.add_argument('--parent', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--batch', type=int, default=16)
    parser.add_argument('--manifest', required=True)
    args = parser.parse_args()
    manifest_path = ROOT / args.manifest
    manifest = json.loads(manifest_path.read_text())
    assert manifest['status'] == 'FROZEN_SOURCE_TEACHER_EXTRACTION'
    assert manifest['sources'] == args.source and manifest['parent'] == args.parent
    assert manifest['output'] == args.output and manifest['batch'] == args.batch
    for rel, digest in manifest['input_sha256'].items():
        assert sha(ROOT/rel) == digest, rel
    assert 1 <= args.batch <= 32
    out = ROOT / args.output
    out.mkdir(exist_ok=False)
    torch.set_num_threads(4)
    torch.manual_seed(11341202)
    torch.backends.cuda.matmul.allow_tf32 = False
    torch.backends.cudnn.allow_tf32 = False
    start = time.monotonic()
    tok = Qwen3Tokenizer(tokenizer_file_path=ROOT / 'code/data/tokenizer-base.json')
    rows = []
    sources = {}
    for rel in args.source:
        path = ROOT / rel
        sources[rel] = sha(path)
        for line in path.read_text().splitlines():
            row = json.loads(line)
            split = split_trace(row, tok)
            ids = row['prompt'] + split['prefix'] + split['separator']
            assert len(ids) <= 4096 and split['separator'] == [271]
            rows.append({'index': row['index'], 'teacher_ids': ids,
                         'source': rel, 'canonical_answer': split['canonical_answer'],
                         'canonical_content_mask': split['canonical_content_mask'],
                         'prefix_boxes': split['prefix_boxes']})
    assert len(rows) == len({row['index'] for row in rows})
    base = FastModel(dict(QWEN_CONFIG_06_B, context_length=4096, dtype=torch.float32)).cuda()
    base.load_state_dict(torch.load(ROOT / args.parent, map_location='cpu', weights_only=True))
    model = LatentDeliberation(base, 0).eval().requires_grad_(False)
    sorted_rows = sorted(rows, key=lambda row: len(row['teacher_ids']))
    features, work = {}, {'trunk_positions': 0, 'attention_pairs': 0, 'head_positions': 0}
    torch.cuda.reset_peak_memory_stats()
    for offset in range(0, len(sorted_rows), args.batch):
        if time.monotonic() - start > 3000:
            raise TimeoutError('Teacher extraction readiness estimate exceeded; no partial success artifact')
        batch = sorted_rows[offset:offset+args.batch]
        width = max(len(row['teacher_ids']) for row in batch)
        ids = torch.tensor([[tok.eos_token_id]*(width-len(row['teacher_ids'])) + row['teacher_ids']
                            for row in batch], device='cuda')
        mask = torch.tensor([[False]*(width-len(row['teacher_ids'])) + [True]*len(row['teacher_ids'])
                             for row in batch], dtype=torch.bool, device='cuda')
        h, f = model.hidden(ids=ids, mask=mask, boundary_features=True)
        assert torch.isfinite(f).all()
        # Store28 transformer outputs including final normalized state;
        # exclude the common delimiter's embedding, which has no context.
        f = f[:, 1:].cpu().contiguous()
        for row, feature in zip(batch, f):
            features[row['index']] = feature.clone()
        work['trunk_positions'] += len(batch)*width
        work['attention_pairs'] += len(batch)*width*width
        del h, f, ids, mask
        if offset % (16*args.batch) == 0:
            print(json.dumps({'complete':len(features), 'seconds':time.monotonic()-start}), flush=True)
    keys = sorted(features)
    tensor = torch.stack([features[key] for key in keys])
    assert tensor.shape == (len(rows), 28, 1024) and not tensor.requires_grad
    torch.save({'indices': keys, 'features': tensor, 'source_sha256': sources,
                'teacher_checkpoint_sha256': sha(ROOT/args.parent)}, out/'features.pt')
    (out/'contexts.jsonl').write_text(''.join(json.dumps(row)+'\n' for row in rows))
    result = {'status': 'COMPLETE', 'scope': 'frozen own-family training features; no model update or native result',
              'questions': len(rows), 'feature_shape': list(tensor.shape),
              'manifest_sha256': sha(manifest_path),
              'source_sha256': sources, 'checkpoint_sha256': sha(ROOT/args.parent),
              'features_sha256': sha(out/'features.pt'), 'contexts_sha256':sha(out/'contexts.jsonl'),
              'code_sha256': {name: sha(ROOT/'code'/name) for name in
                             ['extract_latent_teacher.py','latent_deliberation.py','audit_latent_source_boundary.py']},
              'work': work, 'peak_allocated_bytes': torch.cuda.max_memory_allocated(),
              'seconds_including_setup_and_save': time.monotonic()-start,
              'layer_std_min_mean_max': [float(tensor.std(-1,correction=0).min()),
                                        float(tensor.std(-1,correction=0).mean()),
                                        float(tensor.std(-1,correction=0).max())]}
    (out/'summary.json').write_text(json.dumps(result,indent=2)+'\n')
    print(json.dumps(result,indent=2))


if __name__ == '__main__':
    main()
