#!/usr/bin/env python3
"""LearnAI Course 14 local study. Python 3.10+, no dependencies.

Original LearnAI teaching code and synthetic data; see expanded-study-NOTICE.txt; no separate open-source license asserted. No external dataset or service.
Default mode uses learner TODO functions; --reference selects the supplied solution.
Public software checks are not independent proficiency or human review.
"""
import argparse
import hashlib
import sys
from pathlib import Path
import json
import math


def rows(task, monitor=False):
    return [dict(id=(0 if task == 'old' else 100)+(1000 if monitor else 0)+i,
                 x=(-.94+i*1.88/15 if monitor else -1+i*2/15), task=task)
            for i in range(16)]


def target(row):
    return row['x'] if row['task'] == 'old' else -row['x']+.5


def initial():
    return [math.sin(i*7+2)*.25 for i in range(13)]


def predict(weights, row):
    task = -1 if row['task'] == 'old' else 1
    return weights[12]+sum(weights[9+j]*math.tanh(weights[3*j]*row['x']+
        weights[3*j+1]*task+weights[3*j+2]) for j in range(3))


def loss(weights, data):
    return sum(.5*(predict(weights, r)-target(r))**2 for r in data)/len(data)


def learner_gradient(weights, data):
    """TODO: return all 13 derivatives of mean half-squared error.

    Accumulate against the SAME weights, divide by len(data), never update here.
    """
    raise NotImplementedError('Implement learner_gradient or inspect --reference')


def reference_gradient(weights, data):
    gradient = [0.]*13
    for row in data:
        task = -1 if row['task'] == 'old' else 1
        error = (predict(weights, row)-target(row))/len(data)
        gradient[12] += error
        for j in range(3):
            h = math.tanh(weights[3*j]*row['x']+weights[3*j+1]*task+weights[3*j+2])
            d = error*weights[9+j]*(1-h*h)
            gradient[3*j] += d*row['x']
            gradient[3*j+1] += d*task
            gradient[3*j+2] += d
            gradient[9+j] += error*h
    return gradient


def stream(shift=True):
    result = []
    for i in range(160):
        old = (i % 5 != 0) if i < 80 or not shift else (i % 5 == 0)
        row = rows('old' if old else 'new')[(i*7+i//5) % 16]
        result.append(dict(row, time=i))
    return result


def learner_buffer(data, size, seed):
    """TODO: reservoir sampling using a bounded size and reproducible seed.

    Use the same LCG as reference_buffer for browser parity.
    """
    raise NotImplementedError('Implement learner_buffer or inspect --reference')


def reference_buffer(data, size, seed):
    result = list(data[:size])
    for index in range(size, len(data)):
        seed = (1664525*seed+1013904223) & 0xffffffff
        replacement = int(seed / 2**32 * (index+1))
        if replacement < size:
            result[replacement] = data[index]
    return result


def checks(gradient, buffer):
    weights, data = initial(), rows('old')[:3]+rows('new')[:4]
    actual = gradient(weights, data)
    assert len(actual) == 13
    for i in range(13):
        plus, minus = list(weights), list(weights)
        plus[i] += 1e-5
        minus[i] -= 1e-5
        oracle = (loss(plus, data)-loss(minus, data))/(2e-5)
        assert abs(actual[i]-oracle) < 1e-7, (i, actual[i], oracle)
    held = buffer(stream(), 16, 17)
    assert len(held) == 16 and len({r['time'] for r in held}) == 16
    assert held == buffer(stream(), 16, 17), 'Buffer is not reproducible'
    assert all(0 <= r['time'] < 160 for r in held)
    # Independent protocol oracle: source IDs and monitoring IDs are disjoint.
    assert {r['id'] for r in rows('old')+rows('new')}.isdisjoint(
        {r['id'] for r in rows('old', True)+rows('new', True)})


def train(weights, data, gradient, updates):
    for _ in range(updates):
        g = gradient(weights, data)
        candidate = [w-.08*d for w, d in zip(weights, g)]
        assert all(math.isfinite(v) and abs(v) <= 10 for v in candidate)
        weights = candidate
    return weights


def metrics(weights):
    return {task:2*loss(weights, rows(task, True)) for task in ['old', 'new']}


def study(gradient, buffer):
    baseline = train(initial(), rows('old')*2, gradient, 100)
    before = metrics(baseline)
    runs = []
    for seed in [17, 23, 41]:
        replay = buffer(stream(), 16, seed)
        for strategy in ['none', 'new-only', 'reservoir']:
            data = rows('new')*2 if strategy == 'new-only' else rows('new')+replay
            weights, curve = list(baseline), [dict(updates=0, **before)]
            for update in range(1, 101):
                if strategy != 'none':
                    weights = train(weights, data, gradient, 1)
                if update % 10 == 0:
                    curve.append(dict(updates=update, **metrics(weights)))
            after = metrics(weights)
            passed = after['old'] <= before['old']+.02 and before['new']-after['new'] >= .01
            runs.append(dict(seed=seed, strategy=strategy, baseline=baseline,
                candidate=weights, curve=curve, promoted=passed,
                fitting_rows=0 if strategy == 'none' else 3200,
                comparison_slots=100, memory_receipts=16 if strategy == 'reservoir' else 0))
    return dict(format='learnai-adaptation-study-v1', runs=runs,
                limits='All nine runs, including failures. Synthetic repeated development; '
                'buffer seeds are not new task families. No independent review or proficiency claim.')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--reference', action='store_true')
    parser.add_argument('--output', help='Optional local JSON output file')
    args = parser.parse_args()
    gradient = reference_gradient if args.reference else learner_gradient
    buffer = reference_buffer if args.reference else learner_buffer
    checks(gradient, buffer)
    result = study(gradient, buffer)
    result['mode'] = 'supplied-reference' if args.reference else 'learner-implementation'
    source = Path(__file__).read_bytes()
    result['source_sha256'] = hashlib.sha256(source).hexdigest()
    result['source_bytes'] = len(source)
    result['python_version'] = sys.version
    result['provenance'] = 'Original LearnAI teaching code and synthetic data; see expanded-study-NOTICE.txt; no separate open-source license asserted'
    text = json.dumps(result, indent=2, allow_nan=False)
    if args.output:
        with open(args.output, 'x', encoding='utf-8') as output:
            output.write(text+'\n')
    else:
        print(text)