import os
from acme_diags.run import runner
from acme_diags.parameter.core_parameter import CoreParameter
from acme_diags.parameter.area_mean_time_series_parameter import AreaMeanTimeSeriesParameter
from acme_diags.parameter.enso_diags_parameter import EnsoDiagsParameter
from acme_diags.parameter.qbo_parameter import QboParameter
from acme_diags.parameter.diurnal_cycle_parameter import DiurnalCycleParameter
from acme_diags.parameter.streamflow_parameter import StreamflowParameter

param = CoreParameter()

input_prefix ='/global/cfs/cdirs/e3sm/acme_diags'
obs_climo = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology'
obs_ts = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/time-series'
test_prefix = '/global/cfs/cdirs/e3sm/zhang40/postprocessing_for_e3sm_diags'

#case = '20180215.DECKv1b_H1.ne30_oEC.edison'
case = '20180129.DECKv1b_piControl.ne30_oEC.edison'
casename = case.split('.')[1]+'.'+case.split('.')[2]
climo_path = test_prefix +'/climo/'+case+'/171-200/rgr'
ts_path = test_prefix +'/monthly_ts/'+case+'/171-200/rgr'
dc_climo_path = test_prefix+'/diurnal_climo/'+case+'/171-200/rgr'
param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology'
param.test_data_path = climo_path
param.test_name = case
param.test_short_name = casename
#param.seasons = ["ANN","JJA"]

prefix = '/global/cfs/cdirs/e3sm/www/zhang40/tests/'
param.results_dir = os.path.join(prefix, casename+'_171-200_v230')
#param.multiprocessing = True
#param.num_workers = 15

#Additional parameters:
#param.short_test_name = 'beta0.FC5COSP.ne30'
#param.run_type = 'model_vs_model'
#param.diff_title = 'Difference'
#param.output_format = ['png']
#param.output_format_subplot = ['pdf']
#param.save_netcdf = True


#Set specific parameters for new sets
qbo_param = QboParameter()
qbo_param.reference_data_path = obs_ts
qbo_param.test_data_path = ts_path
qbo_param.test_name = casename
qbo_param.test_start_yr = '171'
qbo_param.test_end_yr = '200'
qbo_param.ref_start_yr = '1985'
qbo_param.ref_end_yr = '2014'

dc_param = DiurnalCycleParameter()
dc_param.reference_data_path = obs_climo
dc_param.test_data_path = dc_climo_path
dc_param.test_name = case
dc_param.short_test_name = casename
dc_param.normalize_test_amp = False

#enso_param = EnsoDiagsParameter()
#enso_param.reference_data_path = obs_ts
#enso_param.test_data_path = ts_path
#enso_param.test_name = casename
#enso_param.start_yr = '171'
#enso_param.end_yr = '200'
#
#ts_param = AreaMeanTimeSeriesParameter()
#ts_param.reference_data_path = obs_ts
#ts_param.test_data_path = ts_path 
#ts_param.test_name = casename
#ts_param.start_yr = '171'
#ts_param.end_yr = '200'
#
#streamflow_param = StreamflowParameter()
#streamflow_param.reference_data_path = obs_ts
#streamflow_param.test_data_path = ts_path 
#streamflow_param.test_start_yr = '171'
#streamflow_param.test_end_yr = '200'
#streamflow_param.ref_start_yr = '1986'
#streamflow_param.ref_end_yr = '1995'
#
#runner.sets_to_run = ['lat_lon','zonal_mean_xy', 'zonal_mean_2d', 'polar', 'cosp_histogram', 'meridional_mean_2d','enso_diags', 'qbo','area_mean_time_series','enso_diags','diurnal_cycle','streamflow']
#runner.run_diags([param, qbo_param, dc_param,enso_param,ts_param,streamflow_param])
runner.sets_to_run = ['diurnal_cycle','qbo']
runner.run_diags([param, dc_param,qbo_param])
#runner.run_diags([param, qbo_param])