#!/usr/bin/python

from   netCDF4  import Dataset
import numpy as np
import matplotlib.pyplot as plt

# Quick Python script to display time series
# of some variable produced by perturbation
#
# Francois Massonnet
# November 2016
# francois.massonnet@bsc.es

svar = 'qlw'	# Variable to show
year = 1993	# Year to show
jy = 20		# y-coordinate of grid-point to show
jx = 20		# x- 
repo = '/esarchive/releases/fg/ocean/DFS5.2' # where data is located
nmemb=25        # Nb members to look for

# ============
# START SCRIPT
# ============

# Naming changes for two types of files (Asif decided so).
fvar = svar
if svar == 'qlw':
  fvar = 'radlw'
if svar == 'qsw':
  fvar = 'radsw'

fig = plt.figure(figsize = (10, 6))
avg = 0
javg = 1
for  m in np.arange(nmemb, -1, -1):
  print(str(m) + '/' + str(nmemb))
  filein = repo + '/' + svar + '_fc' + str(m).zfill(2) + '_DFS5.2_' + str(year) + '.nc'

  f = Dataset(filein, mode = 'r')
  var = f.variables[fvar][:, jy, jx]
  units = f.variables[fvar].units
  f.close()

  if m == 0:
    color = (0.2, 0.0, 0.0)
    lw = 2
  else:
    color = (0, 0.5, 0)
    lw = 1

  plt.plot(var, color = color, lw = lw)

  # Compute mean incrementally
  if m != 0:
    avg = avg + (var - avg) / javg
    javg = javg + 1

# Mean
plt.plot(avg, color = (0, 0.2, 0), lw = 2)
plt.xlabel('time')
plt.ylabel(units)
plt.title(fvar + ' ' + str(year) + ' ' + 'jy = ' + str(jy) + '; jx = ' + str(jx))
fig.savefig('fig.png', dpi = 300)
plt.close("all")