"""Run the two small examples. Requires numpy 2.4.6, scipy 1.16.2, sklearn 1.7.2."""
from pathlib import Path
from datetime import datetime, timezone
import argparse
import hashlib
import json
import platform
import warnings

import numpy as np
import scipy
import sklearn
from sklearn.linear_model import Perceptron
from sklearn.neural_network import MLPClassifier

def run():
    X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=float)
    y = np.array([0, 1, 1, 0])
    configurations = [
        ('perceptron', Perceptron(max_iter=1000, tol=None, random_state=42, shuffle=False)),
        ('mlp-tanh', MLPClassifier(hidden_layer_sizes=(8,), activation='tanh', solver='lbfgs', max_iter=2000, random_state=42, tol=1e-8, alpha=0.0001)),
        ('mlp-identity', MLPClassifier(hidden_layer_sizes=(8,), activation='identity', solver='lbfgs', max_iter=2000, random_state=42, tol=1e-8, alpha=0.0001)),
    ]
    fitted=[]
    for name, model in configurations:
        with warnings.catch_warnings(record=True) as captured:
            warnings.simplefilter('always')
            model.fit(X, y)
        predicted=model.predict(X)
        fitted.append({'name':name,'parameters':model.get_params(),'predicted':predicted.tolist(),'correct':int(np.sum(predicted==y)),
                       'nIter':int(model.n_iter_),'warnings':[str(w.message) for w in captured]})
    Q = np.array([[1., 0.], [0., 1.], [1., 1.]])
    K = Q.copy()
    V = np.array([[1., 0.], [0., 1.], [1., 1.]])
    attention=[]
    for name,q,v in [('baseline',Q,V),('query-first-row-changed',np.array([[0.,1.],[0.,1.],[1.,1.]]),V),('value-first-row-changed',Q,np.array([[2.,0.],[0.,1.],[1.,1.]]))]:
        scores=q @ K.T / np.sqrt(2)
        weights=np.exp(scores-scores.max(axis=1,keepdims=True))
        weights/=weights.sum(axis=1,keepdims=True)
        output=weights @ v
        attention.append({'name':name,'Q':q.tolist(),'K':K.tolist(),'V':v.tolist(),'scores':scores.tolist(),'weights':weights.tolist(),'rowSums':weights.sum(axis=1).tolist(),'output':output.tolist()})
    return {'completedAt':datetime.now(timezone.utc).isoformat(),'location':'server CPU','python':platform.python_version(),
            'versions':{'numpy':np.__version__,'scipy':scipy.__version__,'sklearn':sklearn.__version__},
            'seed':42,'scriptSha256':hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
            'xor':{'inputs':X.tolist(),'expected':y.tolist(),'evaluatedOnTrainingPoints':True,'runs':fitted},
            'attention':{'manuallyChosenQKV':True,'causalMask':False,'trainedLanguageModelUsed':False,'runs':attention}}

if __name__=='__main__':
    parser=argparse.ArgumentParser();parser.add_argument('--out',default='foundation-results.json');args=parser.parse_args()
    result=run();path=Path(args.out);path.parent.mkdir(parents=True,exist_ok=True)
    path.write_text(json.dumps(result,ensure_ascii=False,indent=2,allow_nan=False)+'\n')
    print(json.dumps({'saved':str(path),'xor':{r['name']:r['predicted'] for r in result['xor']['runs']},'attention':{r['name']:{'firstWeights':r['weights'][0],'firstOutput':r['output'][0]} for r in result['attention']['runs']}},ensure_ascii=False))