#!/usr/bin/env python3
"""LearnAI: dependency-free CPU flow reference. Python 3.10+.
Original synthetic teaching data; no remote services, packages or GPU.
Commands: check; train OUTPUT --updates N; sample CHECKPOINT OUTPUT.
A supplied implementation is not evidence of learner authorship or mastery.
"""
import argparse
import copy
import json
import math
from pathlib import Path

CENTERS = [(-1.2, -1.2), (-1.2, 1.2), (1.2, -1.2), (1.2, 1.2)]

class RNG:
    def __init__(self, seed):
        self.state = seed & 0xffffffff
    def uniform(self):
        self.state = (self.state * 1664525 + 1013904223) & 0xffffffff
        return (self.state + .5) / 4294967296
    def normal(self):
        return math.sqrt(-2 * math.log(self.uniform())) * math.cos(2 * math.pi * self.uniform())

def data(split):
    if split not in ('training', 'development'):
        raise ValueError('Unknown split')
    return [[c[0] + .2 * math.cos(2 * math.pi * (k + (0 if split == 'training' else .5)) / 8),
             c[1] + .2 * math.sin(2 * math.pi * (k + (0 if split == 'training' else .5)) / 8)]
            for c in CENTERS for k in range(8)]

def initial_model():
    rng = RNG(71)
    return dict(weights=[0 if i >= 368 or 48 <= i < 64 or 320 <= i < 336 else .25 * rng.normal()
                         for i in range(370)], first=[0.] * 370, second=[0.] * 370, updates=0)

def validate_model(model):
    if not isinstance(model, dict) or type(model.get('updates')) is not int or not 0 <= model['updates'] <= 1000:
        raise ValueError('Invalid bounded model/update count')
    for key in ('weights', 'first', 'second'):
        a = model.get(key)
        if not isinstance(a, list) or len(a) != 370 or any(type(v) not in (int, float) or not math.isfinite(v) or abs(v) > 1e8 for v in a):
            raise ValueError('Invalid numeric model array')
    if any(abs(w) > 100 for w in model['weights']) or any(v < 0 for v in model['second']):
        raise ValueError('Invalid weights/second moments')
    return {key: copy.deepcopy(model[key]) for key in ('weights', 'first', 'second', 'updates')}

def forward(model, point, time):
    w = model['weights']
    inputs = [*point, time]
    h = [math.tanh(w[48+j] + sum(v*w[3*j+i] for i,v in enumerate(inputs))) for j in range(16)]
    h2 = [math.tanh(w[320+j] + sum(v*w[64+16*j+i] for i,v in enumerate(h))) for j in range(16)]
    output = [w[368+i] + sum(v*w[336+16*i+j] for j,v in enumerate(h2)) for i in range(2)]
    return inputs, h, h2, output

def pair(source, target, time):
    return dict(point=[(1-time)*a+time*b for a,b in zip(source,target)], time=time,
                target=[b-a for a,b in zip(source,target)])

def training_cases(update):
    rng = RNG(2201 + update)
    return [pair([rng.normal(), rng.normal()], target, rng.uniform()) for target in data('training')]

def objective(model, cases):
    if not 1 <= len(cases) <= 160:
        raise ValueError('Use 1–160 cases')
    w = model['weights']
    loss, gradient = 0., [0.] * 370
    for case in cases:
        inputs,h,h2,output = forward(model, case['point'], case['time'])
        dh2,dh1 = [0.] * 16, [0.] * 16
        for i in range(2):
            error = output[i] - case['target'][i]
            d = error / len(cases)
            loss += error*error / (2*len(cases))
            gradient[368+i] += d
            for j in range(16):
                gradient[336+16*i+j] += d*h2[j]
                dh2[j] += d*w[336+16*i+j]
        for j in range(16):
            d = dh2[j]*(1-h2[j]**2)
            gradient[320+j] += d
            for k in range(16):
                gradient[64+16*j+k] += d*h[k]
                dh1[k] += d*w[64+16*j+k]
        for j in range(16):
            d = dh1[j]*(1-h[j]**2)
            gradient[48+j] += d
            for k in range(3):
                gradient[3*j+k] += d*inputs[k]
    return loss, gradient

def train(model, updates, rate=.01):
    model = validate_model(model)
    if type(updates) is not int or updates < 0 or model['updates'] + updates > 1000 or rate not in (.003,.01,.03):
        raise ValueError('Bound training at 1,000 total updates and a supported rate')
    for _ in range(updates):
        _,gradient = objective(model, training_cases(model['updates']))
        norm = math.hypot(*gradient)
        scale = 5/norm if norm > 5 else 1
        t = model['updates'] + 1
        first = [.9*v + .1*gradient[i]*scale for i,v in enumerate(model['first'])]
        second = [.999*v + .001*(gradient[i]*scale)**2 for i,v in enumerate(model['second'])]
        weights = [w-rate*(first[i]/(1-.9**t))/(math.sqrt(second[i]/(1-.999**t))+1e-8)
                   for i,w in enumerate(model['weights'])]
        model = dict(weights=weights, first=first, second=second, updates=t)
    return validate_model(model)

def metrics(model):
    result = []
    for split, seed in [('training',4001),('development',5001)]:
        rng = RNG(seed)
        pairs = [([rng.normal(),rng.normal()],target) for target in data(split)]
        levels = []
        for time in (0,.25,.5,.75,1):
            cases = [pair(source,target,time) for source,target in pairs]
            levels.append(dict(time=time,loss=objective(model,cases)[0],
                               zeroBaseline=sum(sum(v*v for v in c['target'])/(2*len(cases)) for c in cases)))
        result.append(dict(split=split,levels=levels,loss=sum(r['loss'] for r in levels)/len(levels)))
    return result

def integrate(point, time, dt, method, velocity):
    if method not in ('euler','heun') or not (0 <= time < 1 and 0 < dt <= 1 and time+dt <= 1+1e-12):
        raise ValueError('Invalid integration step')
    first = velocity(point,time)
    predictor = [x+dt*v for x,v in zip(point,first)]
    if method == 'euler':
        return predictor
    second = velocity(predictor,time+dt)
    return [x+dt*(a+b)/2 for x,a,b in zip(point,first,second)]

def sample(model, steps=16, method='heun', seed=7):
    frozen = validate_model(model)
    if type(steps) is not int or steps not in (4,8,16,32,64) or method not in ('euler','heun') or type(seed) is not int or seed not in (7,19,42):
        raise ValueError('Unsupported bounded sampling settings')
    rng = RNG(seed)
    frames = [[[rng.normal(),rng.normal()] for _ in range(24)]]
    for i in range(steps):
        frames.append([integrate(p,i/steps,1/steps,method,lambda x,t:forward(frozen,x,t)[3]) for p in frames[-1]])
    return dict(model=frozen,steps=steps,method=method,seed=seed,index=steps,frames=frames,
                calls=24*steps*(2 if method == 'heun' else 1))

def quality(points):
    distances = [[math.dist(p,c) for c in CENTERS] for p in points]
    counts = [0]*4
    for ds in distances:
        if min(ds) <= .45:
            counts[ds.index(min(ds))] += 1
    return dict(nearMode=sum(counts),counts=counts,meanDistance=sum(min(ds) for ds in distances)/len(points))

def read_checkpoint(path):
    path = Path(path)
    if path.stat().st_size > 500000:
        raise ValueError('Checkpoint exceeds 500,000 bytes')
    value = json.loads(path.read_text(encoding='utf8'))
    if not isinstance(value,dict) or value.get('version') != 'flow-practice-v1':
        raise ValueError('Expected a flow-practice-v1 checkpoint')
    # Resume the explicitly exported live model, not a possibly older frozen run.
    return validate_model(value.get('model'))

def write_artifact(path, model, run, rate, started_at):
    artifact = dict(version='flow-practice-v1',model=model,rate=rate,metrics=metrics(model),run=run,
                    quality=quality(run['frames'][-1]),implementation='LearnAI Python standard-library reference v1',
                    resumedFromUpdates=started_at,source='Original LearnAI synthetic four-mode 32/32 data',
                    limits='Supplied code is not learner authorship. Public repeated development, not independent qualification. Resume history is not authenticated; samples/metrics recompute from parameters.')
    # Exclusive creation protects preceding checkpoints; choose a new output filename.
    with Path(path).open('x',encoding='utf8') as handle:
        json.dump(artifact,handle,indent=2,allow_nan=False)
        handle.write('\n')
    return artifact

def check():
    model = initial_model()
    cases = training_cases(3)[:4]
    analytic = objective(model,cases)[1]
    max_error = 0.
    for i in range(370):
        old,h = model['weights'][i],1e-5
        model['weights'][i] = old+h
        plus = objective(model,cases)[0]
        model['weights'][i] = old-h
        minus = objective(model,cases)[0]
        model['weights'][i] = old
        max_error = max(max_error,abs(analytic[i]-(plus-minus)/(2*h)))
    if max_error > 1e-7:
        raise AssertionError('Gradient check failed')
    whole = train(model,25)
    resumed = train(train(model,10),15)
    if whole != resumed:
        raise AssertionError('Optimizer resume mismatch')
    if set(map(tuple,data('training'))) & set(map(tuple,data('development'))):
        raise AssertionError('Split overlap')
    def solve(steps,method):
        p=[1.,0.]
        for i in range(steps):
            p=integrate(p,i/steps,1/steps,method,lambda x,t:x)
        return p[0]
    if not abs(solve(8,'heun')-math.e)<abs(solve(16,'euler')-math.e):
        raise AssertionError('Analytic solver check failed')
    if sample(whole,32,'euler')['calls'] != sample(whole,16,'heun')['calls']:
        raise AssertionError('Equal-call comparison failed')
    print(json.dumps(dict(check='passed',parametersChecked=370,maxGradientError=max_error,
                          limits='Public reference checks; not independent implementation review.'),indent=2))

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    commands=parser.add_subparsers(dest='command',required=True)
    commands.add_parser('check')
    t=commands.add_parser('train')
    t.add_argument('output');t.add_argument('--updates',type=int,default=100)
    t.add_argument('--resume');t.add_argument('--rate',type=float,choices=[.003,.01,.03],default=.01)
    s=commands.add_parser('sample');s.add_argument('checkpoint');s.add_argument('output')
    for command in (t,s):
        command.add_argument('--steps',type=int,choices=[4,8,16,32,64],default=16)
        command.add_argument('--method',choices=['euler','heun'],default='heun')
        command.add_argument('--seed',type=int,choices=[7,19,42],default=7)
    args=parser.parse_args()
    if args.command=='check':
        check();return
    if Path(args.output).exists():
        parser.error('Output exists; use a new filename to preserve the preceding checkpoint')
    try:
        model=read_checkpoint(args.resume) if args.command=='train' and args.resume else initial_model() if args.command=='train' else read_checkpoint(args.checkpoint)
        started_at=model['updates']
        rate=args.rate if args.command=='train' else None
        if args.command=='train':
            model=train(model,args.updates,rate)
        run=sample(model,args.steps,args.method,args.seed)
        artifact=write_artifact(args.output,model,run,rate,started_at)
        print(json.dumps(dict(output=args.output,updates=model['updates'],calls=run['calls'],quality=artifact['quality']),indent=2))
    except (ValueError,KeyError,TypeError,OSError) as error:
        parser.error(str(error))

if __name__=='__main__':
    main()