import sys
import matplotlib
matplotlib.use('QtAgg')
from PyQt5 import QtWidgets
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
from jupedsim.internal.notebook_utils import read_sqlite_file

def visualize(trajectory_file):
    traj_data,walkable_area = read_sqlite_file(trajectory_file)
    df = traj_data.data
    fig,ax = plt.subplots(figsize=(12,8))
    if hasattr(walkable_area,"polygon"):
        x,y = walkable_area.polygon.exterior.xy
        ax.plot(x,y,'k-',alpha=0.3)
    scatter = ax.scatter([],[],s=20,c="blue",alpha=0.7)
    frame_text = ax.text(0.02,0.98,'',transform=ax.transAxes)
    frames=sorted(df["frame"].unique())
    frames=frames[::5]
    
    def update(frame):
        frame_data = df[df["frame"]==frame]
        if len(frame_data)>0:
            positions=np.column_stack(
                [frame_data["x"].values,frame_data["y"].values])
            scatter.set_offsets(positions)
        frame_text.set_text(f"Frame: {frame}, Agents: {len(frame_data)}")
        return scatter, frame_text

    anim = FuncAnimation(fig,update,frames=frames,blit=False,interval=1)
    app = QtWidgets.QApplication(sys.argv)
    plt.show(block=False)
    sys.exit(app.exec())

if __name__ == "__main__":

    file_1 = "SRS_evac_2025-11-23_23-03-09.sqlite"
    file_2 = "SRS_evac_2025-11-24_01-03-06.sqlite"
    file_3 = "SRS_evac_2025-11-30_19-43-15.sqlite" 
    file_4 = "crosswalk_sim.sqlite"
    file_5 = "_SRS_sim_232009_.sqlite"
    file_6 = "_SRS_sim_231809_.sqlite"
    file_7 = "_SRS_sim_231233_.sqlite"
    file_8 = "_SRS_sim_010542_.sqlite"
    file_9 = "_SRS_sim_013834_.sqlite"
    visualize(file_3)