"""Controlled recursive native lineage audit, not an on-policy learning recipe."""
import copy
import hashlib
import json
import os
from pathlib import Path
import random
import resource
import subprocess
import time
import torch
from lightnest.learned_patch import encode, load_predictor
from lightnest.pack_feedback import build, SCHEMA
from lightnest.train_patch_rl import evaluate, execution_reward, apply_action, passes
from lightnest.train_robust_patch_rl import make_task
from lightnest import rv32

ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "work/recursive-lineage-reproduction"
HEADS = ('pointer.', 'none.', 'replacement.')
CHECKPOINT = ROOT / 'work/patch-rl-v2/patch-rl.pt'

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

def write(name, obj):
    (OUT / name).write_text(json.dumps(obj, indent=2, sort_keys=True) + '\n')

def invoke(exe, *args):
    result = subprocess.run([str(exe), *map(str, args)], capture_output=True, timeout=15)
    if result.returncode:
        raise RuntimeError(result.stderr.decode())
    return result.stdout

def decode(blob, model):
    assert blob[:8] == b'LNPATCH1' and len(blob) == 32 + 4 * 107994
    values = torch.frombuffer(bytearray(blob[32:]), dtype=torch.float32)
    state, offset = {}, 0
    for name in SCHEMA:
        value = model.state_dict()[name]
        state[name] = values[offset:offset + value.numel()].reshape(value.shape).clone()
        offset += value.numel()
    assert offset == values.numel()
    return state

def logits(model, task):
    c, b = encode(task.command, task.binary)
    return model(torch.tensor([c]), torch.tensor([b]))

def greedy(model, task):
    with torch.no_grad():
        scores, bits = logits(model, task)
        pos = int(scores.argmax(-1)[0])
        byte = 0 if pos == 256 else sum(int(x >= 0) << i for i, x in enumerate(bits[0, pos]))
        return pos, byte

def reference_step(model, task, pos, byte, reward, lr):
    model.zero_grad(set_to_none=True)
    scores, bits = logits(model, task)
    logp = scores.log_softmax(-1)[0, pos]
    if pos != 256:
        target = torch.tensor([(byte >> i) & 1 for i in range(8)], dtype=torch.bool)
        logp = logp + torch.where(target, torch.nn.functional.logsigmoid(bits[0, pos]), torch.nn.functional.logsigmoid(-bits[0, pos])).sum()
    logp.backward()
    with torch.no_grad():
        if reward != 0:
            for name, parameter in model.named_parameters():
                if name.startswith(HEADS) and parameter.grad is not None:
                    parameter.add_((lr * reward * parameter.grad).clamp(-.1, .1)).clamp_(-16, 16)
    return float(logp.detach())

def detailed(model, frozen):
    rows = []
    for family, entries in frozen.items():
        for index, (task, cases) in enumerate(entries):
            with torch.no_grad():
                scores, bits = logits(model, task)
                probabilities = scores.softmax(-1)[0]
                pos = int(probabilities.argmax())
                byte = 0
                confidence = float(probabilities[pos])
                if pos != 256:
                    bp = bits[0, pos].sigmoid()
                    byte = sum(int(x >= .5) << i for i, x in enumerate(bp))
                    confidence = min(confidence, float(torch.maximum(bp, 1-bp).min()))
                gated = 256 if pos == 256 or confidence < .55 or byte == task.binary[pos] else pos
                raw_binary = apply_action(task.binary, pos, byte)
                gated_binary = apply_action(task.binary, gated, byte)
                rows.append(dict(family=family, index=index, position=pos, byte=byte, confidence=confidence, gated_position=gated,
                    raw_success=passes(raw_binary, cases), gated_success=passes(gated_binary, cases),
                    expected=[y for _, y in cases], raw_outputs=[rv32.simulate(raw_binary, xs) for xs, _ in cases], gated_outputs=[rv32.simulate(gated_binary, xs) for xs, _ in cases]))
    return rows

def run():
    OUT.mkdir(parents=True, exist_ok=True)
    if any(OUT.iterdir()):
        raise RuntimeError('Refusing to overwrite a frozen run')
    resource.setrlimit(resource.RLIMIT_CPU, (110, 120))
    torch.set_num_threads(2)
    torch.set_num_interop_threads(2)
    torch.manual_seed(771249)
    started = time.monotonic()
    rng = random.Random(771249)
    erng = random.Random(771250)
    frozen = {family: [make_task(erng, family, True) for _ in range(32)] for family in ('simple', 'already_correct', 'negation', 'keep')}
    serial = {family: [dict(command=t.command, binary_hex=t.binary.hex(), cases=c) for t, c in rows] for family, rows in frozen.items()}
    write('frozen-cases.json', serial)
    # Fix all contexts/action rules before evaluating or updating any model.
    schedule = []
    for cycle in range(8):
        for family, rule in [('already_correct', 'none'), ('simple', 'none'), ('simple', 'greedy'), ('simple', 'zero_first_byte')]:
            task, cases = make_task(rng, family)
            schedule.append((task, cases, family, rule))
    task, cases = make_task(rng, 'simple')
    schedule.append((task, cases, 'simple', 'zero_reward_identity'))
    write('training-contexts.json', [dict(command=t.command, binary_hex=t.binary.hex(), cases=c, family=f, action_rule=r) for t,c,f,r in schedule])
    protocol = dict(seed=771249, evaluation_seed=771250, steps=33, learning_rate=.03, threads=2,
        cpu_budget_seconds=120, address_space_budget_gib=4, wall_budget_seconds=100,
        amendment='Initial 13-step pilot preserved separately; parent requested 24-64 steps before inspecting its findings. This33-step extension uses the same seeds and evaluation, so it is not a new independent efficacy trial.',
        hypothesis='Three or more successive native descendants preserve cumulative C/autograd gradient parity under mixed positive, negative, no-op and zero-reward updates.',
        checkpoint_sha256=digest(CHECKPOINT), frozen_sha256=digest(OUT/'frozen-cases.json'), contexts_sha256=digest(OUT/'training-contexts.json'),
        source_sha256=digest(ROOT/'linux/feedback_nest.c'), driver_sha256=digest(__file__),
        compile_policy='Exactly one initial build; every subsequent generation is invoked from its immediate parent executable, with no repacking or recompilation.',
        action_policy='Eight cycles: NONE on already-correct task; NONE on edit-needed task; current-native greedy action on edit-needed task; position 0 byte 0 on edit-needed task. Last transition applies zero reward to a greedy action.',
        reward='Independent arithmetic cases and RV32 simulator via execution_reward; last zero reward is an explicit mechanism control, not judged task reward.',
        reference='Independent cumulative Torch autograd initialized only once; never reset from native descendants.',
        selection='Initial and final only. No best-checkpoint selection. Frozen .55 gate and raw argmax.',
        evaluation='128 new cases, 32 per family, 12 arithmetic inputs each, test-only registers and held-out command templates. Existing generator is familiar; no broad novelty claim.',
        interpretation='Controlled off-policy score-gradient mechanism experiment, not an unbiased REINFORCE learning trial. Task success reported separately. Fixed architecture/trainer and frozen encoder; no autonomous bootstrap.')
    write('protocol.json', protocol)
    # Guard total address-space after torch startup (shared libraries consume VSZ).
    resource.setrlimit(resource.RLIMIT_AS, (4*1024**3, 4*1024**3))
    model = load_predictor(CHECKPOINT).model.eval()
    reference = copy.deepcopy(model)
    initial_state = copy.deepcopy(model.state_dict())
    baseline = evaluate(model, frozen)
    write('baseline.json', baseline)
    write('baseline-cases.json', detailed(model, frozen))
    initial = build(CHECKPOINT, OUT/'policy-000')
    # Disallow another compiler/packer execution after initial construction.
    real_run = subprocess.run
    def guarded_run(argv, *args, **kwargs):
        assert str(argv[0]).startswith(str(OUT/'policy-')), 'Only emitted lineage executables permitted after initial build'
        return real_run(argv, *args, **kwargs)
    subprocess.run = guarded_run
    current = initial
    initial_hash = digest(initial)
    lineage = []
    all_parents = {str(initial): initial_hash}
    for step, (task, cases, family, rule) in enumerate(schedule, 1):
        if time.monotonic() - started > 100:
            raise TimeoutError('Wall budget exceeded')
        parent_bytes = current.read_bytes()
        parent_blob = invoke(current, '--dump-weights')
        model.load_state_dict(decode(parent_blob, model))
        if rule == 'none':
            pos, byte = 256, 0
        elif rule == 'zero_first_byte':
            pos, byte = 0, 0
        else:
            pos, byte = greedy(model, task)
        judged_reward = execution_reward(task.binary, pos, byte, cases)
        reward = 0.0 if rule == 'zero_reward_identity' else judged_reward
        reference_logp = reference_step(reference, task, pos, byte, reward, .03)
        window = OUT/f'window-{step:03d}.bin'
        window.write_bytes(task.binary)
        child = OUT/f'policy-{step:03d}'
        native = json.loads(invoke(current, '--feedback', task.command, window, pos, byte, reward, .03, '-o', child))
        child_blob = invoke(child, '--dump-weights')
        actual = decode(child_blob, model)
        errors = {key: float((actual[key]-reference.state_dict()[key]).abs().max()) for key in SCHEMA}
        for key in SCHEMA:
            torch.testing.assert_close(actual[key], reference.state_dict()[key], atol=2e-6, rtol=3e-5)
            if not key.startswith(HEADS):
                assert torch.equal(actual[key], initial_state[key])
        at = parent_bytes.index(parent_blob)
        child_bytes = child.read_bytes()
        assert parent_bytes.count(parent_blob) == 1
        assert parent_bytes[:at] == child_bytes[:at] and parent_bytes[at+len(parent_blob):] == child_bytes[at+len(parent_blob):]
        assert current.read_bytes() == parent_bytes
        if reward == 0:
            assert child_bytes == parent_bytes
        assert abs(native['log_probability'] - reference_logp) < 2e-4
        lineage.append(dict(step=step, parent=current.name, child=child.name, parent_sha256=digest(current), child_sha256=digest(child),
            parent_blob_sha256=hashlib.sha256(parent_blob).hexdigest(), child_blob_sha256=hashlib.sha256(child_blob).hexdigest(),
            family=family, action_rule=rule, position=pos, byte=byte, judged_reward=judged_reward, applied_reward=reward,
            native_log_probability=native['log_probability'], reference_log_probability=reference_logp,
            cumulative_max_weight_error=max(errors.values()), per_tensor_max_error=errors,
            changed_weight_floats=sum(int((actual[k] != model.state_dict()[k]).sum()) for k in SCHEMA),
            outside_blob_unchanged=True, parent_unchanged=True))
        all_parents[str(child)] = digest(child)
        write('lineage.json', lineage)
        current = child
    model.load_state_dict(actual)
    final = evaluate(model, frozen)
    reference_final = evaluate(reference, frozen)
    write('reference-final.json', reference_final)
    native_rows = detailed(model, frozen)
    reference_rows = detailed(reference, frozen)
    write('reference-final-cases.json', reference_rows)
    decision_disagreements = sum(any(a[k] != b[k] for k in ('position','byte','gated_position','raw_success','gated_success')) for a,b in zip(native_rows, reference_rows))
    write('final-cases.json', detailed(model, frozen))
    # Native forward cross-check on every final held-out raw action; task judging remains independent.
    probe_errors = []
    for family, entries in frozen.items():
        for index, (task, _) in enumerate(entries):
            pos, byte = greedy(model, task)
            scores, bits = logits(model, task)
            logp = scores.log_softmax(-1)[0, pos]
            if pos != 256:
                target = torch.tensor([(byte >> i) & 1 for i in range(8)], dtype=torch.bool)
                logp = logp + torch.where(target, torch.nn.functional.logsigmoid(bits[0,pos]), torch.nn.functional.logsigmoid(-bits[0,pos])).sum()
            window = OUT/'eval-window.bin'
            window.write_bytes(task.binary)
            observed = json.loads(invoke(current, '--prob', task.command, window, pos, byte))
            error = abs(float(logp.detach())-observed['log_probability'])
            assert error < 2e-4
            probe_errors.append(dict(family=family,index=index,error=error))
    write('native-heldout-parity.json', probe_errors)
    assert digest(CHECKPOINT) == protocol['checkpoint_sha256']
    assert all(digest(p) == h for p,h in all_parents.items())
    assert digest(OUT/'frozen-cases.json') == protocol['frozen_sha256']
    report = dict(protocol=protocol, baseline=baseline, final=final, transitions=len(lineage), reference_final=reference_final, native_reference_decision_disagreements=decision_disagreements, compiler_calls=1,
        positive_rewards=sum(x['applied_reward']>0 for x in lineage), negative_rewards=sum(x['applied_reward']<0 for x in lineage), zero_rewards=sum(x['applied_reward']==0 for x in lineage),
        max_cumulative_weight_error=max(x['cumulative_max_weight_error'] for x in lineage),
        max_cumulative_log_probability_error=max(abs(x['native_log_probability']-x['reference_log_probability']) for x in lineage),
        max_heldout_native_log_probability_error=max(x['error'] for x in probe_errors),
        all_parent_executables_unchanged=True, source_checkpoint_unchanged=True, frozen_encoder_unchanged=True,
        initial_executable_sha256=initial_hash, final_executable_sha256=digest(current),
        elapsed_seconds=time.monotonic()-started, process_cpu_seconds=time.process_time(), peak_rss_kib=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
        limitations=['Controlled actions are off-policy; no claim of unbiased REINFORCE.', 'Final zero-reward transition intentionally has identical bytes/hash despite being separately emitted.', 'External orchestration, reward and sampling/argmax remain required.', 'Fixed executable code and architecture, output-head changes only.', '128 finite generated cases are not correctness proof or broad generalization.'])
    write('report.json', report)
    manifest = {p.name:digest(p) for p in OUT.iterdir() if p.is_file() and p.name not in ('manifest.json','run.log')}
    write('manifest.json',manifest)
    print(json.dumps(report, indent=2), flush=True)
    return report

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=OUT)
    parser.add_argument('--checkpoint', type=Path, default=CHECKPOINT)
    args = parser.parse_args()
    OUT = args.output.resolve()
    CHECKPOINT = args.checkpoint.resolve()
    run()