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

if __name__ == '__main__':
    
    # Initialize OGGM and set up the default run parameters
    cfg.initialize(logging_level='WORKFLOW', future=True)

    cfg.PARAMS['border'] = 80

    # 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)
    
    # Some params 
    cfg.PARAMS['continue_on_error'] = True
    cfg.PARAMS['cfl_number'] = 0.01  # to be a bit more stable
    cfg.PARAMS['prcp_scaling_factor'] = 1.6  # because that's how we calibrated it

    # This is now switched off so we turn it on
    cfg.PARAMS['use_kcalving_for_inversion'] = True
    cfg.PARAMS['use_kcalving_for_run'] = True

    # fraction of the original (non-calving) mu* for this glacier that
    # you are ready to allow for. Totally arbitrary.
    cfg.PARAMS['calving_min_mu_star_frac'] = 0.7
    # calving constant of proportionality k after Oerlemans and Nick (2005)
    # units yr-1. This one is for the ice thickness inversion
    # Oerlemans and Nick (2005) use 2.4 yr-1, but qualitative tests and
    # Recinos et al., (2019) indicate that is should be much smaller.
    # We set it to 0.6 according to Recinos et al 2019 for a start
    cfg.PARAMS['inversion_calving_k'] = 0.6
    # And this one is for the forward model
    cfg.PARAMS['calving_k'] = 0.6
    
    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))    
    
    # Since we will recalibrate the mass-balance, we need the tstars
    ref_tstars_base_url = 'https://cluster.klima.uni-bremen.de/~oggm/ref_mb_params/RGIV62/ERA5/elev_bands/qc3/pcp1.6/'
    workflow.download_ref_tstars(base_url=ref_tstars_base_url)
    
    # RGI glaciers
    rgi_ids = gpd.read_file(utils.get_rgi_region_file(rgi_reg))

    # 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/ERA5/elev_bands/qc3/pcp1.6/match_geod/'
    gdirs = workflow.init_glacier_directories(rgi_ids, from_prepro_level=3, prepro_base_url=base_url)

    # Ice thick have been calibrated without calving - we recalibrate with calving this time - by matching the consensus?
    workflow.calibrate_inversion_from_consensus(gdirs,
                                                apply_fs_on_mismatch=True,
                                                error_on_mismatch=False)
    
    # Do we want to match geodetic estimates?
    # This affects only the bias so we can actually do this *after*
    # the inversion, but we really want to take calving into account here
    opath = os.path.join(OUTPUT_DIR, 'fixed_geometry_mass_balance_before_match_{}.csv'.format(rgi_reg))
    utils.compile_fixed_geometry_mass_balance(gdirs, path=opath)
    workflow.match_regional_geodetic_mb(gdirs, rgi_reg, dataset='hugonnet')
    
    workflow.execute_entity_task(tasks.init_present_time_glacier, gdirs)

    
    # Get end date. The first gdir might have blown up, try some others
    i = 0
    while True:
        if i >= len(gdirs):
            raise RuntimeError('Found no valid glaciers!')
        try:
            y0 = gdirs[i].get_climate_info()['baseline_hydro_yr_0']
            # One adds 1 because the run ends at the end of the year
            ye = gdirs[i].get_climate_info()['baseline_hydro_yr_1'] + 1
            break
        except BaseException:
            i += 1
    
    # OK - run
    workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                                 min_ys=y0, ye=ye,
                                 output_filesuffix='_historical')
    
    # Run output
    opath = os.path.join(OUTPUT_DIR, 'historical_run_output_{}.nc'.format(rgi_reg))
    utils.compile_run_output(gdirs, path=opath, input_filesuffix='_historical')

    # Glacier statistics we recompute here for error analysis
    opath = os.path.join(OUTPUT_DIR, 'glacier_statistics_{}.csv'.format(rgi_reg))
    utils.compile_glacier_statistics(gdirs, path=opath)
    opath = os.path.join(OUTPUT_DIR, 'fixed_geometry_mass_balance_{}.csv'.format(rgi_reg))
    utils.compile_fixed_geometry_mass_balance(gdirs, path=opath)
    
    # Add the extended files
    pf = os.path.join(OUTPUT_DIR, 'historical_run_output_{}.nc'.format(rgi_reg))
    mf = os.path.join(OUTPUT_DIR, 'fixed_geometry_mass_balance_{}.csv'.format(rgi_reg))
    # This is crucial - extending calving only possible with L3 data!!!
    sf = os.path.join(OUTPUT_DIR, 'glacier_statistics_{}.csv'.format(rgi_reg))
    opath = os.path.join(OUTPUT_DIR, 'historical_run_output_extended_{}.nc'.format(rgi_reg))
    utils.extend_past_climate_run(past_run_file=pf,
                                  fixed_geometry_mb_file=mf,
                                  glacier_statistics_file=sf,
                                  path=opath)
    
    log.workflow('OGGM Done')