"""LearnAI P05 — PyTorch CPU autodiff, batches, modes and exact resume.
Keep numerical-network.py beside this file. Python 3.11+, torch==2.9.1,
numpy==2.3.5. No GPU, server or downloaded model is required.

python framework-debugging.py --compare
python framework-debugging.py --diagnose
python framework-debugging.py --demo --updates 20 --dropout .5 --out first.json
python framework-debugging.py --demo --resume first.json --updates 20 --out resumed.json
python framework-debugging.py --check  # learner_step TODO initially fails
python framework-debugging.py --train --updates 40 --out learner.json

Checkpoints are bounded inert JSON, not pickle. Synthetic public grids and
visible tests are practice; all reports remain unreviewed learner evidence.
"""
import argparse
import copy
import hashlib
import importlib.util
import json
import math
import platform
import unittest
from pathlib import Path
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.set_num_threads(1)
torch.use_deterministic_algorithms(True)
KEYS=('W1','b1','W2','b2')
PROTOCOL='torch-cpu-adam-batches-v1'


def numpy_reference():
    path=Path(__file__).with_name('numerical-network.py')
    spec=importlib.util.spec_from_file_location('numpy_network',path)
    m=importlib.util.module_from_spec(spec);spec.loader.exec_module(m)
    return m


class Network(nn.Module):
    def __init__(self,dropout=0.):
        super().__init__()
        if dropout not in [0.,.5]:raise ValueError('Dropout must be 0 or 0.5.')
        for key,value in numpy_reference().initial().items():setattr(self,key,nn.Parameter(torch.tensor(value,dtype=torch.float64)))
        self.dropout=nn.Dropout(dropout)
    def forward(self,x,detach_hidden=False):
        h=torch.tanh(x@self.W1+self.b1)
        if detach_hidden:h=h.detach()  # Labelled diagnostic fault only.
        return self.dropout(h)@self.W2+self.b2


def learner_step(model,optimizer,x,y):
    """Implement one update. Set train mode; clear accumulated gradients;
    compute predictions with exact y shape; compute mean squared loss;
    backpropagate; update optimizer; return the pre-update loss as a float.
    Do not detach hidden values. Do not refit inputs or manually change seeds.
    The caller supplies bounded tensors and the deterministic randomness scope.
    """
    raise NotImplementedError('Implement learner_step, then --check.')


def supplied_step(model,optimizer,x,y):
    model.train();optimizer.zero_grad(set_to_none=True)
    prediction=model(x)
    if prediction.shape!=y.shape:raise ValueError('Prediction and target shapes must match exactly.')
    loss=((prediction-y)**2).mean();loss.backward();optimizer.step()
    return float(loss.detach())


def optimizer_for(model,rate):
    return torch.optim.Adam(model.parameters(),lr=rate,betas=(.9,.999),eps=1e-8,weight_decay=0.,foreach=False,fused=False)


def tensors():
    train,development=numpy_reference().data()
    return tuple(tuple(torch.tensor(a,dtype=torch.float64) for a in part) for part in [train,development])


def next_batch(completed):
    if type(completed) is not int or not 0<=completed<1000:raise ValueError('Batch position out of range.')
    (x,y),_=tensors();epoch,offset=divmod(completed,4)
    generator=torch.Generator(device='cpu').manual_seed(177+epoch)
    loader=DataLoader(TensorDataset(x,y,torch.arange(16)),batch_size=4,shuffle=True,num_workers=0,generator=generator)
    return list(loader)[offset]


def measure(model):
    previous=model.training;model.eval()
    try:
        with torch.no_grad():return {name:float(((model(x)-y)**2).mean()) for name,(x,y) in zip(['training','development'],tensors())}
    finally:model.train(previous)


def fresh(rate=.01,dropout=0.):
    if type(rate) not in [int,float] or type(dropout) not in [int,float] or rate not in [.001,.01,.05] or dropout not in [0.,.5]:raise ValueError('Unknown rate/dropout.')
    model=Network(dropout)
    return dict(model=model,optimizer=optimizer_for(model,rate),completed=0,rate=rate,dropout=dropout,history=[dict(step=0,**measure(model))],stopped=None)


def advance(session,count,step_fn=supplied_step):
    if type(count) is not int or not 1<=count<=100 or session['completed']+count>1000:raise ValueError('Use 1–100 updates per command and 1,000 total.')
    if session['stopped']:raise ValueError('Stopped run: inspect it and restart explicitly.')
    model,optimizer=session['model'],session['optimizer']
    for _ in range(count):
        before=copy.deepcopy(model.state_dict());memory=copy.deepcopy(optimizer.state_dict());mode=model.training
        x,y,ids=next_batch(session['completed'])
        try:
            # Absolute-step seed schedule is the RNG state for this tiny recipe.
            # It isolates training dropout from unrelated inspection calls.
            with torch.random.fork_rng(devices=[]):
                torch.manual_seed(9101+session['completed'])
                batch_loss=step_fn(model,optimizer,x,y)
            if not isinstance(batch_loss,(float,int)) or not math.isfinite(batch_loss) or not 0<=batch_loss<=1e6:raise ValueError('Invalid batch loss.')
            if any(not torch.isfinite(p).all() or torch.max(torch.abs(p))>1000 for p in model.parameters()):raise ValueError('Parameter bound exceeded.')
            if any(p.grad is None or not torch.isfinite(p.grad).all() for p in model.parameters()):raise ValueError('Missing/nonfinite gradient path.')
            scores=measure(model)
            if any(not math.isfinite(v) or not 0<=v<=1e6 for v in scores.values()):raise ValueError('Loss bound exceeded.')
        except (ValueError,RuntimeError) as error:
            model.load_state_dict(before);optimizer.load_state_dict(memory);model.train(mode);optimizer.zero_grad(set_to_none=True)
            session['stopped']='Rejected update; preceding model and optimizer retained: '+str(error)[:200];break
        session['completed']+=1
        session['history'].append(dict(step=session['completed'],batch_ids=ids.tolist(),batch_loss=float(batch_loss),**scores))
    return session


def recipe():
    return dict(protocol=PROTOCOL,torch=torch.__version__,numpy=np.__version__,python=platform.python_version(),
                script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                reference_sha256=hashlib.sha256(Path(__file__).with_name('numerical-network.py').read_bytes()).hexdigest())


def checkpoint(session,implementation='supplied'):
    memory={}
    for name,p in session['model'].named_parameters():
        state=session['optimizer'].state.get(p,{})
        memory[name]=None if not state else dict(step=int(state['step'].item()),mean=state['exp_avg'].tolist(),square=state['exp_avg_sq'].tolist())
    completed=session['completed']
    return dict(recipe=recipe(),implementation=implementation,rate=session['rate'],dropout=session['dropout'],completed=completed,
                model={k:p.detach().tolist() for k,p in session['model'].named_parameters()},optimizer=memory,
                cursor=dict(epoch=completed//4,batch=completed%4,dropout_seed=9101+completed),
                history=copy.deepcopy(session['history']),stopped=session['stopped'],
                limits='CPU float64, 13 parameters, 4 training rows/update, public disjoint 16/16 grids. History is learner-editable, not authenticated training evidence. No reserved test; human review pending.')


def bounded_array(value,shape,bound,nonnegative=False):
    def numeric_tree(v):return all(numeric_tree(x) for x in v) if isinstance(v,list) else type(v) in [int,float]
    if not isinstance(value,list) or not numeric_tree(value):raise ValueError('Expected numerical JSON arrays.')
    a=np.asarray(value,dtype=np.float64)
    if a.shape!=shape or not np.isfinite(a).all() or np.max(np.abs(a))>bound or (nonnegative and np.any(a<0)):raise ValueError('Invalid checkpoint numerical array.')
    return torch.tensor(a,dtype=torch.float64)


def restore(record):
    if not isinstance(record,dict) or record.get('implementation') not in ['supplied','learner-unreviewed'] or record.get('recipe')!=recipe():raise ValueError('Checkpoint code/environment does not match.')
    n=record.get('completed')
    if type(n) is not int or not 0<=n<=1000 or record.get('cursor')!=dict(epoch=n//4,batch=n%4,dropout_seed=9101+n):raise ValueError('Invalid update count or data/randomness cursor.')
    session=fresh(record.get('rate'),record.get('dropout'))
    if not isinstance(record.get('model'),dict) or set(record['model'])!=set(KEYS) or not isinstance(record.get('optimizer'),dict) or set(record['optimizer'])!=set(KEYS):raise ValueError('Missing model/optimizer fields.')
    for name,p in session['model'].named_parameters():
        values=bounded_array(record['model'][name],tuple(p.shape),1000)
        with torch.no_grad():p.copy_(values)
        state=record['optimizer'][name]
        if n==0:
            if state is not None:raise ValueError('Untrained checkpoint has optimizer moments.')
        else:
            if not isinstance(state,dict) or set(state)!={'step','mean','square'} or type(state['step']) is not int or state['step']!=n:raise ValueError('Optimizer update count differs.')
            session['optimizer'].state[p]=dict(step=torch.tensor(float(n),dtype=torch.float32),
                exp_avg=bounded_array(state['mean'],tuple(p.shape),1e12),exp_avg_sq=bounded_array(state['square'],tuple(p.shape),1e24,True))
    history=record.get('history')
    if not isinstance(history,list) or len(history)!=n+1:raise ValueError('Invalid bounded history.')
    for i,row in enumerate(history):
        if not isinstance(row,dict) or type(row.get('step')) is not int or row['step']!=i or any(type(row.get(k)) not in [int,float] or not math.isfinite(row[k]) or not 0<=row[k]<=1e6 for k in ['training','development']):raise ValueError('Invalid loss history.')
        if i and (row.get('batch_ids')!=next_batch(i-1)[2].tolist() or type(row.get('batch_loss')) not in [float,int] or not math.isfinite(row['batch_loss']) or not 0<=row['batch_loss']<=1e6):raise ValueError('Invalid batch trace.')
    for key,value in measure(session['model']).items():
        if not math.isclose(value,history[-1][key],rel_tol=1e-12,abs_tol=1e-12):raise ValueError('Final losses disagree with restored model.')
    stopped=record.get('stopped')
    if stopped is not None and (not isinstance(stopped,str) or len(stopped)>300):raise ValueError('Invalid stop reason.')
    session.update(completed=n,history=copy.deepcopy(history),stopped=stopped)
    return session


def compare():
    reference=numpy_reference();p=reference.initial();x,y=reference.teaching_case()
    model=Network();prediction=model(torch.tensor(x));loss=((prediction-torch.tensor(y))**2).mean();loss.backward()
    expected=reference.supplied_backward(p,x,y);details=[]
    for key,param in model.named_parameters():
        difference=float(np.max(np.abs(param.grad.numpy()-expected[key])))
        details.append(dict(parameter=key,max_absolute_difference=difference,passed=difference<1e-12))
    forward_error=float(np.max(np.abs(prediction.detach().numpy()-reference.supplied_forward(p,x)['prediction'])))
    return dict(forward_max_absolute_difference=forward_error,gradients=details,
                passed=forward_error<1e-12 and all(d['passed'] for d in details),limits='Agreement on a fixed example does not certify every implementation case.')


def diagnose():
    x,y=numpy_reference().teaching_case();x=torch.tensor(x);y=torch.tensor(y)
    model=Network(.5);model.eval();model.zero_grad(set_to_none=True)
    ((model(x,detach_hidden=True)-y)**2).mean().backward()
    detached={name:p.grad is not None for name,p in model.named_parameters()}
    with torch.random.fork_rng(devices=[]),torch.no_grad():
        torch.manual_seed(321);model.train();a=model(x).tolist();b=model(x).tolist()
        model.eval();c=model(x).tolist();d=model(x).tolist()
    return dict(detached_hidden_gradient_present=detached,training_mode_predictions=[a,b],evaluation_mode_predictions=[c,d],
                explanation='Detach removes hidden-parameter gradient paths while output gradients remain. no_grad does not disable dropout; eval changes module mode. Neither inspection trains the model.')


def resume_example():
    saved=checkpoint(advance(fresh(dropout=.5),7))
    full,weights_only=restore(saved),restore(saved)
    before=measure(full['model'])
    # Diagnostic ablation only; not a valid continuation checkpoint.
    weights_only['optimizer']=optimizer_for(weights_only['model'],weights_only['rate'])
    advance(full,1);advance(weights_only,1)
    def values(session):return {k:p.detach().tolist() for k,p in session['model'].named_parameters()}
    a,b=values(full),values(weights_only)
    return dict(recipe=recipe(),checkpoint_updates=7,next_batch_ids=next_batch(7)[2].tolist(),cursor=saved['cursor'],
                before=saved['model'],before_metrics=before,
                full=dict(parameters=a,metrics=measure(full['model'])),
                weights_only=dict(parameters=b,metrics=measure(weights_only['model'])),
                max_parameter_difference=max(float(np.max(np.abs(np.array(a[k])-np.array(b[k])))) for k in KEYS),
                limits='Recorded CPU diagnostic: both branches start with identical weights and use the same next batch/dropout seed. Only Adam memory is reset in the ablation. These are not two independently trained models or reviewed learner evidence.')


class StepPractice(unittest.TestCase):
    def test_one_step_matches_supplied(self):
        a=fresh(dropout=.5);b=fresh(dropout=.5);b['model'].eval()
        advance(a,1);advance(b,1,learner_step)
        self.assertIsNone(b['stopped'])
        for key,p in a['model'].named_parameters():torch.testing.assert_close(p,dict(b['model'].named_parameters())[key],rtol=0,atol=1e-12)
    def test_two_steps_clear_old_gradients(self):
        a=fresh();b=fresh();advance(a,2);advance(b,2,learner_step)
        self.assertIsNone(b['stopped'])
        for key,p in a['model'].named_parameters():torch.testing.assert_close(p,dict(b['model'].named_parameters())[key],rtol=0,atol=1e-12)
    def test_exact_target_shape(self):
        s=fresh();x,y,_=next_batch(0)
        with self.assertRaises(ValueError):learner_step(s['model'],s['optimizer'],x,y[:,0])


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    modes=parser.add_mutually_exclusive_group(required=True)
    for flag in ['compare','diagnose','resume-example','demo','train','check']:modes.add_argument('--'+flag,action='store_true')
    parser.add_argument('--updates',type=int,default=20);parser.add_argument('--rate',type=float,default=.01);parser.add_argument('--dropout',type=float,default=0.)
    parser.add_argument('--resume');parser.add_argument('--out')
    args=parser.parse_args()
    if args.check:
        result=unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(StepPractice));raise SystemExit(0 if result.wasSuccessful() else 1)
    if args.compare or args.diagnose or args.resume_example:
        result=compare() if args.compare else diagnose() if args.diagnose else resume_example();print(json.dumps(result,indent=2,allow_nan=False))
        if args.compare and not result['passed']:raise SystemExit(1)
        return
    if not args.out:parser.error('Supply a new checkpoint filename with --out.')
    if args.resume:
        with Path(args.resume).open('rb') as stream:raw=stream.read(500001)
        if len(raw)>500000:raise ValueError('Checkpoint exceeds 500 KB.')
        record=json.loads(raw);session=restore(record)
        if record['implementation']!=('supplied' if args.demo else 'learner-unreviewed'):raise ValueError('Resume must retain supplied versus learner implementation provenance.')
        # Resume uses saved controls; avoid silently applying new CLI values.
        if args.rate!=.01 or args.dropout!=0.:parser.error('On resume, omit --rate and --dropout; saved controls apply.')
    else:session=fresh(args.rate,args.dropout)
    advance(session,args.updates,supplied_step if args.demo else learner_step)
    out=checkpoint(session,'supplied' if args.demo else 'learner-unreviewed')
    with Path(args.out).open('x') as stream:json.dump(out,stream,indent=2,allow_nan=False)
    print(f"Saved {args.out}; {session['completed']} updates; {session['stopped'] or 'no numerical stop'}. Human review pending.")

if __name__=='__main__':main()