"""Lesson 09: paired observations, finite differences, and held-out rollouts."""
import argparse
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from sklearn.neighbors import KNeighborsRegressor
from common import make_dataset, prepare_embedding


def run(output_dir: Path) -> dict:
    output_dir.mkdir(parents=True, exist_ok=True)
    data = make_dataset()
    _, Z, _, _ = prepare_embedding(data)
    current = np.flatnonzero(data["next_index"] >= 0)
    following = data["next_index"][current]
    dt = data["time"][following] - data["time"][current]
    assert np.all(dt > 0)
    assert np.all(data["session"][current] == data["session"][following])
    velocity = (Z[following] - Z[current]) / dt[:, None]
    features = np.column_stack([Z[current], data["time"][current]])
    fit = data["train"][current]
    model = KNeighborsRegressor(n_neighbors=40, weights="uniform")
    model.fit(features[fit], velocity[fit])
    predicted = Z[current[~fit]] + dt[~fit, None] * model.predict(features[~fit])
    mse = float(np.mean((predicted - Z[following[~fit]]) ** 2))
    persistence = float(np.mean((Z[current[~fit]] - Z[following[~fit]]) ** 2))
    test_sessions = np.unique(data["session"][~data["train"]])
    rollout_errors, persistence_errors = [], []
    shown = None
    for sid in test_sessions:
        ids = np.flatnonzero(data["session"] == sid)
        ids = ids[np.argsort(data["time"][ids])]
        roll = [Z[ids[0]].copy()]
        for a, b in zip(ids[:-1], ids[1:]):
            h = data["time"][b] - data["time"][a]
            v = model.predict([[*roll[-1], data["time"][a]]])[0]
            roll.append(roll[-1] + h * v)
        roll = np.asarray(roll)
        rollout_errors.append(np.mean((roll[1:] - Z[ids[1:]]) ** 2))
        persistence_errors.append(np.mean((Z[ids[0]] - Z[ids[1:]]) ** 2))
        if shown is None:
            shown = (ids, roll)
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.4))
    ax[0].scatter(Z[:, 0], Z[:, 1], s=5, c=data["time"], cmap="viridis", alpha=.3)
    arrows = current[~fit][::9]
    V = model.predict(np.column_stack([Z[arrows], data["time"][arrows]]))
    ax[0].quiver(Z[arrows, 0], Z[arrows, 1], .12 * V[:, 0], .12 * V[:, 1],
                 angles="xy", scale_units="xy", scale=1, color="#dc5a32")
    ax[0].set_title("Projected field: arrows show 0.12 time units")
    ids, roll = shown
    ax[1].plot(*Z[ids].T, "o-", label="Held-out observations", markersize=3)
    ax[1].plot(*roll.T, "o--", label="Free rollout", markersize=3)
    ax[1].legend()
    ax[1].set_title("One unseen session, initial state only")
    for a in ax:
        a.set(xlabel="PC 1", ylabel="PC 2")
    fig.tight_layout()
    fig.savefig(output_dir / "lesson09.png", dpi=160)
    plt.close(fig)
    return {"one_step_mse": mse, "one_step_persistence_mse": persistence,
            "rollout_mse": float(np.mean(rollout_errors)),
            "rollout_persistence_mse": float(np.mean(persistence_errors)),
            "test_sessions": int(len(test_sessions))}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=Path("outputs"))
    args = parser.parse_args()
    metrics = run(args.output)
    print(json.dumps(metrics, indent=2, allow_nan=False))