"""LearnAI P10 reference. Python 3, standard library only. Synthetic public data.
Run: python research-study.py --output my-study.json
Compare every run with the browser export (absolute tolerance 1e-10).
Learner extension: copy this file, add one declared feature or penalty, freeze
its protocol before evaluating new independently reserved tasks. Never overwrite
the original output. This supplied reference is not learner implementation evidence.
"""
import argparse
import hashlib
import json
import math
from pathlib import Path

VERSION = 'regression-study-v1'
SEEDS = (7, 19, 41)

def rows(task, partition):
    offset = {'training': 0, 'development': 1, 'final': 2}[partition]
    result = []
    for i in range(12):
        index = i * 3 + offset
        x = -1 + 2 * index / 35
        result.append({'id': f'task{task}/row{index}', 'x': x,
                       'y': .3 + (task+1)*.2*x + (task-1)*.6*x*x + .04*math.sin(index*17+task)})
    return result

def predict(weights, x):
    return weights[0] + weights[1]*x + weights[2]*x*x

def train(task, seed, quadratic=True, regularization=False, updates=80):
    if task not in range(3) or seed not in SEEDS or updates not in (40, 80):
        raise ValueError('Unsupported bounded recipe')
    w = [.02*math.sin(seed), .02*math.cos(seed), .02*math.sin(seed*2) if quadratic else 0.]
    for _ in range(updates):
        g = [0., 0., 0.]
        for row in rows(task, 'training'):
            error = predict(w, row['x']) - row['y']
            for j in range(3 if quadratic else 2):
                g[j] += error * row['x']**j / 12
        if regularization:
            g[1] += .08*w[1]
            if quadratic:
                g[2] += .08*w[2]
        w = [v-.15*d for v,d in zip(w,g)]
    return w

def score(w, values):
    return sum(abs(predict(w,r['x'])-r['y']) for r in values)/len(values)

def study(quadratic=True, regularization=False, updates=80, partition='final'):
    output=[]
    for task in range(3):
        for seed in SEEDS:
            w=train(task,seed,quadratic,regularization,updates)
            baseline=score(train(task,seed,False,False,updates),rows(task,partition))
            candidate=score(w,rows(task,partition))
            output.append(dict(task=task,seed=seed,baseline=baseline,candidate=candidate,
                               difference=baseline-candidate,weights=w,examples=updates*24))
    return output

def checks():
    for task in range(3):
        ids=[r['id'] for part in ('training','development','final') for r in rows(task,part)]
        assert len(set(ids))==36
    assert predict([1,2,3],2)==17
    assert train(0,7,False)[2]==0
    # Independent closed-loop replay: output records are scored directly from stored weights.
    for r in study():
        expected=sum(abs(r['weights'][0]+r['weights'][1]*x['x']+r['weights'][2]*x['x']**2-x['y']) for x in rows(r['task'],'final'))/12
        assert abs(expected-r['candidate'])<1e-12
    return True

if __name__=='__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument('--output',type=Path)
    parser.add_argument('--updates',type=int,choices=(40,80),default=80)
    parser.add_argument('--linear',action='store_true')
    parser.add_argument('--regularization',action='store_true')
    args=parser.parse_args()
    checks()
    payload=dict(version=VERSION,source_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
                 recipe=dict(quadratic=not args.linear,regularization=args.regularization,updates=args.updates),
                 partition='final',seeds=SEEDS,results=study(not args.linear,args.regularization,args.updates),
                 provenance='Authored synthetic public data; supplied reference; not independent learner evidence')
    text=json.dumps(payload,indent=2)
    if args.output:
        with args.output.open('x',encoding='utf-8') as f:
            f.write(text+'\n')
    else:
        print(text)