import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

if len(sys.argv) < 2:
    sys.exit("Use it like the following: python visualization.py positions.csv (out.mp4)")

df = pd.read_csv(sys.argv[1])
frames = sorted(df["time"].unique())

# Replace p.e with direction-of-next-hop: where the ped actually moves to next frame.
# Falls back to p.e for the last frame and for teleport jumps.
df = df.sort_values(["id", "time"]).reset_index(drop=True)
df["dx"] = df.groupby("id")["x"].shift(-1) - df["x"]
df["dy"] = df.groupby("id")["y"].shift(-1) - df["y"]
mag = np.sqrt(df["dx"]**2 + df["dy"]**2)
TELEPORT_THRESHOLD = 5.0  # any per-step jump larger than this is the teleport, not motion
fallback = (mag > TELEPORT_THRESHOLD) | mag.isna() | (mag < 1e-9)
df.loc[fallback, "dx"] = df.loc[fallback, "p.e.x"]
df.loc[fallback, "dy"] = df.loc[fallback, "p.e.y"]
mag = np.sqrt(df["dx"]**2 + df["dy"]**2).replace(0, 1.0)
df["dx"] = df["dx"] / mag
df["dy"] = df["dy"] / mag

fig, ax = plt.subplots(figsize=(10, 2.2))
ax.set_xlim(0, 50); ax.set_ylim(0, 10)
ax.set_aspect("equal"); ax.set_facecolor("white")
ax.set_xticks([]); ax.set_yticks([])
for s in ("top", "right", "left", "bottom"):
    ax.spines[s].set_color("gray")

s_white = ax.scatter([], [], s=30, facecolors="black", edgecolor="black", zorder=2)
s_black = ax.scatter([], [], s=30, facecolors="white", edgecolor="black", zorder=2)
title = ax.set_title("")

# Keep a reference to quiver objects so we can remove them each frame
quivers = []

def update(t):
    sub = df[df["time"] == t]
    w = sub[sub["color"] == 1]
    b = sub[sub["color"] == 0]

    s_white.set_offsets(w[["x", "y"]].values)
    s_black.set_offsets(b[["x", "y"]].values)

    # Remove old quiver arrows
    for q in quivers:
        q.remove()
    quivers.clear()

    # Draw fresh arrows each frame
    arrow_scale = 1.5  # Tune this to change arrow length visually
    if len(w):
        q_white = ax.quiver(w["x"], w["y"], w["dx"], w["dy"],
                            color="black", angles="xy", scale_units="xy",
                            scale=1/arrow_scale, width=0.003,
                            headwidth=4, headlength=5, zorder=3)
        quivers.append(q_white)
    if len(b):
        q_black = ax.quiver(b["x"], b["y"], b["dx"], b["dy"],
                            color="gray", angles="xy", scale_units="xy",
                            scale=1/arrow_scale, width=0.003,
                            headwidth=4, headlength=5, zorder=3)
        quivers.append(q_black)

    title.set_text(f"t = {t}")
    return s_white, s_black, title

anim = FuncAnimation(fig, update, frames=frames, interval=50, blit=False)

if len(sys.argv) >= 3:
    anim.save(sys.argv[2], writer="ffmpeg", fps=5)
    print(f"Saved {sys.argv[2]}")
else:
    plt.show()