import os
import logging 
import sys

# Libs
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
from oggm.shop import gcm_climate

# Initialize OGGM and set up the default run parameters
cfg.initialize(logging_level='ERROR')
rgi_version = '62'

# 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['border'] = 240
cfg.PARAMS['prcp_scaling_factor'] = 1.6
cfg.PARAMS['min_ice_thick_for_length'] = 1.
cfg.PARAMS['store_model_geometry'] = False
cfg.PARAMS['store_diagnostic_variables'] =  ['volume', 'volume_bsl', 'area']

rgi_reg = os.environ.get('OGGM_RGI_REG', '')
if rgi_reg not in ['{:02d}'.format(r) for r in range(1, 20)]:
    raise RuntimeError('Need an RGI Region')

# 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=rgi_version))

# For greenland we omit connectivity level 2
if rgi_reg == '05':
    rgi_ids = rgi_ids.loc[rgi_ids['Connect'] != 2]
    
# rgi_ids = rgi_ids.loc[:16]
    
# Go - get the pre-processed glacier directories
base_url = 'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.4/L3-L5_files/ERA5/elev_bands/qc3/pcp1.6/match_geod/'
gdirs = workflow.init_glacier_directories(rgi_ids, from_prepro_level=5, prepro_base_url=base_url, prepro_rgi_version=rgi_version)


workflow.execute_entity_task(gcm_climate.process_lmr_data, gdirs, year_range=('1980', '1999'));
workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                             ys=1,
                             climate_filename='gcm_data',  # use gcm_data, not climate_historical
                             climate_input_filesuffix='',  # use a different scenario
                             return_value=False,
                             );

utils.compile_run_output(gdirs, path=os.path.join(OUTPUT_DIR, 'RGI' + rgi_reg + '.nc'))

log.workflow('OGGM Done')