#!/usr/bin/env python3

import matplotlib
import sys
from matplotlib import colors, ticker, cm
from matplotlib.backends.backend_pdf import PdfPages
from netCDF4 import Dataset
import matplotlib.pyplot as plt
import numpy as np
import math
from matplotlib import animation
import matplotlib.gridspec as gridspec
from matplotlib.lines import Line2D
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
matplotlib.rcParams['text.usetex'] = True


input_file_name  = './corners_agt.059.nc'

file_content     = Dataset(input_file_name, mode='r') # load entire netCDF file
print(file_content)
time          = file_content.variables['time'][:] # extract variable
agx           = file_content.variables['ag_x'][:,:]
agy           = file_content.variables['ag_y'][:,:]
group           = file_content.variables['ag_group'][:,:]
file_content.close() # close netCDF file

input_file_name  = './topo.txt'
f     = open(input_file_name, mode='r')
data = []
building_id = []
vertex_id = []
polygon = []
gons = []
p_last = 1
for line in f:
  line = line.strip()
  columns = line.split()
  if (int(columns[0]) > p_last):
    gons.append(polygon)
    polygon = []
  polygon.append([float(columns[2]),float(columns[3])])
  p_last = int(columns[0])
gons.append(polygon)
nop = int(columns[0])
patches = []

for i in range(nop):
  poly = Polygon(gons[i-1], True)
  patches.append(poly)
p = PatchCollection(patches,facecolors='grey')

input_file_name  = './topo_nav.txt'
f     = open(input_file_name, mode='r')
t_x = []
t_y = []

l_count = 0
for line in f:
  l_count = l_count + 1
  line = line.strip()
  columns_t = line.split()
  t_x.append(columns_t[2])
  t_y.append(columns_t[3])

Nt = np.shape(time)[0]
print(Nt)
Nt = Nt -1
Nags = np.shape(agx[1,:])
Nx = 20
Ny = 20

x = np.linspace(0.0, 20.0, Nx)
y = np.linspace(0.0, 20.0, Ny)
X, Y = np.meshgrid(x, y)

fig,ax = plt.subplots()
fig.set_size_inches(7., 7.)
xrange = [0, 20]
yrange = [0, 20]
colors = ['yellow','red','blue']


def animate(i):
    ax.clear()
    ax.add_collection(p)
    x = agx[i,:]
    y = agy[i,:]
    ax.set_xlim(*xrange)
    ax.set_ylim(*yrange)
    ax.scatter(x,y,s=3,c=group[i,:])
    ax.scatter(t_x,t_y,s=1,c='red')
    plt.title(r'simulated time: %i s' % time[i])
    k = i/Nt*100.
    print("Done: %6.2f percent\r" % k, file=sys.stdout, end=" ")
    return ax,

anim = animation.FuncAnimation(fig, animate, frames=Nt, interval=60, blit=False,save_count=0)
anim.save('ags.mp4', dpi=200, writer='ffmpeg')
#anim.save('basic_animation.gif', fps=1, extra_args=['-vcodec', 'libx264'])

#plt.show()