import os
from e3sm_diags.parameter.core_parameter import CoreParameter
from e3sm_diags.run import runner

param = CoreParameter()

## observational data
# param.reference_data_path = '/global/cfs/cdirs/e3sm/e3sm_diags/obs_for_e3sm_diags/climatology/'
param.reference_data_path = '/global/cfs/cdirs/e3sm/e3sm_diags/postprocessed_e3sm_v1_data_for_e3sm_diags/'
param.test_data_path = '/global/cfs/cdirs/e3sm/e3sm_diags/test_model_data_for_e3sm_diags/'
#param.test_name = '20161118.beta0.FC5COSP.ne30_ne30.edison'
param.test_name = '20180215.DECKv1b_H1.ne30_oEC.edison'
# param.seasons = ["ANN"]   
param.seasons = ["ANN", "DJF", "MAM", "JJA", "SON"]

prefix = '/global/cfs/cdirs/e3sm/www/cowherd/deckv1b/'
param.results_dir = os.path.join(prefix, 'june1')
# Use the following if running in parallel:
param.multiprocessing = True
param.num_workers = 32
'''
def runoff_calc(QOVER, SNOW_FROM_ATM, RAIN_FROM_ATM):
    """ RUNOFF_RATIO"""
    var =  QOVER/ (SNOW_FROM_ATM + RAIN_FROM_ATM) 
    var.units = "dimensionless"
    var.long_name = "runoff ratio"
    return var

derived_variables = {
    'RUNOFF_RATIO': {
        ('QOVER', 'SNOW_FROM_ATM','RAIN_FROM_ATM'): runoff_calc
    }
}
'''

# variables = ["QOVER","SNOW","H2OSNO","H2OCAN","SNOCAN","RAIN","TBOT"]

# Use below to run all core sets of diags:
#runner.sets_to_run = ['lat_lon','zonal_mean_xy', 'zonal_mean_2d', 'polar', 'cosp_histogram', 'meridional_mean_2d']
# runner.sets_to_run = ['lat_lon' , 'zonal_mean_xy']
# Use below to run lat_lon map only:
# runner.sets_to_run = ['lat_lon']
runner.run_diags([param])