import os
import sys
import logging

import geopandas as gpd

import oggm.cfg as cfg
from oggm import utils, workflow, tasks
from oggm import entity_task

from func_add_terrafirma import (process_terrafirma_ukesm_data, run_with_terrafirma_data)
log = logging.getLogger(__name__)


#########
scenarios = ['hist',  # same model (UKESM1.2-LL) being forced through the CMIP historical forcing
             'up2p0', # ramp-up phase
             'up2p0-gwl2p0',  # stabilisation at 2°C
             'up2p0-gwl2p0-50y-dn2p0', # "ramp-down" from 2C stabilisation after 50 years'
             'up2p0-gwl4p0',  # stabilisation at 4C
             "up2p0-gwl4p0-50y-dn2p0",  # "ramp-down" from 4C stabilisation after 50 years'
             #"piControl"  # same model with pre-industrial forcing'
            ]
# Go - get the pre-processed glacier directories
base_url = 'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.6/L3-L5_files/2025.6/elev_bands/W5E5/per_glacier_spinup/'
rgi_version = '62'
test = False
######
rgi_reg = str(sys.argv[1])



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

cfg.PARAMS['border'] = 160 # changed for OGGM v16

# I got a strange KeyError: 'dl_verify_data_cluster.klima.uni-bremen.de' (only for RGI reg 03 and  'gfdl-esm4' 'ssp370')
cfg.PARAMS['dl_verify'] = False 
cfg.PARAMS['continue_on_error'] = True
cfg.PARAMS['use_multiprocessing']=True
cfg.PARAMS['store_model_geometry'] = True
           
# 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)


OGGM_GLACIER_JOB = os.environ.get('OGGM_GLACIER_JOB', '')
if test:
    n_glacier=32
else:
    n_glacier = 1000
id0 = (int(OGGM_GLACIER_JOB)-1)*n_glacier
id1 = (int(OGGM_GLACIER_JOB))*n_glacier

# Module logger

log.workflow(f'Starting run for RGI reg {rgi_reg}: glaciers [{id0}:{id1}]')    
log.workflow(f'using Terrafirma UKESM data, and starting at a dynamical spinup glacier state of 1979, but using the climate from the first simulation year (if ramp-up "up2p0" scenario) or from 1975 (if "hist" scenario)')
log.workflow(f'for all scenarios except "hist": The timestamp of the projections was switched to start at calendar year 1975 (but the projections use the ramp-up climate imulation years from "0" onwards...) ')
# RGI glaciers
rgi_ids = gpd.read_file(utils.get_rgi_region_file(rgi_reg, version=rgi_version))
if rgi_reg == '05':
    log.workflow('Remove connectivity 2  glaciers')
    rgi_ids = rgi_ids.loc[(rgi_ids['Connect'] == 0) | (rgi_ids['Connect'] ==1)]
rgi_ids = rgi_ids[id0:id1]    



ALL_DIAGS = ['volume', 'volume_bsl', #'volume_bwl',
             'area', 'length',  'calving', 'calving_rate',
             'off_area', 'on_area', 'melt_off_glacier',
             'melt_on_glacier', 'liq_prcp_off_glacier', 'liq_prcp_on_glacier',
             'snowfall_off_glacier', 'snowfall_on_glacier', 'model_mb',
             'residual_mb', 'snow_bucket']
# Add debug vars
cfg.PARAMS['store_diagnostic_variables'] = ALL_DIAGS

gdirs = workflow.init_glacier_directories(rgi_ids, from_prepro_level=5, prepro_border=160,
                                        prepro_base_url=base_url, prepro_rgi_version=rgi_version)
### Process terrafirma gcm data
for scen in scenarios:
    workflow.execute_entity_task(process_terrafirma_ukesm_data,
                                 gdirs, scenario= scen,
                                  y0=None, y1=None,
                                  filesuffix=None, # overwrites default filesuffix
                                  year_range=('1975','2014'),
                                  shift_timeseries_by_year_range = True)

### let's now do the actual projections
workflow.execute_entity_task(run_with_terrafirma_data,gdirs)

# compile and store run output
eq_dir = os.path.join(OUTPUT_DIR, 'RGI' + rgi_reg)
utils.mkdir(eq_dir)
for scenario in scenarios:
    rid = f'_terrafirma_UKESM1-2-LL_esm_{scenario}_bc_1975_2014'
    if scenario == 'hist':
        rid_new = rid
        utils.compile_run_output(gdirs, input_filesuffix=f'_future_run{rid}', 
                                path=os.path.join(eq_dir,                                                   f'run{rid_new}_starting_at_1979_glacier_state_from_1975_sim_year_climate_Batch_{id0}_{id1}.nc.nc')) 
    else:
        rid_new = f'_terrafirma_UKESM1-2-LL_esm_{scenario}_bc_1975_2014_to_0_40_sim_year'
        utils.compile_run_output(gdirs, input_filesuffix=f'_future_run{rid}', 
                                path=os.path.join(eq_dir,                                                  f'run{rid_new}_starting_at_1979_glacier_state_from_0_sim_year_climate_Batch_{id0}_{id1}.nc.nc')) 

log.workflow('OGGM Done')