import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
import tempfile
from matplotlib.patches import Circle


def create_dummy_animation():
    fig, ax = plt.subplots(figsize=(7, 7))
    xdata, ydata = [], []
    ln, = plt.plot([], [], 'b-', animated=True)

    def init():
        ax.set_xlim(0, 2*np.pi)
        ax.set_ylim(-1.1, 1.1)
        return ln,

    def update(frame):
        xdata.append(frame)
        ydata.append(np.sin(frame))
        ln.set_data(xdata, ydata)
        return ln,

    ani = animation.FuncAnimation(
        fig, update, frames=np.linspace(0, 2*np.pi, 100),
        init_func=init, blit=True, repeat=False
    )

    # Save to MP4
    temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    ani.save(temp_video.name, writer='ffmpeg', fps=20)
    plt.close(fig)

    return temp_video.name


def create_animation():
    path = [(i,i) for i in range(50)]
    targets_x = [20, 80, 80, 20]
    targets_y = [20, 20, 80, 80]
    RADIUS_COVERAGE = 10
    fig, ax = plt.subplots(figsize=(7, 7))

    # agent
    ln1, = plt.plot([path[0][0]], [path[0][1]], marker='o', color='b', alpha=0.5, linewidth=5, markersize=15)

    # targets
    ln2, = plt.plot(targets_x, targets_y, marker='X', color='orange', alpha=0.5, linestyle='none', markersize=15)
    for t_x, t_y in zip(targets_x, targets_y):
        circle = Circle((t_x, t_y), RADIUS_COVERAGE, color='orange', fill=True, alpha=0.3)
        ax.add_patch(circle)

    def init():
        ax.set_xlim([0, 100])
        ax.set_ylim([0, 100])
        ax.set_title(f'Warehouse Env', fontweight="bold", size=10)
        return ln1,

    def update(frame):
        # for each frame, update the data stored on each artist.
        x = [path[frame][0]]
        y = [path[frame][1]]

        ln1.set_data(x, y)
        return ln1,

    ani = animation.FuncAnimation(fig, update, frames=40,
                                  init_func=init, blit=True, repeat=False)
    # plt.show()

    # Save to MP4
    temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
    ani.save(temp_video.name, writer='ffmpeg', fps=20)
    plt.close(fig)
    return temp_video.name


def main():
    create_animation()


if __name__ == '__main__':
    main()