#!/bin/env python
#
#  Collection of python functions to extract variables from CESM history files.
#
#  TODO list: 
#     add support for daily files
#

import sys
import subprocess
import os
import glob

def driver(base,case,midfix,vars,ys,ye,outdir):
   infiles = []
   for year in range(int(ys), int(ye)+1): # ye+1 is not included in range
      print('year = '+format(year,'04d'))
      infiles_year = (glob.glob(base+'/'+case+midfix+format(year,'04d')+'*.nc'))
      if (len(infiles_year) != 12):
         print(infiles_year)
         sys.exit('ERROR: not enough monthly files found for year = '+str(year))
      infiles.extend(infiles_year)

   infiles = sorted(infiles) # sort by month

   for var in vars:
      monfile=outdir+"/"+var+"_"+case+'_moncat.nc'
      if not os.path.isfile(monfile):
         concatenate_var_cdo(var,infiles,monfile)

      mean1=outdir+"/"+var+"_"+case+'_ymonmean.nc'
      mean2=outdir+"/"+var+"_"+case+'_timmean.nc'
      mean3=outdir+"/"+var+"_"+case+'_yearmean.nc'

      if not os.path.isfile(mean1):
         arg = ['cdo','ymonmean',monfile,mean1]
         subprocess.check_call(arg,stderr=subprocess.STDOUT)
         print("INFO: created "+mean1)

      if not os.path.isfile(mean2):
         arg = ['cdo','timmean',monfile,mean2]
         subprocess.check_call(arg,stderr=subprocess.STDOUT)
         print("INFO: created "+mean2)

      if not os.path.isfile(mean3):
         arg = ['cdo','yearmean',monfile,mean3]
         subprocess.check_call(arg,stderr=subprocess.STDOUT)
         print("INFO: created "+mean3)

def concatenate_var_cdo(var, infiles, outfile):
   print(var)

   arg = ['cdo','-L','-cat']
   for file in infiles:
      arg.append('-selvar,'+var)
      arg.append(file)
   arg.append(outfile)

   subprocess.check_call(arg,stderr=subprocess.STDOUT)
   #subprocess.call(arg,stderr=subprocess.STDOUT)
   print("INFO: created "+outfile)


if __name__ == "__main__":

   # general
   case = os.path.expandvars('$CASE')
   outdir = os.path.expandvars('$OUTDIR')
   ys = os.path.expandvars('$YS')
   ye = os.path.expandvars('$YE')

   # CLM
   clm_base = os.path.expandvars('$CLM_BASE')
   clm_vars = os.path.expandvars('$CLM_VARS').split()
   driver(clm_base,case,'.clm2.h0.',clm_vars,ys,ye,outdir)

   # CAM
   cam_base = os.path.expandvars('$CAM_BASE')
   cam_vars = os.path.expandvars('$CAM_VARS').split()
   #driver(cam_base,case,'.cam.h0.',cam_vars,ys,ye,outdir)