#!/usr/bin/env python
################################################################################## 
# Utility to produce animation frames directly from the WAD_TEST_CASE output.
# It can be used without arguments, i.e.:
#     python plotframes.py
# in which case a frame is created every tenth time level from SSH data extracted from
# WAD_1ts_00010101_00010101_grid_T.nc along the centre of the basin (j=17). A closed 
# basin is assumed and frames are named wadfr0000.png etc. Bathymetry information is 
# extracted from the mesh_mask.nc file. The frames are annotated with a timestamp that
# assumes an 18s baroclinic timestep.
# 
# All these settings can be overridden with command-line arguments. See:
#     python plotframes.py -h 
# for details. For example:
#     python plotframes.py  -nt 300 -stride 30 -froot mywad
# 
# Two major variations are also supported for specific test cases:
#     python plotframes.py -obc
# plots the right-hand side of the basin as an open boundary (test case 7) and:
#     python plotframes.py -use_sal
# colours each gridcell according to its salinity value (test case 6)
################################################################################## 
import os, sys
from argparse import ArgumentParser
import numpy as np
import netCDF4

import matplotlib.pyplot as plt
import matplotlib
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
from netCDF4 import Dataset

#
# Turn off the unhelpful warning about open figures
#
matplotlib.rcParams['figure.max_open_warning'] = 0

if __name__ == '__main__':
    parser = ArgumentParser(description=
    """
    produce frames for the animation of results from the WAD_TEST_CASES.
    Mostly this can be run without arguments but command line arguments may be
    used to override defaults. These are necessary in cases with open boundaries
    (e.g. nn_wad_test=7) and cases where it is desired to show variations in salinity
    (e.g. nn_wad_test=6).
     e.g. plotframes.py -tfile <T-grid file> -bfile  <bathymetry file>
    -froot <root name for frames>
    -nt <maximum number of time frames to process>
    -stride <stride through time frames> 
    -rdt <length of baroclinic timestep (s)> 
    -obc
    -use_sal
    """)
    parser.add_argument('-tfile',dest='tfile',help='T-grid file if not WAD_1ts_00010101_00010101_grid_T.nc', default='WAD_1ts_00010101_00010101_grid_T.nc')
    parser.add_argument('-bfile',dest='bfile',help='Bathymetry file if not mesh_mask.nc', default='mesh_mask.nc')
    parser.add_argument('-j',dest='jrow',help='jrow; j-row to extract and plot (default: 17 (fortran index))', type=int, default=16)
    parser.add_argument('-froot',dest='froot',help='froot; root name for frames (default: wadfr)', default='wadfr')
    parser.add_argument('-nt',dest='nfmax',help='nfmax; maximum number of frames to produce', type=int, default=None)
    parser.add_argument('-stride',dest='tinc',help='tinc; stride through time frames (default: 10)', type=int, default=10)
    parser.add_argument('-rdt',dest='rdt',help='rdt; length of baroclinic timestep (s) (default: 18.0)', type=float, default=18.0)
    parser.add_argument('-obc',help='Right-hand side boundary is open', action="store_true")
    parser.add_argument('-use_sal',help='colour polygons according to salinity variations', action="store_true")
    args   = parser.parse_args()
    tfile  = args.tfile
    bfile  = args.bfile
    jrow   = args.jrow
    froot  = args.froot
    nfmax  = args.nfmax
    stride = args.tinc
    rdt    = args.rdt
    obc    = args.obc
    use_sal= args.use_sal


fw = Dataset(tfile)
ssh = fw.variables['sossheig'][:,jrow,:]
vot = fw.variables['sosaline'][:,jrow,:]
if use_sal:
 sal = fw.variables['vosaline'][:,:,jrow,:]
 nz = sal.shape[1]
fw.close()

fw = Dataset(bfile)
bat = fw.variables['ht_wd'][0,jrow,:]
if use_sal:
 mbat = fw.variables['mbathy'][0,jrow,:]
fw.close()

#print "ssh"
#print ssh.shape
#print "bat"
#print bat.shape 
#print "vot"
#print vot.shape 
nt = ssh.shape[0]
nx = ssh.shape[1]
#print nx,nt

bat = -1.*bat
batmin = np.amin(bat)
batmax = np.amax(bat)
brange = batmax - batmin
tol = 0.1*brange
#print batmin,batmax,' ho'
if obc:
 batrhs = batmin
else:
 batrhs = batmax

if nfmax is None:
  nfmax = nt

nf = 0
ntmax = np.minimum(nt,nfmax*stride)

if not use_sal:
#
# plot solid single colour polygons just showing ssh variation
#
  for t in range(0,ntmax,stride):
   wadfr = froot+"{:0>4d}.png".format(nf)
   nf = nf + 1
  
   tfac = rdt/3600.0
   t24  = np.int(np.mod(t*tfac,24))
   dy   = np.int(t*tfac/24.0)
   mn   = np.int(np.rint((np.mod(t*tfac,24) - t24 )*60))
   hour  = "t={:0>2d}:{:0>2d}:{:0>2d} ".format(dy,t24,mn)
   hour2  = "  (days:hrs:mins)"
   batpts = np.zeros((nx+4,2))
   sshpts = np.zeros((2*nx,2))
   votpts = np.zeros((nx,2))
  
   for pt in range(nx):
      batpts[pt+2,0] = pt
      batpts[pt+2,1] = bat[pt]
      sshpts[pt,0] = pt
      sshpts[pt,1] = ssh[t,pt]
      votpts[pt,0] = pt
      votpts[pt,1] = np.minimum(36.,vot[t,pt])
      votpts[pt,1] = np.maximum(30.0,votpts[pt,1])
      votpts[pt,1] = batmin +0.2*brange + (votpts[pt,1]-30.)*brange/6.0
  
   batpts[nx+1,1] = batrhs
   batpts[nx+2,0] = nx-1
   batpts[nx+2,1] = batrhs
   batpts[nx+3,0] = nx-1
   batpts[nx+3,1] = batmin
   batpts[0,0] = 0.0
   batpts[0,1] = batmin
   batpts[1,0] = 0.0
   batpts[1,1] = batmax
   batpts[2,1] = batmax
   sshpts[nx-1,0] = nx-1
   sshpts[nx-1,1] = sshpts[nx-2,1]
   sshpts[0,0] = 0.0
   votpts[nx-2,0] = nx-1
   votpts[nx-1,0] = nx-1
   votpts[0,0] = 0.0
   votpts[nx-1,1] = batmax + tol
   votpts[0,1] = batmax + tol
   sshpts[0,1] = sshpts[1,1]
   for pt in range(nx):
      sshpts[pt+nx,0]=batpts[nx+1-pt,0]
      sshpts[pt+nx,1]=batpts[nx+1-pt,1]
   
  
   xs, ys = zip(*votpts)
   fig, ax = plt.subplots()
  
   patches = []
   polygon = Polygon(batpts, True)
   patches.append(polygon)
   p = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0)
   p.set_facecolors(['#f1a9a9'])
  
   patches = []
   polygon2 = Polygon(sshpts, True)
   patches.append(polygon2)
   p2 = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0)
   p2.set_facecolors(['#44a1ff'])
   
# Maximum depth set here to -10m
   ax.set_ylim([-10., 6.0])
   ax.set_xlim([0., 51.0])
   ax.add_collection(p2)
   ax.add_collection(p)
   ax.plot(xs,ys, '--', color='black', ms=10)
  
   plt.annotate(hour,xy=(2,batmin+0.1*brange))
   plt.annotate(hour2,xy=(2,batmin+0.05*brange))
   plt.savefig(wadfr)

else:
#
# plot each gridcell coloured according to its salinity value
#
  for t in range(0,ntmax,stride):
   wadfr = froot+"{:0>4d}.png".format(nf)
   nf = nf + 1

   tfac = rdt/3600.0
   t24  = np.int(np.mod(t*tfac,24))
   dy   = np.int(t*tfac/24.0)
   mn   = np.int(np.rint((np.mod(t*tfac,24) - t24 )*60))
   hour  = "t={:0>2d}:{:0>2d}:{:0>2d} ".format(dy,t24,mn)
   hour2  = "  (days:hrs:mins)"
   batpts = np.zeros((nx+4,2))
   votpts = np.zeros((nx,2))
   salpts = np.zeros((nx*nz,6,2))
   salmin = 28.
   salmax = 37.
   salrange = salmax - salmin
   faccol = np.zeros((nx*nz))
   cl = 0
   for pt in range(nx):
      batpts[pt+2,0] = pt
      batpts[pt+2,1] = bat[pt]
      votpts[pt,0] = pt
      votpts[pt,1] = np.minimum(35.,vot[t,pt])
      votpts[pt,1] = np.maximum(30.0,votpts[pt,1])
      votpts[pt,1] = batmin +0.2*brange + (votpts[pt,1]-30.)*brange/6.0
   batpts[nx+1,1] = batmax
   batpts[nx+2,0] = nx-1
   batpts[nx+2,1] = batmax
   batpts[nx+3,0] = nx-1
   batpts[nx+3,1] = batmin
   batpts[0,0] = 0.0
   batpts[0,1] = batmin
   batpts[1,0] = 0.0
   batpts[1,1] = batmax
   batpts[2,1] = batmax
   votpts[nx-2,0] = nx-1
   votpts[nx-1,0] = nx-1
   votpts[0,0] = 0.0
   votpts[nx-1,1] = batmax + tol
   votpts[0,1] = batmax + tol
  
   cl = 0
   for pt in range(nx):
      mz = np.maximum(1,mbat[pt])
      im1 = np.maximum(pt-1,1)
      ip1 = np.minimum(pt+1,nx-1)
      dz  = (ssh[t,pt] - batpts[pt+2,1] )/mz
      dz1 = 0.5*(ssh[t,pt] + ssh[t,im1] - batpts[pt+2,1] - batpts[im1+2,1] )/mz
      dz2 = 0.5*(ssh[t,pt] + ssh[t,ip1] - batpts[pt+2,1] - batpts[ip1+2,1] )/mz
      dz  = np.maximum(dz ,0.0)
      dz1 = np.maximum(dz1,0.0)
      dz2 = np.maximum(dz2,0.0)
      bat1 = 0.5*( batpts[pt+2,1] + batpts[im1+2,1] )
      bat2 = 0.5*( batpts[pt+2,1] + batpts[ip1+2,1] )
      ptm = np.maximum(pt-0.5,0.0)
      ptx = np.minimum(pt+0.5,nx-1)
      for z in range(mz):
        if ( sal[t,mz-1-z,pt] > 0.0 ):
          salpts[cl,0,0] = pt
          salpts[cl,1,0] = ptm
          salpts[cl,2,0] = ptm
          salpts[cl,3,0] = pt
          salpts[cl,4,0] = ptx
          salpts[cl,5,0] = ptx
          salpts[cl,0,1] = batpts[pt+2,1] +dz*z
          salpts[cl,1,1] = bat1 +dz1*z
          salpts[cl,2,1] = bat1 +dz1*(z+1)
          salpts[cl,3,1] = batpts[pt+2,1] +dz*(z+1)
          salpts[cl,4,1] = bat2 +dz2*(z+1)
          salpts[cl,5,1] = bat2 +dz2*z
          faccol[cl] = 100*(sal[t,mz-1-z,pt] - salmin) / salrange
          faccol[cl] = np.maximum(faccol[cl],0.0)
          faccol[cl] = np.minimum(faccol[cl],100.0)
          cl = cl + 1
   
  
   votpts2 = votpts[2:nx-4,:]
   xs, ys = zip(*votpts2)
   fig, ax = plt.subplots()
   patches = []
  
   for pt in range(cl-1):
      polygon = Polygon(salpts[pt,:,:], True)
      patches.append(polygon)
  
   p = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0)
   
   p.set_array(faccol)
   p.set_edgecolor('face')
  
   patches = []
   polygon = Polygon(batpts, True)
   patches.append(polygon)
   p2 = PatchCollection(patches, cmap=matplotlib.cm.jet, alpha=1.0)
   p2.set_facecolors(['#f1a9a9'])
  
# Maximum depth set here to -8m (suitable for test case 6 only)
   ax.set_ylim([-8., 6.0])
   ax.set_xlim([0., 51.0])
   ax.add_collection(p)
   ax.add_collection(p2)
   ax.plot(xs,ys, '--', color='black', ms=10)
  
   plt.annotate(hour,xy=(2,batmin+0.1*brange))
   plt.annotate(hour2,xy=(2,batmin+0.05*brange))
   plt.savefig(wadfr)