#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: Ada Gjermundsen
year: 2019 - 2021
This script is used to calculate the eddy-induced overturning in CESM2 and NorESM2 (LM and MM) south of 50S
for the CMIP experiments piControl and abrupt-4xCO2 after 150 
the average time is 30 years
The result is used in FIGURE 4
"""

import sys
sys.path.insert(1, '/scratch/adagj/CMIP6/CLIMSENS/CMIP6_UTILS')
import CMIP6_ATMOS_UTILS as atmos
import CMIP6_SEAICE_UTILS as ocean
from read_modeldata_cmip6 import ecs_models_cmip6, make_filelist_cmip6, Modelinfo
import numpy as np
from dask.diagnostics import ProgressBar
import warnings
warnings.simplefilter('ignore')
import xarray as xr
xr.set_options(enable_cftimeindex=True)

def make_attributes(da, var, expid):
    da.attrs['long_name']='Global Ocean Meridional Overturning Mass Streamfunction Due to Parameterized Mesoscale Advection'
    da.attrs['name']='eddymoc'
    da.attrs['units']='kg s-1'
    da.attrs['standard_name']='global_ocean_meridional_overturning_mass_streamfunction_due_to_parameterized_mesoscale_eddy_advection'
    da.attrs['expid']=expid
    ds = da.to_dataset(name = var)
    return ds

def extract_global_moc(modelname, da, dac, var):
    if 'sector' in da.coords:
        da = da.drop('sector')
    if 'sector' in dac.coords:
        dac = dac.drop('sector')
    da = da.isel(basin=-1)
    dac = dac.isel(basin=-1)
    return da, dac
                                                                                                                      
def make_reference_slice(model, ds, var, endyr):
    ds = ocean.consistent_naming(ds)
    ds = atmos.fix_time(ds, 1)
    return ds

def make_yearly_avg(model, ds, var, endyr):
    da = atmos.yearly_avg(ds[var])
    if model.expid in ['piControl']:
        da = da.isel(year=slice(model.branchtime_year+endyr-30, model.branchtime_year+endyr))
    else:
        da = da.isel(year=slice(endyr-30, endyr))
    da = da.mean(dim='year')
    return da

def make_modelobj(modelname, expinfo, expid='piControl'):
    model = Modelinfo(name = modelname, institute = expinfo['institute'], expid = expid, realm = 'Omon',
                 realiz=expinfo['variant_labels'][0], grid_atmos = expinfo['grid_label_atmos'][0], grid_ocean = expinfo['grid_label_ocean'], branchtime_year=expinfo['branch_yr'])
    return model

def read_files(model, var):
    if model.name in ['NorESM2-LM', 'NorESM2-MM']:
        make_filelist_cmip6(model, var,  component = 'ocean', activity_id='CMIP',path_to_data = '/projects/NS9034K/CMIP6/')
    else:
        make_filelist_cmip6(model, var,  component = 'ocean')
    print(model.filenames)
    if model.filenames:
        if len(model.filenames)>1:
            ds =  xr.open_mfdataset(model.filenames, combine='nested', concat_dim='time', parallel=True, chunks={"time":1})
        else:
            ds =  xr.open_dataset(model.filenames[0], chunks={"time":1})
        print('%s loaded for model: %s, experiment: piControl . Lenght of simulation: %.1f years'%(var,model.name, len(ds[var].time.values)/12))
    else:
        print('%s not loaded for model %s, experiment: piControl. Skipping model! Please check!'%(var,model.name))
    return ds

def make_last_30yrs_avg(models, var, outpath, endyr=150):
    print('global eddy moc: \n')
    for modelname,expinfo in models.items():
        print(modelname)
        if var in ['msftmzsmpa'] and modelname in ['NorESM2-LM']:
            continue
        modelctrl = make_modelobj(modelname, expinfo, expid='piControl')
        dsc = read_files(modelctrl, var)
        dsc = make_reference_slice(modelctrl, dsc, var, endyr)
        model4xco2 = make_modelobj(modelname, expinfo, expid='abrupt-4xCO2')
        ds = read_files(model4xco2, var)
        ds = make_reference_slice(model4xco2, ds, var, endyr)
        ds, dsc = extract_global_moc(modelname, ds, dsc, var)
        da = make_yearly_avg(model4xco2, ds, var, endyr)
        dac = make_yearly_avg(modelctrl, dsc, var, endyr)
        dsout_ctrl = make_attributes(dac, var, 'piControl')
        dsout_case = make_attributes(da, var, 'abrupt-4xCO2')
        print(dsout_ctrl)
        print(dsout_case)
        dsout_ctrl = dsout_ctrl.to_netcdf(outpath + var +'_' + modelctrl.realm +'_' + modelctrl.name + '_' + modelctrl.expid + '_' + modelctrl.realiz + '_'+str(endyr) + '_30yravg.nc', compute=False)
        dsout_case = dsout_case.to_netcdf(outpath + var +'_' + model4xco2.realm +'_' + model4xco2.name + '_' + model4xco2.expid + '_' + model4xco2.realiz + '_'+str(endyr) + '_30yravg.nc', compute=False)
        with ProgressBar():
            result = dsout_ctrl.compute()
            result = dsout_case.compute()
        del model4xco2, modelctrl, dsc, ds, dac, da, dsout_ctrl, dsout_case


if __name__ == '__main__':
    outpath = 'path_to_outdata/' 
    models = ecs_models_cmip6()
    models = {'NorESM2-LM':models['NorESM2-LM'], 'CESM2':models['CESM2']}
    for var in ['msftmzsmpa', 'msftmzmpa']:
        make_last_30yrs_avg(models, var=var, outpath=outpath, endyr=150)
   

[((13, 0, 13, 63), 'sys.path.insert', 'sys.path.insert', ({(13, 16, 13, 17): '(1)', (13, 19, 13, 62): '"""/scratch/adagj/CMIP6/CLIMSENS/CMIP6_UTILS"""'}, {}), "(1, '/scratch/adagj/CMIP6/CLIMSENS/CMIP6_UTILS')", False, 'import sys\n'), ((20, 0, 20, 31), 'warnings.simplefilter', 'warnings.simplefilter', ({(20, 22, 20, 30): '"""ignore"""'}, {}), "('ignore')", False, 'import warnings\n'), ((22, 0, 22, 39), 'xarray.set_options', 'xr.set_options', (), '', True, 'import xarray as xr\n'), ((43, 9, 43, 36), 'CMIP6_SEAICE_UTILS.consistent_naming', 'ocean.consistent_naming', ({(43, 33, 43, 35): 'ds'}, {}), '(ds)', True, 'import CMIP6_SEAICE_UTILS as ocean\n'), ((44, 9, 44, 30), 'CMIP6_ATMOS_UTILS.fix_time', 'atmos.fix_time', ({(44, 24, 44, 26): 'ds', (44, 28, 44, 29): '1'}, {}), '(ds, 1)', True, 'import CMIP6_ATMOS_UTILS as atmos\n'), ((48, 9, 48, 34), 'CMIP6_ATMOS_UTILS.yearly_avg', 'atmos.yearly_avg', ({(48, 26, 48, 33): 'ds[var]'}, {}), '(ds[var])', True, 'import CMIP6_ATMOS_UTILS as atmos\n'), ((57, 12, 58, 178), 'read_modeldata_cmip6.Modelinfo', 'Modelinfo', (), '', False, 'from read_modeldata_cmip6 import ecs_models_cmip6, make_filelist_cmip6, Modelinfo\n'), ((106, 13, 106, 31), 'read_modeldata_cmip6.ecs_models_cmip6', 'ecs_models_cmip6', ({}, {}), '()', False, 'from read_modeldata_cmip6 import ecs_models_cmip6, make_filelist_cmip6, Modelinfo\n'), ((63, 8, 63, 123), 'read_modeldata_cmip6.make_filelist_cmip6', 'make_filelist_cmip6', (), '', False, 'from read_modeldata_cmip6 import ecs_models_cmip6, make_filelist_cmip6, Modelinfo\n'), ((65, 8, 65, 61), 'read_modeldata_cmip6.make_filelist_cmip6', 'make_filelist_cmip6', (), '', False, 'from read_modeldata_cmip6 import ecs_models_cmip6, make_filelist_cmip6, Modelinfo\n'), ((69, 18, 69, 123), 'xarray.open_mfdataset', 'xr.open_mfdataset', (), '', True, 'import xarray as xr\n'), ((71, 18, 71, 72), 'xarray.open_dataset', 'xr.open_dataset', (), '', True, 'import xarray as xr\n'), ((98, 13, 98, 26), 'dask.diagnostics.ProgressBar', 'ProgressBar', ({}, {}), '()', False, 'from dask.diagnostics import ProgressBar\n')]