"""LearnAI P04: a vectorized 2–3–1 tanh network in NumPy, CPU only.
Python 3.11+; numpy==2.3.5. No downloads or frameworks at runtime.
python numerical-network.py --demo --updates 200 --out supplied-run.json
python numerical-network.py --check  # learner TODO functions initially fail
python numerical-network.py --gradient-check  # finite differences of learner code
python numerical-network.py --train --updates 200 --out learner-run.json

Supplied execution and learner implementation are labelled separately. Public
synthetic data and visible tests are practice, not independent proficiency.
"""
import argparse
import hashlib
import json
import math
import platform
import unittest
from pathlib import Path
import numpy as np

KEYS=('W1','b1','W2','b2')


def initial():
    return dict(W1=np.array([[.4,-.7,.2],[-.3,.5,.8]],dtype=np.float64),
                b1=np.array([.1,-.2,.05],dtype=np.float64),
                W2=np.array([[.2],[-.3],[.4]],dtype=np.float64),b2=np.array([0.],dtype=np.float64))


def validate(p,x,y=None):
    if not isinstance(p,dict) or set(p)!=set(KEYS):raise ValueError('Four named parameter arrays required.')
    if any(not isinstance(v,np.ndarray) or v.dtype!=np.float64 or not np.isfinite(v).all() or np.max(np.abs(v),initial=0)>1000 for v in p.values()):raise ValueError('Finite bounded float64 parameter arrays required.')
    if p['W1'].ndim!=2 or not all(1<=n<=8 for n in p['W1'].shape):raise ValueError('Use 1–8 features and hidden units.')
    features,hidden=p['W1'].shape
    if p['b1'].shape!=(hidden,) or p['W2'].shape!=(hidden,1) or p['b2'].shape!=(1,):raise ValueError('Parameter shapes disagree.')
    if not isinstance(x,np.ndarray) or x.dtype!=np.float64 or x.ndim!=2 or not 1<=len(x)<=64 or x.shape[1]!=features or not np.isfinite(x).all() or np.max(np.abs(x))>100:raise ValueError('Use 1–64 rows of bounded float64 features.')
    if y is not None and (not isinstance(y,np.ndarray) or y.dtype!=np.float64 or y.shape!=(len(x),1) or not np.isfinite(y).all() or np.max(np.abs(y))>100):raise ValueError('Targets must have exact shape (rows,1); reject accidental broadcasting.')


def learner_forward(p,x):
    """Validate inputs, then return {'Z': X@W1+b1, 'H': tanh(Z),
    'prediction': H@W2+b2}. Shapes: X(n,d), W1(d,h), b1(h),
    Z/H(n,h), W2(h,1), b2(1), prediction(n,1). Return new arrays;
    do not mutate p or x. Support all shapes allowed by validate.
    """
    raise NotImplementedError('Implement learner_forward, then --check.')


def learner_backward(p,x,y):
    """Gradient of MEAN squared error, not half-squared or summed loss.
    Validate p/x/y and return a new dict with the same keys/shapes as p.
    All derivatives must use the SAME pre-update parameters. Do not mutate.
    Use the forward cache; start with D = 2*(prediction-y)/n. Then derive
    dW2=H.T@D, db2=sum(D,axis=0), dZ=(D@W2.T)*(1-H*H),
    dW1=X.T@dZ, db1=sum(dZ,axis=0). No extra division after D.
    """
    raise NotImplementedError('Implement learner_backward, then --check.')


def supplied_forward(p,x):
    validate(p,x)
    z=x@p['W1']+p['b1'];h=np.tanh(z)
    return dict(Z=z,H=h,prediction=h@p['W2']+p['b2'])


def supplied_backward(p,x,y):
    validate(p,x,y);cache=supplied_forward(p,x)
    d=2*(cache['prediction']-y)/len(x)
    dz=(d@p['W2'].T)*(1-cache['H']**2)
    return dict(W1=x.T@dz,b1=dz.sum(axis=0),W2=cache['H'].T@d,b2=d.sum(axis=0))


def loss(p,x,y,forward=supplied_forward):
    validate(p,x,y);prediction=forward(p,x)['prediction']
    if not isinstance(prediction,np.ndarray) or prediction.shape!=y.shape or prediction.dtype!=np.float64 or not np.isfinite(prediction).all():raise ValueError('Forward prediction shape or values invalid.')
    return float(np.mean((prediction-y)**2))


def checked_gradients(p,gradients):
    if not isinstance(gradients,dict) or set(gradients)!=set(KEYS) or any(not isinstance(gradients[k],np.ndarray) or gradients[k].shape!=p[k].shape or gradients[k].dtype!=np.float64 or not np.isfinite(gradients[k]).all() for k in KEYS):raise ValueError('Backward arrays must match every parameter shape and be finite float64.')
    return gradients


def simultaneous_step(p,x,y,rate,backward=supplied_backward):
    if not math.isfinite(rate) or not 0<=rate<=.1:raise ValueError('Use rate 0–0.1.')
    validate(p,x,y)
    gradients=checked_gradients(p,backward({k:v.copy() for k,v in p.items()},x.copy(),y.copy()))
    proposed={k:p[k]-rate*gradients[k] for k in KEYS}
    validate(proposed,x,y)
    return proposed


def teaching_case(rows=3):
    if rows not in (1,3):raise ValueError('Use one or three worked rows.')
    return np.array([[1,2],[-1,1],[2,-1]],dtype=np.float64)[:rows],np.array([[1],[-.5],[.5]],dtype=np.float64)[:rows]


def data():
    train_values=[-1,-1/3,1/3,1];dev_values=[-.75,-.25,.25,.75]
    def make(values):
        x=np.array([[a,b] for a in values for b in values],dtype=np.float64)
        y=(.4*x[:,0]-.3*x[:,1]+.2*x[:,0]*x[:,1]).reshape(-1,1)
        return x,y
    return make(train_values),make(dev_values)


def gradient_check(p,x,y,forward=learner_forward,backward=learner_backward):
    analytic=checked_gradients(p,backward({k:v.copy() for k,v in p.items()},x.copy(),y.copy()))
    h=1e-5;rows=[]
    for key in KEYS:
        for index in np.ndindex(p[key].shape):
            plus={k:v.copy() for k,v in p.items()};minus={k:v.copy() for k,v in p.items()}
            plus[key][index]+=h;minus[key][index]-=h
            numeric=(loss(plus,x,y,forward)-loss(minus,x,y,forward))/(2*h)
            exact=float(analytic[key][index]);error=abs(exact-numeric);tolerance=1e-6+1e-4*max(abs(exact),abs(numeric))
            rows.append(dict(parameter=key,index=list(index),analytic=exact,finite_difference=numeric,error=error,tolerance=tolerance,passed=error<=tolerance))
    return dict(method='central difference, h=1e-5; absolute tolerance 1e-6 plus relative tolerance 1e-4',passed=all(r['passed'] for r in rows),checks=rows)


def train(updates,rate,forward,backward):
    if type(updates) is not int or not 1<=updates<=1000:raise ValueError('Use 1–1,000 full-batch updates.')
    if rate not in [.001,.01,.05,.1]:raise ValueError('Choose a listed rate.')
    (x,y),(dx,dy)=data();p=initial();history=[];stopped=None
    def measured(step):return dict(step=step,training=loss(p,x,y,forward),development=loss(p,dx,dy,forward))
    history.append(measured(0))
    for step in range(1,updates+1):
        try:
            proposed=simultaneous_step(p,x,y,rate,backward)
            a,b=loss(proposed,x,y,forward),loss(proposed,dx,dy,forward)
            if not math.isfinite(a+b) or max(a,b)>1e6:raise ValueError('Loss exceeds numerical bound.')
        except ValueError as error:
            stopped='Proposal rejected; preceding model retained: '+str(error);break
        p=proposed;history.append(dict(step=step,training=a,development=b))
    return dict(protocol='numpy-two-layer-mse-v1',implementation='supplied' if forward is supplied_forward else 'learner-unreviewed',
                source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                environment=dict(python=platform.python_version(),numpy=np.__version__),
                config=dict(rate=rate,requested_updates=updates,accepted_updates=len(history)-1,parameters=13,training_rows=16,development_rows=16),
                parameters={k:v.tolist() for k,v in p.items()},history=history,stopped=stopped,
                training_mean_baseline_development_mse=float(np.mean((float(y.mean())-dy)**2)),
                examples_processed=(len(history)-1)*len(x),
                limitations='Fixed public synthetic grids, disjoint inputs; repeated development evidence, no reserved test or deployment claim. Supplied teacher formula is not learned. Human review pending.')


class NetworkPractice(unittest.TestCase):
    def test_known_zero_hidden(self):
        p=dict(W1=np.zeros((2,2)),b1=np.zeros(2),W2=np.array([[2.],[-1.]]),b2=np.zeros(1))
        x=np.array([[1.,2.]]);y=np.array([[1.]])
        out=learner_forward(p,x)
        self.assertEqual(set(out),{'Z','H','prediction'})
        for key,shape in [('Z',(1,2)),('H',(1,2)),('prediction',(1,1))]:
            self.assertEqual(out[key].shape,shape);self.assertEqual(out[key].dtype,np.float64)
        np.testing.assert_allclose(out['prediction'],[[0]])
        g=learner_backward(p,x,y)
        np.testing.assert_allclose(g['W1'],[[-4,2],[-8,4]]);np.testing.assert_allclose(g['b1'],[-4,2])
        np.testing.assert_allclose(g['W2'],[[0],[0]]);np.testing.assert_allclose(g['b2'],[-2])
    def test_all_parameter_finite_differences(self):
        x,y=teaching_case();self.assertTrue(gradient_check(initial(),x,y,learner_forward,learner_backward)['passed'])
    def test_target_shape_and_input_integrity(self):
        p=initial();x,y=teaching_case();before={k:v.copy() for k,v in p.items()};original=x.copy()
        with self.assertRaises(ValueError):learner_backward(p,x,y[:,0])
        learner_forward(p,x);learner_backward(p,x,y)
        for key in KEYS:np.testing.assert_array_equal(p[key],before[key])
        np.testing.assert_array_equal(x,original)
    def test_mean_gradient_not_sum(self):
        x,y=teaching_case(1);p=initial()
        a=learner_backward(p,x,y);b=learner_backward(p,np.repeat(x,3,axis=0),np.repeat(y,3,axis=0))
        for key in KEYS:np.testing.assert_allclose(a[key],b[key],rtol=1e-10,atol=1e-12)


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    modes=parser.add_mutually_exclusive_group(required=True)
    modes.add_argument('--demo',action='store_true');modes.add_argument('--train',action='store_true');modes.add_argument('--check',action='store_true');modes.add_argument('--gradient-check',action='store_true')
    parser.add_argument('--updates',type=int,default=200);parser.add_argument('--rate',type=float,default=.05);parser.add_argument('--out')
    args=parser.parse_args()
    if args.check:
        result=unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(NetworkPractice))
        raise SystemExit(0 if result.wasSuccessful() else 1)
    if args.gradient_check:
        x,y=teaching_case();result=gradient_check(initial(),x,y)
        print(json.dumps(result,indent=2,allow_nan=False));raise SystemExit(0 if result['passed'] else 1)
    if not args.out:parser.error('Provide a new --out JSON filename.')
    result=train(args.updates,args.rate,supplied_forward if args.demo else learner_forward,supplied_backward if args.demo else learner_backward)
    with Path(args.out).open('x') as stream:json.dump(result,stream,indent=2,allow_nan=False)
    print(f"Saved {args.out}: {result['implementation']}; {result['config']['accepted_updates']} updates. Human review pending.")

if __name__=='__main__':main()