#! python3
########################################################################################

import netCDF4
import datetime
import matplotlib
import matplotlib.pyplot as plt
import numpy as np

########################################################################################

def read_pmc_nc(file_nc):
  nc = netCDF4.Dataset(file_nc, 'r')

  rad     = nc.variables['nc_rad'][:,:,:]
  hgt     = nc.variables['nc_hgt'][:,:,:]
  flag    = nc.variables['nc_flag'][:,:,:]
  dttxt   = nc.variables['nc_dt'][:]
  ew      = nc.variables['nc_ew'][:]
  lat     = nc.variables['nc_lat'][:]
  lon     = nc.variables['nc_lon'][:,:]
  orate   = nc.variables['nc_orate'][:,:,:]
  dhgt    = nc.variables['nc_dhgt'][:,:,:]
  datetxt = nc.variables['nc_date'][:]
  gain    = nc.variables['nc_gain'][:]
  dgain   = nc.variables['nc_dgain'][:]

  nc.close()
  return(rad,hgt,flag,dttxt,ew,lat,lon,orate,dhgt,datetxt,gain,dgain)

########################################################################################

def main(file_nc,file_png):
  #file_nc = './h08_b01_s01s02_orate.nc'
  #file_png = './h08_b01_s01s02_orate.png'
  rad,hgt,flag,dttxt,ew,lat,lon,orate,dhgt,datetxt,gain,dgain = read_pmc_nc(file_nc)

  bnd = file_nc[-18:-16]
  sg1 = file_nc[-14:-12]
  sg2 = file_nc[-11:-9]

  ccode='dummy'
  if bnd == '01':
    ccode='Blues'
    colors = ["lightsteelblue","deepskyblue","royalblue","blue","darkblue"]
  if bnd == '02':
    ccode='Greens'
    colors = ["darkseagreen","lightgreen","limegreen","green","darkgreen"]
  if bnd == '03':
    ccode='Reds'
    colors = ["rosybrown","lightcoral","tomato","red","darkred"]

  cmap= matplotlib.colors.ListedColormap(colors)
  cmap.set_under("white")
  cmap.set_over("black")
  norm= matplotlib.colors.Normalize(vmin=0,vmax=100)

  tcom = 'dummy'
  if sg1=='01' and sg2=='02':
    tcom = '(NH, Band-'+bnd+')'
    #lat = (+1)*lat
  if sg1=='09' and sg2=='10':
    tcom = '(SH, Band-'+bnd+')'
    lat = (-1)*lat

  # x
  for i in range(len(datetxt)):
    d = datetime.datetime.strptime(datetxt[i]+' 12:00:00','%Y-%m-%d %H:%M:%S')
    if i==0:
      x = np.copy(d)
    else:
      x = np.append(x,d)

  # y
  y = np.copy(lat)

  # z
  orate_w = orate[:,0,:].T
  orate_e = orate[:,1,:].T

  # grid
  xm,ym = np.meshgrid(x,y)

  #figure
  fs = 14
  fig,(ax) = plt.subplots(2,1,figsize=(12,6))
  plt.subplots_adjust(wspace=0.0)
  fmt = matplotlib.dates.DateFormatter('%H:%M\n%b %d\n%Y')
  xr = [datetime.datetime(2015,7,7,0,0,0),np.max(x)+datetime.timedelta(days=0.5)]
  yr = [np.min(y)-0.5,np.max(y)+0.5]

  ax[0].set_title(tcom,fontsize=fs)
  ax[0].set_xlim(xr[0],xr[1])
  ax[0].set_ylim(yr[1],yr[0])
  ax[0].xaxis.set_major_formatter(fmt)
  ax[0].tick_params(labelbottom=False,bottom=False)
  ax[0].set_ylabel('Lat. (East)',fontsize=fs)
  ax[0].set_yticks([45,55,65,75])
  ax[0].tick_params(labelsize=fs)

  ax[1].set_xlim(xr[0],xr[1])
  ax[1].set_ylim(yr[0],yr[1])
  ax[1].xaxis.set_major_formatter(fmt)
  ax[1].set_ylabel('Lat. (West)',fontsize=fs)
  ax[1].set_yticks([45,55,65,75])
  ax[1].tick_params(labelsize=fs)

  with np.errstate(invalid='ignore'):
    # occurrence rate
    ze = np.where((orate_e > 0), orate_e, np.nan)
    zw = np.where((orate_w > 0), orate_w, np.nan)
    cm0 = ax[0].contourf(xm,ym,ze,cmap=cmap,levels=[0,10,20,30,40,50,60,70,80,90,100],extend='both')
    cm1 = ax[1].contourf(xm,ym,zw,cmap=cmap,levels=[0,10,20,30,40,50,60,70,80,90,100],extend='both')
    cb0 = fig.colorbar(cm0,ax=ax[0],extend='both')
    cb0.set_label('Occurrence rate (%)',fontsize=fs)
    cb0.ax.tick_params(labelsize=fs)

    # no data
    ze = np.where((orate_e > -1), np.nan, 0.6)
    zw = np.where((orate_w > -1), np.nan, 0.6)
    cm0 = ax[0].contourf(xm,ym,ze,cmap='Greys',vmin=0,vmax=1,alpha=0.2)
    cm0 = ax[1].contourf(xm,ym,zw,cmap='Greys',vmin=0,vmax=1,alpha=0.2)

  ax[0].set_position([0.1,0.3,0.6,0.2])
  ax[1].set_position([0.1,0.1,0.6,0.2])
  cb0.ax.set_position([0.72,0.1,0.02,0.4])

  ax[0].grid(True)
  ax[1].grid(True)

  plt.savefig(file_png,bbox_inches="tight")
  print(file_png)

########################################################################################

if __name__ == "__main__":
  file_nc = './h08_b01_s01s02_orate.nc'
  file_png = './h08_b01_s01s02_orate.png'
  main(file_nc,file_png)

########################################################################################