import os, numpy as np, optparse
from acme_diags.parameter.core_parameter import CoreParameter
from acme_diags.run import runner

# start interactive job:
# salloc --nodes=1 --partition=regular --time=02:00:00 -C haswell

# E3SM unified environment
# source /global/cfs/cdirs/e3sm/software/anaconda_envs/load_latest_e3sm_unified.sh

# ------------------------------------------------------------------------------
# Parse the command line options
parser = optparse.OptionParser()
parser.add_option('--case',dest='case',default=None,help='Sets the case name to run diagnostics for')
# parser.add_option('--ref',dest='ref_case',default=None,help='Sets the reference case, defaults to obs if not set')
(opts, args) = parser.parse_args()
# ------------------------------------------------------------------------------

param = CoreParameter()

# param.debug = True

if opts.case is not None: test_case = opts.case

# test_case = 'E3SM.ne30pg2_r05_oECv3.F-MMF1.01'
# test_case = 'E3SM.ne30pg2_r05_oECv3.F-MMF1.CRM-AC.RAD-AC.01'
# test_case = 'E3SM.PGVAL.ne30_r05_oECv3.F2010SC5-CMIP6.master-cbe53b'

#--------------------------------
# RGMA
#-------------------------------
### MMF vs Obs
if test_case=='RGMA_LR_vs_Obs':
  test_case = 'E3SM.RGMA.ne30pg2_r05_oECv3.FC5AV1C-L.00'
  param.test_name, param.reference_name, param.diff_title = 'E3SM-LR','Obs','E3SM-LR - Obs'
  param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/RGMA/{param.test_name}_vs_Obs'
  param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology/'
  param.run_type = 'model_vs_obs'

### MMF vs LR
if test_case=='RGMA_MMF_vs_LR':
  test_case = 'E3SM.RGMA.ne30pg2_r05_oECv3.F-MMF1.CRMNX_64.CRMDX_2000.RADNX_4.00'
  ref_case = 'E3SM.RGMA.ne30pg2_r05_oECv3.FC5AV1C-L.00'
  param.test_name, param.reference_name, param.diff_title = 'E3SM-MMF','E3SM-LR','E3SM-MMF - E3SM-LR'
  param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/RGMA/{param.test_name}_vs_LR'
  param.reference_data_path = f'/global/cscratch1/sd/whannah/e3sm_scratch/cori-knl/{ref_case}/clim/'
  param.run_type = 'model_vs_model'

### MMF vs Obs
if test_case=='RGMA_MMF_vs_Obs':
  test_case = 'E3SM.RGMA.ne30pg2_r05_oECv3.F-MMF1.CRMNX_64.CRMDX_2000.RADNX_4.00'
  param.test_name, param.reference_name, param.diff_title = 'E3SM-MMF','Obs','E3SM-MMF - Obs'
  param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/RGMA/{param.test_name}_vs_Obs'
  param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology/'
  param.run_type = 'model_vs_obs'

### HR vs LR
if test_case=='RGMA_HR_vs_LR':
  test_case = 'E3SM.RGMA.ne120pg2_r05_oECv3.FC5AV1C-H01A.00'
  ref_case = 'E3SM.RGMA.ne30pg2_r05_oECv3.FC5AV1C-L.00'
  param.test_name, param.reference_name, param.diff_title = 'E3SM-HR','E3SM-LR','E3SM-HR - E3SM-LR'
  param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/RGMA/{param.test_name}_vs_LR'
  param.reference_data_path = f'/global/cscratch1/sd/whannah/e3sm_scratch/cori-knl/{ref_case}/clim/'
  param.run_type = 'model_vs_model'

### HR vs Obs
if test_case=='RGMA_HR_vs_Obs':
  test_case = 'E3SM.RGMA.ne120pg2_r05_oECv3.FC5AV1C-H01A.00'
  param.test_name, param.reference_name, param.diff_title = 'E3SM-HR','Obs','E3SM-HR - Obs'
  param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/RGMA/{param.test_name}_vs_Obs'
  param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology/'
  param.run_type = 'model_vs_obs'

param.test_name = test_case

#--------------------------------
# Physgrid validation
#--------------------------------

# param.short_ref_name  = 'maint1-4ad402'
# param.short_test_name = 'master-12cf5f'
# param.diff_title = f'{param.short_test_name} - {param.short_ref_name}'


# if opts.ref_case is None:
#   param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology/'
# else:
#   param.reference_data_path = f'/global/cscratch1/sd/whannah/e3sm_scratch/cori-knl/{opts.ref_case}/clim/'

# param.reference_data_path = '/global/cfs/cdirs/e3sm/acme_diags/obs_for_e3sm_diags/climatology/'
# ref_case = 'E3SM.PGVAL.ne30_r05_oECv3.F2010SC5-CMIP6.master-cbe53b'
# ref_case = 'E3SM.PGVAL.conusx4v1_r05_oECv3.F2010SC5-CMIP6.master-cbe53b'

#--------------------------------
#--------------------------------


param.test_data_path      = f'/global/cscratch1/sd/whannah/e3sm_scratch/cori-knl/{test_case}/clim/'


# all seasons + annual will run if param.seasons not set
param.seasons = ['ANN','DJF','JJA']

# if 'PGVAL' in opts.case:
#   results_root = '/global/cfs/cdirs/e3sm/www/whannah/physgrid_validation_diags'
#   if 'obs_for_e3sm' in param.reference_data_path:
#     param.results_dir = f'{results_root}/{param.test_name}_vs_obs'
#   else:
#     param.results_dir = f'{results_root}/{param.test_name}_vs_np4'
#     param.run_type = 'model_vs_model'
#     param.reference_title = ref_case
#     if 'ne30pg2_' in param.test_name: param.diff_title = 'pg2 - np4'
#     if 'ne30pg3_' in param.test_name: param.diff_title = 'pg3 - np4'
#     if 'ne30pg4_' in param.test_name: param.diff_title = 'pg4 - np4'
# else:
#   param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/{param.test_name}'

# param.run_type = 'model_vs_model'

# param.results_dir = f'/global/cscratch1/sd/whannah/diags_tmp/{param.test_name}'


# Use the following if running in parallel:
param.multiprocessing = True
param.num_workers = 30

# param.sets = ['lat_lon']
param.sets = ['lat_lon','zonal_mean_xy','zonal_mean_2d','polar']


# Find max char length of param attributes
fmt_key_len = 0
for key in param.__dict__.keys(): fmt_key_len = np.max([fmt_key_len,len(key)])
# Print out all attributes of the params object
str_out = ''
for key in param.__dict__.keys(): 
  attribute = getattr(param,key)
  if attribute!='' and not isinstance(attribute,dict) :
      str_out += f'  {key:{fmt_key_len}} : {attribute}\n'
  # if isinstance(attribute,dict):
  #   str_out += f'  {key:{fmt_key_len}}:\n'
  #   for k in attribute.keys(): 
  #     str_out += ' '*(fmt_key_len+4)+f'{k:8}  {attribute[k]}\n'
print(f'\n{str_out}\n')


runner.sets_to_run = param.sets
runner.run_diags([param])