"""Illustrative reconstruction for this portfolio, not original lab code/data.
Run: python -m pip install jax matplotlib
Then: python friction-example.py
A sliding particle on a horizontal surface; all samples precede stopping.
"""
from pathlib import Path
import jax
import jax.numpy as jnp
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Known initial speed; noise-free synthetic observations, not video measurements.
g, v0 = 9.81, 3.0
t = jnp.linspace(0.0, 0.8, 17)
true_mu = 0.25

def position(mu):
    return v0 * t - 0.5 * mu * g * t**2

observed = position(true_mu)

def loss(mu):
    return jnp.mean((position(mu) - observed)**2)

mu = jnp.array(0.08)
for _ in range(100):
    mu = mu - 0.1 * jax.grad(loss)(mu)

assert abs(float(mu) - true_mu) < 1e-5
assert float(t[-1]) < v0 / (true_mu * g)
print(f"Recovered friction: {float(mu):.4f}; MSE: {float(loss(mu)):.2e}")

fig, ax = plt.subplots(figsize=(8, 4.8), layout="constrained")
fig.patch.set_facecolor("#fffdf9")
ax.set_facecolor("#fffdf9")
ax.plot(t, position(0.08), "--", color="#74716a", label="Initial guess: 0.08", linewidth=2)
ax.plot(t, position(mu), color="#bf5700", label=f"Recovered: {float(mu):.2f}", linewidth=2.4)
ax.scatter(t, observed, facecolors="#fffdf9", edgecolors="#262723", s=32, label="Synthetic observations", zorder=3)
ax.set(xlabel="Time (s)", ylabel="Position (m)", title="Recovering friction from observed motion")
ax.spines[["top", "right"]].set_visible(False)
ax.grid(alpha=0.12)
ax.legend(frameon=False, fontsize=9)
fig.savefig(Path(__file__).with_name("friction-example.png"), dpi=180)