#!/usr/bin/env python3
"""Run Lab 11's scalar SFT/DPO arithmetic; no model or assistant behavior."""
import argparse
import json
import math
from pathlib import Path
import platform


def sigmoid(z):
    if z >= 0:
        return 1/(1+math.exp(-z))
    e = math.exp(z)
    return e/(1+e)


def loss(z, beta):
    scaled = beta*z
    return max(0, -scaled) + math.log1p(math.exp(-abs(scaled)))


def run():
    experiments = {}
    for name, beta in [('sft', 1.0), ('dpo-1', 1.0), ('dpo-half', 0.5)]:
        z, rows = 0.0, []
        for step in range(21):
            pa = sigmoid(z)
            rows.append({'step': step, 'z': z, 'p_A': pa, 'p_B': 1-pa,
                         'loss': loss(z, beta), 'reference_relative_margin': z})
            assert abs(rows[-1]['p_A']+rows[-1]['p_B']-1) < 1e-15
            if step < 20:
                z -= 0.1*beta*(sigmoid(beta*z)-1)
        experiments[name] = rows
    assert experiments['sft'] == experiments['dpo-1']
    assert experiments['sft'][1]['z'] == 0.05
    assert experiments['dpo-half'][1]['z'] == 0.025
    assert all(abs(r[0]['loss']-math.log(2)) < 1e-15 for r in experiments.values())
    return {'python': platform.python_version(), 'steps': 20, 'learning_rate': 0.1,
            'scope': 'scalar arithmetic only; no model loaded or behavioral comparison run',
            'experiments': experiments}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    if any(p.is_symlink() for p in [args.output, *args.output.parents]):
        raise ValueError('Output path must not traverse symlinks')
    result = run()
    with args.output.open('x') as stream:
        json.dump(result, stream, indent=2, allow_nan=False);stream.write('\n')
    print(json.dumps(result, indent=2))