"""Gudao AI Lab / Python 3.10+ / standard library only / original teaching examples.
Run: python ai_lab.py --lesson vectors
Choices: vectors model loss gradient split generalize project all
All data are synthetic. No downloads, accounts, network or paid services.
"""
import argparse
import json
import math
from pathlib import Path


def mse(y, p):
    if len(y) != len(p) or not y:
        raise ValueError('equal, nonempty arrays required')
    return sum((a-b)**2 for a,b in zip(y,p))/len(y)


def predict(x, w, b):
    return [w*t+b for t in x]


def train(x, y, rate=0.03, steps=3000):
    w=b=0.0
    for _ in range(steps):
        errors=[p-t for p,t in zip(predict(x,w,b),y)]
        dw=2*sum(e*t for e,t in zip(errors,x))/len(x)
        db=2*sum(errors)/len(x)
        w,b=w-rate*dw,b-rate*db
        if not math.isfinite(w+b) or abs(w)+abs(b)>1e12:
            raise ArithmeticError('diverged: try a smaller learning rate')
    return w,b


def vectors():
    # Teaching tokenizer/vectors/weights, not values extracted from an LLM.
    tokens=["北京","是","中国","的"]
    embeddings=[[2,1],[0,1],[4,2],[1,0]]
    weights=[[3,5],[1,-1]]
    def transform(x):
        if any(len(row)!=len(x) for row in weights):
            raise ValueError("vector and weight dimensions must match")
        return [sum(a*b for a,b in zip(row,x)) for row in weights]
    outputs=[transform(x) for x in embeddings]
    assert outputs==[[11,1],[5,-1],[22,2],[3,1]]
    changed=[transform([3,1]),*outputs[1:]]
    assert changed[0]==[14,2] and changed[1:]==outputs[1:]
    return dict(sentence="北京是中国的……", teaching_only=True, tokens=tokens,
                ids=list(range(4)),embedding_shape=[4,2],embeddings=embeddings,
                weights=weights,outputs=outputs,after_editing_beijing=changed)


def model():
    x=[0,1,2,3];y=[1,3,5,7]
    p=predict(x,2,1)
    assert p==y
    return dict(w=2,b=1,predictions=p,new_x=4,new_prediction=9)


def loss():
    y=[1,3];p=[1,5]
    squared=mse(y,p);absolute=sum(abs(a-b) for a,b in zip(y,p))/len(y)
    assert squared==2 and absolute==1
    return dict(mse=squared,mae=absolute,large_error_mse=mse(y,[1,7]))


def gradient():
    x=[0,1,2,3];y=[1,3,5,7];w=.4;b=.2;eps=1e-6
    dw=2*sum((w*t+b-v)*t for t,v in zip(x,y))/len(x)
    numerical=(mse(y,predict(x,w+eps,b))-mse(y,predict(x,w-eps,b)))/(2*eps)
    assert abs(dw-numerical)<1e-5
    w,b=train(x,y)
    assert abs(w-2)<.001 and abs(b-1)<.001
    failed=False
    try: train(x,y,rate=.5,steps=100)
    except ArithmeticError: failed=True
    assert failed, 'negative control must diverge'
    return dict(w=w,b=b,gradient_error=abs(dw-numerical),large_rate_diverged=failed)


def split():
    groups=[[0,2,4,6,8,9],[1,5],[3,7]]
    assert len(set(sum(groups,[])))==10
    x=list(range(10))
    train_mean=sum(x[i] for i in groups[0])/6
    return dict(train=groups[0],validation=groups[1],test=groups[2],
                train_mean=train_mean,all_data_mean=sum(x)/len(x))


def generalize():
    x=[0,1,2,3];y=[1.1,2.9,5.1,6.9]
    memo=dict(zip(x,y));vx=[.5,1.5,2.5];vy=[2,4,6]
    w,b=train(x,y)
    a=mse(vy,[memo.get(t,0) for t in vx]);bb=mse(vy,predict(vx,w,b))
    assert mse(y,[memo[t] for t in x])==0 and bb<a
    return dict(memorizer_train_mse=0,memorizer_validation_mse=a,linear_validation_mse=bb)


def project():
    x=[i/10 for i in range(40)]
    y=[2*t+1+.12*math.sin(7*t) for t in x]
    train_ids=[i for i in range(40) if i%5 not in (1,3)]
    val_ids=[i for i in range(40) if i%5==1]
    test_ids=[i for i in range(40) if i%5==3]
    tx=[x[i] for i in train_ids];ty=[y[i] for i in train_ids]
    candidates=[]
    for rate in [.01,.03]:
        w,b=train(tx,ty,rate=rate,steps=2000)
        v=mse([y[i] for i in val_ids],predict([x[i] for i in val_ids],w,b))
        candidates.append((v,rate,w,b))
    v,rate,w,b=min(candidates)
    targets=[y[i] for i in test_ids]
    test=mse(targets,predict([x[i] for i in test_ids],w,b))
    baseline=mse(targets,[sum(ty)/len(ty)]*len(targets))
    assert test<baseline and test<.03
    result=dict(data='synthetic: x=i/10, y=2*x+1+0.12*sin(7*x), i=0..39',
                split=dict(train=train_ids,validation=val_ids,test=test_ids),
                selected_rate=rate,steps=2000,w=w,b=b,validation_mse=v,
                test_mse=test,baseline_test_mse=baseline,
                limitations='One synthetic dataset. Not evidence of real-world accuracy.')
    Path('report.json').write_text(json.dumps(result,indent=2),encoding='utf-8')
    return result


if __name__=='__main__':
    tasks={f.__name__:f for f in [vectors,model,loss,gradient,split,generalize,project]}
    p=argparse.ArgumentParser();p.add_argument('--lesson',choices=[*tasks,'all'],default='all')
    name=p.parse_args().lesson
    for key in tasks if name=='all' else [name]:
        print(key, json.dumps(tasks[key](),ensure_ascii=False))