import os
import logging 
import sys

# Libs
import xarray as xr
import pandas as pd
import numpy as np
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'

# That's for OGGM - VAS can use 80 when the gdirs are ready
cfg.PARAMS['border'] = 240

# 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_model_geometry'] = False
cfg.PARAMS['store_diagnostic_variables'] =  ['volume', 'volume_bsl', 'area', 'length']
cfg.PARAMS['hydro_month_nh'] = 1
cfg.PARAMS['hydro_month_sh'] = 1
cfg.PARAMS['min_mu_star'] = 20
cfg.PARAMS['max_mu_star'] = 600
cfg.PARAMS['climate_qc_months'] = 0

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')

temp_bias = float(sys.argv[1])

# Module logger
log = logging.getLogger(__name__)
log.workflow('Starting run for RGI reg {} and bias {:+.1f}'.format(rgi_reg, temp_bias))    

# RGI glaciers
rgi_ids = gpd.read_file(utils.get_rgi_region_file(rgi_reg, version=rgi_version))

# # Test
# rgi_ids = rgi_ids.iloc[:8]
# rgi_ids = ['RGI60-11.00897']

# For greenland we omit connectivity level 2
if rgi_reg == '05':
    rgi_ids = rgi_ids.loc[rgi_ids['Connect'] != 2]

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

# run constant climate scenario with different temperature biases
filesuffix = "bias{:+.1f}".format(temp_bias)

# start model run
workflow.execute_entity_task(tasks.run_random_climate, gdirs, 
                             unique_samples=True,
                             nyears=3000, y0=2009, halfsize=10,
                             temperature_bias=temp_bias,
#                              init_model_filesuffix='_historical',
                             output_filesuffix=filesuffix,
                             return_value=False)

# compile and store run output
eq_dir = os.path.join(OUTPUT_DIR, 'RGI' + rgi_reg)
utils.mkdir(eq_dir)
utils.compile_run_output(gdirs, input_filesuffix=filesuffix,
                         path=os.path.join(eq_dir, filesuffix + '.nc'))

log.workflow('OGGM Done')