#!/usr/bin/env python3
# Copyright 2021 - Barcelona Supercomputing Center
# Author: Rodrigo Martín Posada
# MIT License

import matplotlib.pyplot as plt
from matplotlib import animation
import netCDF4
import os
import sys

if len(sys.argv) != 3:
    raise Exception(
        'Invalid args. Usage: ./nc_to_gif.py <work_dir> <file>. Try ./nc_to_gif.py work_python_4_1_1990-01-01_1991-01-01_0 A_Ice_frac')

dir_path = os.path.dirname(os.path.realpath(__file__))
work_path = sys.argv[1]
work_path_abs = os.path.join(dir_path, work_path)
file = sys.argv[2]

nc = netCDF4.Dataset('%s/%s.nc' % (work_path_abs, file))  # pylint: disable=no-member

topo = nc.variables[file]
frames = len(topo)

# make image
fig = plt.figure()
ax = plt.axes()
time_text = ax.text(0.95, 0.07, '', horizontalalignment='center', verticalalignment='top', transform=ax.transAxes, color='white')
im = ax.imshow(topo[0])

def init():
    time_text.set_text(0)
    im.set_data(topo[0])
    return [im]


def animate(i):
    time_text.set_text(i)
    im.set_array(topo[i])
    return [im]


anim = animation.FuncAnimation(fig, animate, init_func=init, frames=frames, interval=20, blit=True)
anim.save('%s/%s.gif' % (work_path_abs, file), fps=60)
