import os
import logging 
import sys
import json 

# Libs
import numpy as np
import xarray as xr
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt

# Locals
import oggm.cfg as cfg
from oggm import utils, workflow, tasks
    
# Initialize OGGM and set up the default run parameters
cfg.initialize(logging_level='WORKFLOW')
# cfg.initialize(logging_level='DEBUG')
rgi_version = '62'

cfg.PARAMS['border'] = 160

# Local working directory (where OGGM will write its output)
WORKING_DIR = os.environ.get('OGGM_WORKDIR', '')
if not WORKING_DIR:
    raise RuntimeError('Need a working dir')
utils.mkdir(WORKING_DIR)
cfg.PATHS['working_dir'] = WORKING_DIR

OUTPUT_DIR = os.environ.get('OGGM_OUTDIR', '')
if not OUTPUT_DIR:
    raise RuntimeError('Need an output dir')
utils.mkdir(OUTPUT_DIR)

cfg.PARAMS['continue_on_error'] = True
cfg.PARAMS['store_diagnostic_variables'] = ['volume', 'volume_bsl', 'area', 'length']
cfg.PARAMS['store_model_geometry'] = True
cfg.PARAMS['mp_processes'] = 8

# Init 
workflow.init_mp_pool(True)
    
rgi_reg = int(os.environ.get('OGGM_GLACIER_JOB', ''))

    
# Module logger
log = logging.getLogger(__name__)
log.workflow('Starting run for RGI reg {}'.format(rgi_reg))

# RGI glaciers
rgi_ids = gpd.read_file(utils.get_rgi_region_file(rgi_reg, version='62'))
# rgi_ids = rgi_ids.iloc[:2]

# Sort for more efficient mp
# rgi_ids = rgi_ids.sort_values(by='Area', ascending=False)

# Go - get the pre-processed glacier directories
base_url = 'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.6/L3-L5_files/2023.3/elev_bands/W5E5_spinup/'
gdirs = workflow.init_glacier_directories(rgi_ids, 
                                          from_prepro_level=5, 
                                          prepro_base_url=base_url, 
                                          prepro_rgi_version=rgi_version)

gcms = pd.read_csv('/home/www/oggm/cmip6/all_gcm_list.csv', index_col=0)


from oggm.shop.gcm_climate import process_lmr_data, process_modera_data

# rid = 'lmr_v21_avg'
# log.workflow(f'Bias correcting {rid}') 
# workflow.execute_entity_task(process_lmr_data, gdirs, 
#                              filesuffix='_' + rid,
#                              );

# log.workflow(f'Running {rid}') 
# workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
#                              climate_filename='gcm_data',  # use gcm_data, not climate_historical
#                              climate_input_filesuffix='_' + rid, 
#                              output_filesuffix=rid,  # recognize the run for later
#                              return_value=False,
#                              ys=1000,
#                              );
# workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
#                              output_filesuffix=rid + '_hist',  # recognize the run for later
#                              return_value=False,
#                              init_model_filesuffix=rid,
#                              );

# log.workflow(f'Aggregating {rid}')
# out_nc_path = os.path.join(OUTPUT_DIR, rid + f'_reg{rgi_reg}.nc')
# utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)          
# out_nc_path = os.path.join(OUTPUT_DIR, rid + '_hist' + f'_reg{rgi_reg}.nc')
# utils.compile_run_output(gdirs, input_filesuffix=rid + '_hist', path=out_nc_path)


# rid = 'lmr_online_avg'
# log.workflow(f'Bias correcting {rid}') 
# workflow.execute_entity_task(process_lmr_data, gdirs, 
#                              filesuffix='_' + rid,
#                              version='online',
#                              );



# log.workflow(f'Running {rid}') 
# workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
#                              climate_filename='gcm_data',  # use gcm_data, not climate_historical
#                              climate_input_filesuffix='_' + rid, 
#                              output_filesuffix=rid,  # recognize the run for later
#                              return_value=False,
#                              ys=1000,
#                              );
# workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
#                              output_filesuffix=rid + '_hist',  # recognize the run for later
#                              return_value=False,
#                              init_model_filesuffix=rid,
#                              );

# log.workflow(f'Aggregating {rid}')
# out_nc_path = os.path.join(OUTPUT_DIR, rid + f'_reg{rgi_reg}.nc')
# utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)          
# out_nc_path = os.path.join(OUTPUT_DIR, rid + '_hist' + f'_reg{rgi_reg}.nc')
# utils.compile_run_output(gdirs, input_filesuffix=rid + '_hist', path=out_nc_path)

rid = 'modera_avg'
log.workflow(f'Bias correcting {rid}') 
workflow.execute_entity_task(process_modera_data, gdirs, 
                             filesuffix='_' + rid,
                             );

log.workflow(f'Running {rid}') 
workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                             climate_filename='gcm_data',  # use gcm_data, not climate_historical
                             climate_input_filesuffix='_' + rid, 
                             output_filesuffix=rid,  # recognize the run for later
                             return_value=False,
                             ys=1421,
                             ye=2000,
                             );
workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                             output_filesuffix=rid + '_hist',  # recognize the run for later
                             return_value=False,
                             init_model_filesuffix=rid,
                             );

log.workflow(f'Aggregating {rid}')
out_nc_path = os.path.join(OUTPUT_DIR, rid + f'_reg{rgi_reg}.nc')
utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)          
out_nc_path = os.path.join(OUTPUT_DIR, rid + '_hist' + f'_reg{rgi_reg}.nc')
utils.compile_run_output(gdirs, input_filesuffix=rid + '_hist', path=out_nc_path)