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

from helper import perturbate_mb_params, run_from_climate_data_glen_a_fac
    
# Initialize OGGM and set up the default run parameters
cfg.initialize(logging_level='ERROR')
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)

DELETE_FULL_FILES = True

cfg.PARAMS['continue_on_error'] = True
cfg.PARAMS['store_diagnostic_variables'] = ['volume', 'volume_bsl', 'area']
cfg.PARAMS['store_model_geometry'] = False

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

# RGI glaciers
with open('rgi_chuncks_s.json', 'r') as f:
    rgi_ids = json.load(f)[f'{OGGM_GLACIER_JOB}']

# rgi_ids = rgi_ids[:32]
    
rgi_reg = rgi_ids[0].split('-')[1].split('.')[0]
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))
    
# Sort for more efficient mp
# rgi_ids = rgi_ids.sort_values(by='Area', ascending=False)

# Prepare params
params_df = pd.read_csv('hypercube_oggm_v2.csv', index_col=0)
exps = params_df.index

# Go - get the pre-processed glacier directories
base_url = 'https://cluster.klima.uni-bremen.de/~oggm/gdirs/oggm_v1.6/L3-L5_files/2023.3/elev_bands/W5E5_spinup/'
gdirs = workflow.init_glacier_directories(rgi_ids, 
                                          from_prepro_level=5, 
                                          prepro_base_url=base_url, 
                                          prepro_rgi_version=rgi_version)

gcms = pd.read_csv('/home/www/oggm/cmip6/all_gcm_list.csv', index_col=0)

# Run a new GCM with all scenario
for gcm in ['CESM2-WACCM']:
    df1 = gcms.loc[gcms.gcm == gcm]
    for ssp in ['ssp126', 'ssp585', 'ssp534-over']:
        
        df2 = df1.loc[df1.ssp == ssp]
        assert len(df2) == 2
        ft = df2.loc[df2['var'] == 'tas'].iloc[0]
        fp = df2.loc[df2['var'] == 'pr'].iloc[0].path
        
        rid = ft.fname.replace('_r1i1p1f1_tas.nc', '')
        ft = ft.path
        
        log.workflow(f'Bias correcting {rid}') 
        
        workflow.execute_entity_task(tasks.process_cmip_data, gdirs, 
                                     filesuffix='_' + rid,  # recognize the climate file for later
                                     fpath_temp=ft,  # temperature projections
                                     fpath_precip=fp,  # precip projections
                                     year_range=('2000', '2019'),
                                     y0=2000-10,
                                     );

        for exp in exps:
        
            log.workflow(f'Starting run for GCM {rid} param exp {exp}') 
            
            if exp > 0:
                params = params_df.loc[exp]
                
                cfg.PARAMS['temp_melt'] = params['temp_melt'] - 1
                glen_a_fac = params['glen_a']
                clim_params = {k: params[k] for k in ('temp_bias', 'prcp_fac', 'melt_f')}
                workflow.execute_entity_task(perturbate_mb_params, gdirs, perturbation=clim_params)
            else:
                cfg.PARAMS['temp_melt'] = - 1
                glen_a_fac = 1
                workflow.execute_entity_task(perturbate_mb_params, gdirs, reset_default=True)
                                             
            
            exp = f'{exp:02d}'
            workflow.execute_entity_task(run_from_climate_data_glen_a_fac, gdirs,
                                         glen_a_fac=glen_a_fac,
                                         climate_filename='gcm_data',  # use gcm_data, not climate_historical
                                         climate_input_filesuffix='_' + rid,  # use a different scenario
                                         init_model_filesuffix='_spinup_historical',          
                                         init_model_yr=2000,
                                         output_filesuffix=rid,  # recognize the run for later
                                         return_value=False,
                                         ye=2300,
                                         );

            gcm_dir = os.path.join(OUTPUT_DIR, base_url.split('L3-L5_files/')[-1].split('/')[-1], 'RGI' + rgi_reg, gcm)
            
            out_nc_path = os.path.join(gcm_dir, rid + f'_exp_{exp}_batch_{OGGM_GLACIER_JOB}.nc')
            out_csv_path = os.path.join(gcm_dir, rid + f'_exp_{exp}_batch_{OGGM_GLACIER_JOB}.csv')
            out_json_path = os.path.join(gcm_dir, rid + f'_exp_{exp}_failed_ids_batch_{OGGM_GLACIER_JOB}.json')

            utils.mkdir(gcm_dir)
            
            try:
                utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)

                # Aggregate
                with xr.open_dataset(out_nc_path) as ds:
                    dfo = ds.sum(dim='rgi_id').to_dataframe()[['volume', 'volume_bsl', 'area']]
                    failed_ids = ds.rgi_id[ds.volume.isel(time=0).isnull()].data.tolist()
                    out_json = {exp: failed_ids}
                    dfo.to_csv(out_csv_path)
                    with open(out_json_path, 'w') as f:
                        json.dump(out_json, f)

                if DELETE_FULL_FILES:
                    os.remove(out_nc_path)
                    
            except RuntimeError:
                dfo = pd.DataFrame(columns=['volume', 'volume_bsl', 'area'])
                out_json = {exp: rgi_ids}
                dfo.to_csv(out_csv_path)
                with open(out_json_path, 'w') as f:
                    json.dump(out_json, f)

# Run old GCMs with a new scenario
for gcm in ['IPSL-CM6A-LR', 'MRI-ESM2-0']:
    df1 = gcms.loc[gcms.gcm == gcm]
    for ssp in ['ssp534-over']:
        
        df2 = df1.loc[df1.ssp == ssp]
        assert len(df2) == 2
        ft = df2.loc[df2['var'] == 'tas'].iloc[0]
        fp = df2.loc[df2['var'] == 'pr'].iloc[0].path
        
        rid = ft.fname.replace('_r1i1p1f1_tas.nc', '')
        ft = ft.path
        
        log.workflow(f'Bias correcting {rid}') 
        
        workflow.execute_entity_task(tasks.process_cmip_data, gdirs, 
                                     filesuffix='_' + rid,  # recognize the climate file for later
                                     fpath_temp=ft,  # temperature projections
                                     fpath_precip=fp,  # precip projections
                                     year_range=('2000', '2019'),
                                     y0=2000-10,
                                     );

        for exp in exps:
        
            log.workflow(f'Starting run for GCM {rid} param exp {exp}') 
            
            if exp > 0:
                params = params_df.loc[exp]
                
                cfg.PARAMS['temp_melt'] = params['temp_melt'] - 1
                glen_a_fac = params['glen_a']
                clim_params = {k: params[k] for k in ('temp_bias', 'prcp_fac', 'melt_f')}
                workflow.execute_entity_task(perturbate_mb_params, gdirs, perturbation=clim_params)
            else:
                cfg.PARAMS['temp_melt'] = - 1
                glen_a_fac = 1
                workflow.execute_entity_task(perturbate_mb_params, gdirs, reset_default=True)
                                             
            
            exp = f'{exp:02d}'
            workflow.execute_entity_task(run_from_climate_data_glen_a_fac, gdirs,
                                         glen_a_fac=glen_a_fac,
                                         climate_filename='gcm_data',  # use gcm_data, not climate_historical
                                         climate_input_filesuffix='_' + rid,  # use a different scenario
                                         init_model_filesuffix='_spinup_historical',          
                                         init_model_yr=2000,
                                         output_filesuffix=rid,  # recognize the run for later
                                         return_value=False,
                                         ye=2300,
                                         );

            gcm_dir = os.path.join(OUTPUT_DIR, base_url.split('L3-L5_files/')[-1].split('/')[-1], 'RGI' + rgi_reg, gcm)
            
            out_nc_path = os.path.join(gcm_dir, rid + f'_exp_{exp}_batch_{OGGM_GLACIER_JOB}.nc')
            out_csv_path = os.path.join(gcm_dir, rid + f'_exp_{exp}_batch_{OGGM_GLACIER_JOB}.csv')
            out_json_path = os.path.join(gcm_dir, rid + f'_exp_{exp}_failed_ids_batch_{OGGM_GLACIER_JOB}.json')

            utils.mkdir(gcm_dir)
            
            try:
                utils.compile_run_output(gdirs, input_filesuffix=rid, path=out_nc_path)

                # Aggregate
                with xr.open_dataset(out_nc_path) as ds:
                    dfo = ds.sum(dim='rgi_id').to_dataframe()[['volume', 'volume_bsl', 'area']]
                    failed_ids = ds.rgi_id[ds.volume.isel(time=0).isnull()].data.tolist()
                    out_json = {exp: failed_ids}
                    dfo.to_csv(out_csv_path)
                    with open(out_json_path, 'w') as f:
                        json.dump(out_json, f)

                if DELETE_FULL_FILES:
                    os.remove(out_nc_path)
                    
            except RuntimeError:
                dfo = pd.DataFrame(columns=['volume', 'volume_bsl', 'area'])
                out_json = {exp: rgi_ids}
                dfo.to_csv(out_csv_path)
                with open(out_json_path, 'w') as f:
                    json.dump(out_json, f)

log.workflow('Aggregating done!')