"""Dependency-free Course 12 reference / learner exercise, Python 3.10+.
Original LearnAI teaching code; see expanded-study-NOTICE.txt. Public synthetic data, no personal data.
python world-model-study.py --check
python world-model-study.py --output my-study.json
python world-model-study.py --exercise  # replace learner_predict first
"""
import argparse
import hashlib
import sys
from pathlib import Path
import json
import math


def environment(state, force, mass=1, drag=.2):
    x, v = state
    v += .1 * (force / mass - drag * v - .6 * v * abs(v))
    return (x + .1 * v, v)


def collect():
    rows = []
    for episode, start in enumerate([-.9, -.6, -.3, .3, .6, .9]):
        state = (start, .4 if episode % 2 else -.4)
        for step in range(12):
            force = 2 * math.sin((episode + 1) * 1.37 + step * .83)
            target = environment(state, force)
            rows.append((episode, state, force, target))
            state = target
    return rows


def solve(matrix, rhs):
    a = [list(row) + [b] for row, b in zip(matrix, rhs)]
    for c in range(len(a)):
        p = max(range(c, len(a)), key=lambda i: abs(a[i][c]))
        a[c], a[p] = a[p], a[c]
        divisor = a[c][c]
        if abs(divisor) < 1e-12:
            raise ValueError('Unresolved fit')
        a[c] = [x / divisor for x in a[c]]
        for i in range(len(a)):
            if i != c:
                factor = a[i][c]
                a[i] = [x - factor * y for x, y in zip(a[i], a[c])]
    return [r[-1] for r in a]


def fit(rows):
    design = [[*s, u, 1] for _, s, u, _ in rows]
    gram = [[sum(z[i] * z[j] for z in design) + (1e-6 if i == j else 0)
             for j in range(4)] for i in range(4)]
    return [solve(gram, [sum(z[i] * (target[k] - state[k])
                            for z, (_, state, _, target) in zip(design, rows))
                        for i in range(4)]) for k in range(2)]


def predict(coefficients, state, force):
    z = [*state, force, 1]
    return tuple(state[k] + sum(a * b for a, b in zip(coefficients[k], z)) for k in range(2))


def learner_predict(coefficients, state, force):
    raise NotImplementedError('Implement the affine delta predictor and independent tests')


def rollout(weights, start, step, horizon=15):
    state, cost, trace = start, 0., []
    for _ in range(horizon):
        force = 2 * math.tanh(sum(a * b for a, b in zip(weights, state)))
        state = step(state, force)
        if not all(math.isfinite(v) for v in state) or abs(state[0]) > 20 or abs(state[1]) > 10:
            raise ValueError('Trajectory outside bounds')
        cost += state[0] ** 2 + .1 * state[1] ** 2 + .01 * (force / 2) ** 2
        trace.append({'state': state, 'force': force})
    return cost / horizon, trace


STARTS = [(-1, -.3), (-.5, .3), (.5, -.3), (1, .3)]


def objective(weights, step):
    return sum(rollout(weights, s, step)[0] for s in STARTS) / len(STARTS)


def train(step, updates=20):
    weights = [0., 0.]
    for _ in range(updates):
        gradient = []
        for j in range(2):
            plus, minus = weights[:], weights[:]
            plus[j] += 1e-4
            minus[j] -= 1e-4
            gradient.append((objective(plus, step) - objective(minus, step)) / 2e-4)
        weights = [max(-3, min(3, w - .3 * g)) for w, g in zip(weights, gradient)]
    return weights


def plan(state, step, horizon=3):
    """Stable width-six beam; no external truth enters learned-model search."""
    beam, calls = [(0., [], state)], 0
    for _ in range(horizon):
        candidates = []
        for cost, actions, previous in beam:
            for force in [-2., 0., 2.]:
                predicted = step(previous, force)
                calls += 1
                if not all(math.isfinite(v) for v in predicted) or abs(predicted[0]) > 20 or abs(predicted[1]) > 10:
                    continue
                added = predicted[0] ** 2 + .1 * predicted[1] ** 2 + .01 * (force / 2) ** 2
                candidates.append((cost + added, actions + [force], predicted))
        candidates.sort(key=lambda item: item[0])
        beam = candidates[:6]
        if not beam:
            raise ValueError('No valid plan; unresolved, not successful')
    return beam[0][1][0], calls


def mpc_rollout(start, predictor, external):
    state, cost, calls, trace = start, 0., 0, []
    for _ in range(20):
        force, work = plan(state, predictor)
        calls += work
        state = external(state, force)
        cost += state[0] ** 2 + .1 * state[1] ** 2 + .01 * (force / 2) ** 2
        trace.append({'state': state, 'force': force})
    return cost / 20, trace, calls


def checks():
    assert environment((1, .5), 0) == (1.0475, .475)
    assert solve([[2, 0], [0, 4]], [6, 8]) == [3, 2]
    rows = collect()
    assert len(rows) == 72 and len(set(r[0] for r in rows)) == 6
    coefficients = fit(rows)
    assert predict([[0, 0, 0, .2], [0, 0, 0, -.1]], (1, .5), 2) == (1.2, .4)
    initial = objective([0, 0], environment)
    trained = train(environment, 5)
    assert objective(trained, environment) < initial
    assert all(math.isfinite(v) for row in coefficients for v in row)
    force, calls = plan((1, 0), environment, 1)
    assert force == -2 and calls == 3
    _, calls = plan((1, 0), environment, 3)
    assert calls == 3 + 9 + 18
    print('Seven public arithmetic, fit, split, optimization and MPC checks passed.')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--check', action='store_true')
    parser.add_argument('--exercise', action='store_true')
    parser.add_argument('--output')
    args = parser.parse_args()
    if args.check:
        checks()
        return
    coefficients = fit(collect())
    predictor = learner_predict if args.exercise else predict
    imagined = lambda s, u: predictor(coefficients, s, u)
    policies = {'imagined': train(imagined), 'simulator': train(environment), 'zero': [0, 0]}
    result = {'format': 'world-model-python-study-v1', 'reference': not args.exercise,
              'source_sha256': hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
              'python_version': sys.version,
              'provenance': 'Original LearnAI teaching code and synthetic data; see expanded-study-NOTICE.txt; no separate open-source license asserted',
              'collection': 'six whole episodes / 72 transitions', 'coefficients': coefficients,
              'policy_updates': 20, 'calls_per_policy_training': 20 * 4 * 4 * 15,
              'policies': policies, 'evaluation': [],
              'scope': 'Different declared collection from browser; public development evidence, not authenticated learner work.'}
    for start in [(-.8, .1), (.8, -.1)]:
        for changed in [False, True]:
            step = lambda s, u: environment(s, u, 1.4 if changed else 1, .35 if changed else .2)
            for name, weights in policies.items():
                cost, trace = rollout(weights, start, step, 20)
                result['evaluation'].append({'start': start, 'changed': changed, 'policy': name, 'cost': cost, 'trace': trace})
            cost, trace, calls = mpc_rollout(start, imagined, step)
            result['evaluation'].append({'start': start, 'changed': changed, 'policy': 'learned-model-mpc',
                                         'cost': cost, 'trace': trace, 'model_calls': calls,
                                         'planner': 'width-six beam, horizon three; approximate search'})
    encoded = json.dumps(result, indent=2, allow_nan=False)
    if args.output:
        with open(args.output, 'x', encoding='utf-8') as stream:
            stream.write(encoded + '\n')
    else:
        print(encoded)


if __name__ == '__main__':
    main()