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.core 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

# ERA5 params
cfg.PARAMS['continue_on_error'] = True

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

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]
    
# Test
# rgi_ids = rgi_ids.sample(64)

# Go - get the pre-processed glacier directories
base_url = 'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.4/L1-L2_files/elev_bands'
gdirs = workflow.init_glacier_directories(rgi_ids, from_prepro_level=2, prepro_base_url=base_url)

from oggm.shop.millan22 import add_millan_thickness, compile_millan_statistics

workflow.execute_entity_task(add_millan_thickness, gdirs)
compile_millan_statistics(gdirs, path=f'{OUTPUT_DIR}/RGI{rgi_reg}.csv')

log.workflow('OGGM Done')