# Python imports
import logging
import os

# Libs
import geopandas as gpd

# Locals
import oggm.cfg as cfg
from oggm import utils, workflow, tasks

# For timing the run
import time
start = time.time()

# Module logger
log = logging.getLogger(__name__)

# On cluster we always set this variable in the batch script
ON_CLUSTER = 'WORKDIR' in os.environ

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

if ON_CLUSTER:
    # get the options from Env vars
    rgi_reg = '{:02}'.format(int(os.environ.get('RGI_REG')))
    WORKING_DIR = os.environ["WORKDIR"]
else:
    # Hardcoded stuff
    WORKING_DIR = '/home/mowglie/disk/TMP_Data/run_test_calving/'
    rgi_reg = '01'
    cfg.PARAMS['use_multiprocessing'] = False

# How many grid points around the glacier?
cfg.PARAMS['border'] = 10

# Make it robust
cfg.PARAMS['continue_on_error'] = True

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

# RGI file
path = utils.get_rgi_region_file(rgi_reg, version=rgi_version)
rgidf = gpd.read_file(path)

# Select only tidewater glaciers
rgidf = rgidf.loc[rgidf.TermType == 1]

# For testing no need to do all
if not ON_CLUSTER:
    rgidf = rgidf.iloc[:2]

# Sort for more efficient parallel computing
rgidf = rgidf.sort_values('Area', ascending=False)

log.workflow('Starting OGGM run')
log.workflow('Number of glaciers: {}'.format(len(rgidf)))

# Go - get the pre-processed glacier directories
# Level 3 because we want to redo the inversion with calving
gdirs = workflow.init_glacier_directories(rgidf, from_prepro_level=3)

# Runs
# We start by storing the volume of glaciers before calving
utils.compile_glacier_statistics(gdirs, filesuffix='_nocalving',
                                 inversion_only=True)

# Prepare for run
workflow.execute_entity_task(tasks.init_present_time_glacier, gdirs)

# Run 1 and 2
suffix = '_tstar_nocalving'
workflow.execute_entity_task(tasks.run_random_climate, gdirs, nyears=100,
                             seed=1, y0=None,
                             bias=0, output_filesuffix=suffix)
utils.compile_run_output(gdirs, input_filesuffix=suffix)

suffix = '_commitment_nocalving'
workflow.execute_entity_task(tasks.run_random_climate, gdirs, nyears=100,
                             seed=1, y0=2000,
                             output_filesuffix=suffix)
utils.compile_run_output(gdirs, input_filesuffix=suffix)

# Redo the inversion with calving this time
workflow.execute_entity_task(tasks.find_inversion_calving, gdirs)
utils.compile_glacier_statistics(gdirs, filesuffix='_withcalving')

# Prepare for run
workflow.execute_entity_task(tasks.init_present_time_glacier, gdirs)

# Run 3 and 4
suffix = '_tstar_withcalving'
workflow.execute_entity_task(tasks.run_random_climate, gdirs, nyears=100,
                             seed=1, y0=None,
                             bias=0, output_filesuffix=suffix)
utils.compile_run_output(gdirs, input_filesuffix=suffix)

suffix = '_commitment_withcalving'
workflow.execute_entity_task(tasks.run_random_climate, gdirs, nyears=100,
                             seed=1, y0=2000,
                             output_filesuffix=suffix)
utils.compile_run_output(gdirs, input_filesuffix=suffix)

# Log
m, s = divmod(time.time() - start, 60)
h, m = divmod(m, 60)
log.workflow('OGGM is done! Time needed: %d:%02d:%02d' % (h, m, s))