import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pandas as pd
import numpy as np

fig = plt.figure()

ax = plt.axes(xlim=(0, 25000), ylim=(0, 100000))

line, = ax.plot([], [], marker='o', linewidth=1, markersize=4, color='b')
ax.grid(True)
ax.title.set_text('Pressure vs Altitude')
ax.set_xlabel('Altitude(m)')
ax.set_ylabel('Pressure(mPa)')

plt.rcParams['animation.ffmpeg_path'] = 'C:/Users/dawsh/PycharmProjects/Elijah_Balloon_Space_Music/venv/Lib/site-packages/ffmpeg/bin/ffmpeg.exe'
FFwriter = animation.FFMpegFileWriter(fps=78)

file_location = 'Sample_Data_Testing_Parsing.xlsx'

sheet = pd.ExcelFile(file_location).parse(sheet_name='Parsed_Data', names=['Altitude', 'Pressure'], usecols='B,D',
                                          dtype=float, skiprows=600)

alt_read = sheet.loc[:, 'Altitude']
pressure_read = sheet.loc[:, 'Pressure']

sorted_alt = np.empty(936, dtype=float)
sorted_pressure = np.empty(936, dtype=float)


for label, content in alt_read.items():
    sorted_alt[label] = content

for label, content in pressure_read.items():
    sorted_pressure[label] = content


def init():
    line.set_data([], [])
    return line,


def animate(i):
    x = sorted_alt[i-1]
    y = sorted_pressure[i-1]
    # x = np.linspace(0, 2, 1000)
    # y = np.sin(2 * np.pi * (x-0.01*i))
    line.set_data(x, y)
    return line,


anim = animation.FuncAnimation(fig, animate, init_func=init, frames=936, interval=20, blit=True)
anim.save('pressure_animation.mp4', writer=FFwriter)
plt.show()