average_he

#!/usr/bin/env python

from  __future__ import print_function
from argparse import ArgumentParser
import glob
import subprocess as sp
import numpy as np
import re
import sys
import os
import json
import flo_utils as fu
from netCDF4 import Dataset as nc
# import math

ncea = "ncea"

js_config = json.load(open(os.path.expanduser('~')+"/MPI/offsets.json"))

def average_files(infiles, outfile):
#  infiles = [ "../"+x for x in infiles ] # for use with ln -s
  print ([ncea] + infiles + [outfile])
  fu.qo ([ncea] + infiles + [outfile])

def generate_atm_name(run_id, offset, timestep):
  file_years = round ((offset + timestep)/100)*10
  return ("%s/ECHAM/OUTPUT/10yr/%sATM_%d-9_lm.nc"%(run_id, run_id, file_years))

def generate_bot_name(run_id, offset, timestep):
  file_years = round ((offset + timestep)/100)*10
  return ("%s/ECHAM/OUTPUT/10yr/%sBOT_%d-9_lm.nc"%(run_id, run_id, file_years))

def generate_mpiom_name(run_id, offset, timestep):
  file_years = round ((offset + timestep)/100)*10
  pattern = "%s/MPIOM/OUTPUT/10_year_means/3d_year_%d?-90_lm.nc"%(run_id, file_years)
  files = (glob.glob(pattern))
  if (len(files) == 1):
    return files[0]
  else:
    fu.cerr("ERROR GENERATING MPIOM NAMES! RECEIVED WRONG COUNT OF MATCHES FOR THE PATTERN")
    fu.cerr(pattern)
    fu.cerr (files)
    sys.exit(666)

def generate_mpiom_surf_name(run_id, offset, timestep):
  file_years = round ((offset + timestep)/100)*10
  pattern = "%s/MPIOM/OUTPUT/10_year_means/2d_year_%d?-90_lm.nc"%(run_id, file_years)
  files =  (glob.glob(pattern))
  if (len(files) == 1):
    return files[0]
  else:
    fu.cerr("ERROR GENERATING MPIOM NAMES! RECEIVED WRONG COUNT OF MATCHES FOR THE PATTERN")
    fu.cerr(pattern)
    fu.cerr (files)
    sys.exit(666)

def generate_pism_name(run_id, offset, timestep):
  file_years = round ((offset + timestep)/10)
  return ("%s/mPISM/OUTPUT/main_fields/main_fields%d.nc"%(run_id, file_years))

def generate_mpiom_flux_name(run_id, offset, timestep):
  file_years = round ((offset + timestep)/10)
  return ("%s/mPISM/OUTPUT/fluxes/mpiom_fluxes_%d.nc"%(run_id, file_years))


def average_he(options):
  events = options.events
  experiment_ids = [ x[0:3] for x in events ]
  print (experiment_ids)
  offsets = { x : js_config["offsets"][x] for x in events }
  timesteps = np.arange(-350, 3600, options.interval)

  fu.mkdir(options.output)

  if options.PISM:
    types = ["main_fields"]
    name_generators = {"main_fields":generate_pism_name}
  elif options.mpiom_fluxes:
    types = ["mpiom_fluxes"]
    name_generators = {"mpiom_fluxes":generate_mpiom_flux_name}
  else:
    types = ["atm", "bot", "mpiom", "mpiom_surf"]
    name_generators = {"atm": generate_atm_name, "bot": generate_bot_name, "mpiom": generate_mpiom_name, "mpiom_surf": generate_mpiom_surf_name}

  for timestep in timesteps:
    for ctype in types:
      ng = name_generators[ctype]
      infiles = [ ng(run, offsets[event], timestep) for (run, event) in zip (experiment_ids, events) ]
      infiles = [ x for x in infiles if x ]
      outfile =  "%s/%s_%+05d.nc"%(options.output, ctype,  timestep)
      average_files(infiles, outfile)
      fdb = fu.debug
      if options.PISM:
        of = nc(outfile, "r+")
        of.variables["t"][0]=timestep
      if options.mpiom_fluxes:
        of = nc(outfile, "r+")
        of.variables["time"][0]=timestep
      else:
        fu.debug=True
        fu.qo(["ncatted", "-a", "average_period,global,o,c,%d-%d"%(timestep-options.interval/2, timestep+options.interval/2), outfile])
        fu.debug=fdb



def parse_args():
  parser = ArgumentParser()
  parser.description = "Average HE data"
  parser.add_argument ("events", help='First year to use in average', nargs='*')
  parser.add_argument ("-o", "--output", help='where to dump the output', required=True)
  parser.add_argument ( "--interval", help='interval at which to expect data', type=int, default=100)
  parser.add_argument ( "--PISM", help='Process PISM files', action="store_true")
  parser.add_argument ( "--mpiom_fluxes", help='Process PISM mpiom_flux files', action="store_true")
  parser.add_argument("-v", "--verbose",
                      help='''Be verbose''', action="store_true")
  parser.add_argument("--ncecat",
                      help='''Use ncecat instead of ncea''', action="store_true")

  options = parser.parse_args()
  return options


def main():
  options = parse_args()
  if options.verbose:
    print (dir(options))
  if options.ncecat:
    global ncea
    ncea = "ncecat"
  average_he(options)
if __name__ == "__main__":
  main()