import os
import numpy
from e3sm_diags.parameter.area_mean_time_series_parameter import AreaMeanTimeSeriesParameter
from e3sm_diags.parameter.core_parameter import CoreParameter
from e3sm_diags.parameter.enso_diags_parameter import EnsoDiagsParameter
from e3sm_diags.parameter.qbo_parameter import QboParameter


from e3sm_diags.run import runner

short_name = 'master-wcycle1850'
test_ts = 'ts'
start_yr = int('0001')
end_yr = int('0004')
num_years = end_yr - start_yr + 1
ref_start_yr = 1990

param = CoreParameter()

# Model
param.test_data_path = 'climo'
param.test_name = '05102023.WCYCL1850.v2setting.master.t1'
param.short_test_name = short_name

# Obs
param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology'


# Output dir
param.results_dir = 'model_vs_obs_0001-0004'

# Additional settings
param.run_type = 'model_vs_obs'
param.diff_title = 'Model - Observations'
param.output_format = ['png']
param.output_format_subplot = ['pdf']
param.multiprocessing = True
param.num_workers = 24
#param.fail_on_incomplete = True
params = [param]
enso_param = EnsoDiagsParameter()
enso_param.test_data_path = test_ts
enso_param.test_name = short_name
enso_param.test_start_yr = start_yr
enso_param.test_end_yr = end_yr

# Obs
enso_param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/time-series'
enso_param.ref_start_yr = ref_start_yr
enso_param.ref_end_yr = ref_start_yr + 10

params.append(enso_param)
qbo_param = QboParameter()
qbo_param.test_data_path = test_ts
qbo_param.test_name = short_name
qbo_param.test_start_yr = start_yr
qbo_param.test_end_yr = end_yr
qbo_param.ref_start_yr = ref_start_yr
ref_end_yr = ref_start_yr + num_years - 1
if (ref_end_yr <= 2010):
  qbo_param.ref_end_yr = ref_end_yr
else:
  qbo_param.ref_end_yr = 2010

# Obs
qbo_param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/time-series'

params.append(qbo_param)
ts_param = AreaMeanTimeSeriesParameter()
ts_param.test_data_path = test_ts
ts_param.test_name = short_name
ts_param.start_yr = start_yr
ts_param.end_yr = end_yr

# Obs
ts_param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/time-series'

params.append(ts_param)

# Run
runner.sets_to_run = ['lat_lon', 'zonal_mean_xy', 'zonal_mean_2d', 'zonal_mean_2d_stratosphere', 'polar', 'cosp_histogram', 'meridional_mean_2d', 'enso_diags', 'qbo', 'area_mean_time_series', 'annual_cycle_zonal_mean', 'aerosol_aeronet', 'aerosol_budget']
runner.run_diags(params)