#!/usr/bin/env python3

import netCDF4
import numpy as np
import matplotlib.pyplot as plt
from scipy import ndimage
from numpy import ma
import sys

if len(sys.argv) != 2: sys.exit("Required argument: ADT file")

f= netCDF4.Dataset(sys.argv[1])

longitude = f.variables["longitude"][:]
latitude = f.variables["latitude"][:]
ssh = f.variables["adt"][:].squeeze()

# No extremum on borders:

my_max = ndimage.maximum_filter(ssh, size = 3, mode = "nearest") == ssh
my_max[0, :] = False
my_max[-1, :] = False
my_max[:,0] = False
my_max[:,-1] = False

my_min = ndimage.minimum_filter(ssh, size = 3, mode = "nearest") == ssh
my_min[0, :] = False
my_min[-1, :] = False
my_min[:,0] = False
my_min[:,-1] = False

ilat_max, ilon_max = np.where(my_max)
ilat_min, ilon_min = np.where(my_min)

print("Maxima:")
print("(longitude index, latitude index) (1-based) -- (longitude, latitude):")
print("Values around maxima\n")
for j,i in np.argwhere(my_max):
    print("({}, {}) -- ({}, {})".format(i + 1, j + 1, longitude[i],
                                        latitude[j]))
    print(ssh[j-1:j+2, i-1:i+2], "\n")

print("Minima:")
print("(longitude index, latitude index) (1-based) -- (longitude, latitude):")
print("Values around minima\n")
for j,i in np.argwhere(my_min):
    print("({}, {}) -- ({}, {})".format(i + 1, j + 1, longitude[i],
                                        latitude[j]))
    print(ssh[j-1:j+2, i-1:i+2], "\n")

f.close()
choice = input("Compare extr_map (y/n)? ")
if choice == "y":
    print("From the NetCDF file ((index of latitude, index of longitude) for "
          "each extremum):")
    with netCDF4.Dataset("test_local_extrema.nc") as f:
        set_extr_map = {tuple(x) for x in
                        np.argwhere(f.variables["extr_map"][:] != 0)}
    extr_new = np.vstack((np.column_stack((ilat_max, ilon_max)),
                          np.column_stack((ilat_min, ilon_min))))
    set_new = {tuple(x) for x in extr_new}
    print("Same extrema:", set_new == set_extr_map)

lon_edge = np.hstack((longitude - 0.125, longitude[-1] + 0.125))
lat_edge = np.hstack((latitude - 0.125, latitude[-1] + 0.125))

lon_2d, lat_2d = np.meshgrid(longitude, latitude)

plt.figure()
plt.pcolormesh(lon_edge, lat_edge, ssh)
plt.colorbar()

plt.figure()
c = plt.contour(longitude, latitude, ssh)
plt.clabel(c)
plt.scatter(longitude[ilon_max], latitude[ilat_max], color = "red")
plt.scatter(longitude[ilon_min], latitude[ilat_min], color = "blue")
plt.plot(lon_2d.reshape(-1), lat_2d.reshape(-1), "+", color="gray")
plt.xlabel("longitude (in degrees east)")
plt.ylabel("latitude (in degrees north)")
plt.title("local extrema")

plt.show()