import sys
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.gif)")


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

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")


# Lets define the 2 groups
# WHITE = 1 (left -> right, drawn as filled black circles, like paper)
# BLACK = 0 (right -> left, drawn as empty circles)
s_white = ax.scatter([],[], s=30, facecolors = "black", edgecolor = "black", zorder = 2)
s_black = ax.scatter([],[], s=30, facecolors = "white", edgecolor = "black", zorder = 2)



# update our graphic from t to t+1
title = ax.set_title("")
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)
    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=20)
    print(f"saved {sys.argv[2]}")
else:
    plt.show()