import os
import logging 
import sys
import json 

# Libs
import numpy as np
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
    
# Initialize OGGM and set up the default run parameters
cfg.initialize(logging_level='WORKFLOW')
# cfg.initialize(logging_level='DEBUG')
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_diagnostic_variables'] = ['volume', 'volume_bsl', 'area', 'area_min_h', 'length']
# cfg.PARAMS['store_model_geometry'] = True
cfg.PARAMS['min_ice_thick_for_area'] = 2.
cfg.PARAMS['min_ice_thick_for_length'] = 2.
cfg.PARAMS['evolution_model'] = 'FluxBased'

# Init 
workflow.init_mp_pool(True)
    
rgi_reg = int(os.environ.get('OGGM_GLACIER_JOB', ''))

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

# RGI glaciers
rgi_ids = gpd.read_file(utils.get_rgi_region_file(rgi_reg, version='62'))
# rgi_ids = rgi_ids.loc[rgi_ids.O2Region == '1']
# rgi_ids = rgi_ids.iloc[:4]

# Sort for more efficient mp
# rgi_ids = rgi_ids.sort_values(by='Area', ascending=False)

# Go - get the pre-processed glacier directories
base_url = 'https://cluster.klima.uni-bremen.de/~fmaussion/runs/lmr_2025/lmr_gdirs_v2/'
gdirs = workflow.init_glacier_directories(rgi_ids, 
                                          from_prepro_level=5, 
                                          prepro_base_url=base_url, 
                                          prepro_rgi_version=rgi_version)


rid = 'lmr_mira_hist_cl'
log.workflow(f'Running {rid}') 
workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                             climate_filename='climate_historical', 
                             # climate_input_filesuffix='_' + rid, 
                             output_filesuffix=rid,  # recognize the run for later
                             return_value=False,
                             ys=1,
                             );

log.workflow(f'Aggregating {rid}')
out_nc_path = os.path.join(OUTPUT_DIR, rid + f'_reg{rgi_reg}.nc')
utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)          

rid = 'lmr_mira_nat_cl'
log.workflow(f'Running {rid}') 
from mira_forcing import process_mira_data

workflow.execute_entity_task(process_mira_data, gdirs, nat_forcing=True, output_filesuffix='_nat')

workflow.execute_entity_task(tasks.run_from_climate_data, gdirs,
                             climate_filename='climate_historical',  
                             climate_input_filesuffix='_nat', 
                             output_filesuffix=rid,  # recognize the run for later
                             return_value=False,
                             ys=1,
                             );

log.workflow(f'Aggregating {rid}')
out_nc_path = os.path.join(OUTPUT_DIR, rid + f'_reg{rgi_reg}.nc')
utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)