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

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
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_1990-01-01_1991-01-01_0_L128_true 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.86, 0.06, 'Day: 0', horizontalalignment='left',
                    verticalalignment='top', transform=ax.transAxes, color='white')
im = ax.imshow(topo[0])
fig.tight_layout()

def init():
    time_text.set_text('Day: ' + str(0))
    im.set_data(topo[0])
    return [im]


def animate(i):
    time_text.set_text('Day: ' + str(i))
    im.set_array(topo[i])
    return [im]


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

optimized_gif_path = '%s/%s_optimized.gif' % (work_path_abs, file)
os.system('gifsicle -O2 --colors 256 %s > %s' % (gif_path, optimized_gif_path))
