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, flowline
import oggm_gmip3

    
# Initialize OGGM and set up the default run parameters
cfg.initialize(logging_level='ERROR')
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_model_geometry'] = False
cfg.PARAMS['hydro_month_nh'] = 1
cfg.PARAMS['hydro_month_sh'] = 1
cfg.PARAMS['prcp_scaling_factor'] = 1.6
cfg.PARAMS['store_diagnostic_variables'] =  ['volume', 'volume_bsl', 'area', 'length']

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

gcm = str(sys.argv[1])
ssp = str(sys.argv[2])

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

log.workflow(f'Pre-download data...')    
shuffle_file = 'https://raw.githubusercontent.com/GlacierMIP/GlacierMIP3/main/shuffling/shuffled_years_GlacierMIP3.csv'
utils.file_downloader(shuffle_file)

# 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, random_state=12)

# 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_pergla/'
gdirs = workflow.init_glacier_directories(rgi_ids, from_prepro_level=5, prepro_base_url=base_url, prepro_rgi_version=rgi_version)

# Climate data
clim_fsuff = f'{gcm}_{ssp}'
workflow.execute_entity_task(oggm_gmip3.process_isimip_data, gdirs,
                             filesuffix=clim_fsuff, gcm=gcm, ssp=ssp)

# Periods
if ssp == 'historical':
    periods = ['1851-1870', '1901-1920', '1951-1970', '1995-2014']
else:
    periods = ['2021-2040', '2041-2060', '2061-2080', '2081-2100']

for period in periods:

    run_fsuff = f'{gcm}_{ssp}_{period}_stop2'

    workflow.execute_entity_task(oggm_gmip3.run_isimip_climate, gdirs,
                                 period=period,
                                 climate_input_filesuffix=clim_fsuff,
                                 output_filesuffix=run_fsuff,
                                 stop_criterion=flowline.equilibrium_stop_criterion)

    # 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=run_fsuff,
                             path=os.path.join(eq_dir, run_fsuff + '.nc'))

log.workflow('OGGM Done')