import json
from pathlib import Path
import numpy as np

SEED = 7
RNG = np.random.default_rng(SEED)


def paper_matrix_step(x1, x2, a, c, d, h, e, f, u):
    z = np.array([x1, x2, x1 * x1, 1.0])
    A = np.array([[a, 0, 0, 0], [0, c, d, h], [0, 0, a*a, 0], [0, 0, 0, 1.]])
    B = np.array([[0, 0, 0, 0], [e, 0, 0, f], [0, 0, 0, 0], [0, 0, 0, 0.]])
    matrix_next = A @ z + (B @ z) * u
    direct_next = np.array([a*x1, c*x2 + d*x1*x1 + h + (e*x1 + f)*u, (a*x1)**2, 1.0])
    return matrix_next, direct_next


def plant(x, u):
    return (0.72 + 0.58 * u) * x + 0.08 * u + 0.03


def make_data(n_sequences=160, horizon=24):
    rows, targets = [], []
    starts, controls = [], []
    for _ in range(n_sequences):
        x = RNG.uniform(-1.0, 1.0)
        starts.append(x)
        us = RNG.uniform(-1.0, 1.0, horizon)
        controls.append(us)
        for u in us:
            rows.append((x, u))
            x = plant(x, u)
            targets.append(x)
    return np.asarray(rows), np.asarray(targets), np.asarray(starts), np.asarray(controls)


def design(rows, kind):
    x, u = rows[:, 0], rows[:, 1]
    if kind == 'additive':
        return np.column_stack([x, u, np.ones_like(x)])
    if kind == 'bilinear':
        return np.column_stack([x, u, x*u, np.ones_like(x)])
    raise ValueError(kind)


def fit(rows, targets, kind):
    X = design(rows, kind)
    return np.linalg.lstsq(X, targets, rcond=None)[0]


def predict(x, u, coef, kind):
    return float((design(np.array([[x, u]]), kind) @ coef).item())


def evaluate(starts, controls, coef, kind):
    errors = []
    for x0, us in zip(starts, controls):
        x = x0
        for u in us:
            y = plant(x, u)
            x = predict(x, u, coef, kind)
            errors.append((x-y)**2)
    return float(np.mean(errors)), float(np.sqrt(np.mean(errors)))


def main():
    # Core math sanity check: matrix form equals the explicitly expanded update.
    diffs = []
    for _ in range(1000):
        vals = RNG.normal(size=9)
        diffs.append(np.max(np.abs(paper_matrix_step(*vals)[0] - paper_matrix_step(*vals)[1])))
    identity_error = max(diffs)

    rows, targets, starts, controls = make_data()
    # Hold out complete trajectories, preserving a genuine rollout test.
    n = len(starts)
    split = int(0.75 * n)
    train_rows = []
    train_targets = []
    for i in range(split):
        x = starts[i]
        for u in controls[i]:
            train_rows.append((x, u))
            x = plant(x, u)
            train_targets.append(x)
    train_rows = np.asarray(train_rows)
    train_targets = np.asarray(train_targets)
    test_starts, test_controls = starts[split:], controls[split:]

    results = {'identity_max_abs_error': identity_error, 'seed': SEED,
               'train_sequences': split, 'test_sequences': n-split, 'horizon': controls.shape[1]}
    for kind in ('additive', 'bilinear'):
        coef = fit(train_rows, train_targets, kind)
        one_step = float(np.mean((design(train_rows, kind) @ coef - train_targets)**2))
        mse, rmse = evaluate(test_starts, test_controls, coef, kind)
        results[kind] = {'one_step_mse': one_step, 'rollout_mse': mse,
                         'rollout_rmse': rmse, 'parameter_count': int(len(coef))}
        results[kind]['coefficients'] = coef.tolist()

    Path('results.json').write_text(json.dumps(results, indent=2))
    print(json.dumps(results, indent=2))


if __name__ == '__main__':
    main()