from oggm import cfg
from oggm import tasks, utils, workflow, graphics
from oggm.core import massbalance
import numpy as np
import pandas as pd
import geopandas as gpd
import xarray as xr
import matplotlib.pyplot as plt
import os
import logging
import shutil

# 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')
WORKING_DIR = utils.mkdir(WORKING_DIR)

OUTPUT_DIR = os.environ.get('OGGM_OUTDIR', '')
if not OUTPUT_DIR:
    raise RuntimeError('Need an output dir')
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')
OUTPUT_DIR = os.path.join(OUTPUT_DIR, f'RGI{rgi_reg}')
utils.mkdir(OUTPUT_DIR)

# Initialize OGGM and set up the default run parameters
cfg.initialize(logging_level='INFO')
cfg.PARAMS['use_multiprocessing'] = True
cfg.PARAMS['continue_on_error'] = True

# Local working directory (where OGGM will write its output)
cfg.PATHS['working_dir'] = utils.mkdir(WORKING_DIR)

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

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

# rgi_ids = rgi_ids.sample(64, random_state=1)

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

# This is the url with minimal data
base_url = 'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.6/L3-L5_files/2023.3/elev_bands/W5E5/'

# This gets the data for this glacier
gdirs = workflow.init_glacier_directories(rgi_ids, prepro_base_url=base_url, from_prepro_level=3, prepro_border=80)

# Run the tasks
from mb_helper import binned_mb_ts
workflow.execute_entity_task(binned_mb_ts, gdirs, outdir=OUTPUT_DIR)