# libraries 
import os 
import sys
import shutil 
import numpy as np 
from nco import Nco,NCOException 
nco = Nco()

# settings
idir = '/nird/datalake/NS11071K/www/seaclim/wp2/raw'
odir = '/nird/datalake/NS11071K/www/seaclim/wp2/calibration'
wdir = '/scratch/ingo/make_climatologies_seaclim'
expid = 'noresm2-mm-seaclim_hindcast' 
syears = np.arange(1982,2024+1)
climPeriod = [1993, 2024] 
member1 = 1
membern = 10
nleadMonth = 64
startMonth = 11 
outputGroups = [ 
    ['atm', 'cam.h0'], 
    ['ice', 'cice.h'], 
    ['rof','mosart.h0'],
    ['ocn','blom.hmphyglb'],
    ['ocn','blom.hmphy20n'],
    ['ocn','blom.hmbgcglb'],
]
test = False
if test:
    climPeriod = [1993, 1994]
    membern = 2
    nleadMonth = 2
    outputGroups = [['atm', 'cam.h0']] 

# change to work directory 
os.makedirs(odir,exist_ok=True)
os.makedirs(wdir,exist_ok=True)
os.chdir(wdir)

# main 
for component,outputGroup in outputGroups:
    fpath_clim = f'{odir}/{expid}.{outputGroup}.clim.{climPeriod[0]:d}-{climPeriod[1]:d}.startmonth11.leadmonth1-{nleadMonth:d}.mem{member1:d}-{membern:d}.nc'
    if os.path.exists(fpath_clim):
        print(f'skip existing {fpath_clim}') 
        continue 
    else:
        print(f'create {fpath_clim}')     
    # create climatologies for individual lead months 
    fpath_clim_leadMonths = []
    for leadMonth in range(1,nleadMonth+1):
        fpath_clim_leadMonth = f'{wdir}/{expid}.{outputGroup}.clim.{climPeriod[0]:d}-{climPeriod[1]:d}.startmonth11.leadmonth{leadMonth:d}.mem{member1:d}-{membern:d}.nc'
        fpath_clim_leadMonths.append(fpath_clim_leadMonth) 
        if os.path.exists(fpath_clim_leadMonth):
            print(f'skip existing {fpath_clim_leadMonth}') 
            continue 
        else:
            print(f'create {fpath_clim_leadMonth}')     
        month = (startMonth+leadMonth-2)%12+1
        # create member specific climatogies 
        fpath_clim_members = [] 
        for member in range(member1, membern+1):     
            fpath_clim_member = f'{wdir}/{expid}.{outputGroup}.clim.{climPeriod[0]:d}-{climPeriod[1]:d}.startmonth11.leadmonth{leadMonth:d}.mem{member:d}.nc'
            fpath_clim_members.append(fpath_clim_member)
            if os.path.exists(fpath_clim_member):
                print(f'skip existing {fpath_clim_member}') 
                continue 
            else:
                print(f'create {fpath_clim_member}')             
            ifiles = [] 
            for syear in syears:
                year = syear + int(np.floor((startMonth+leadMonth-2)/12))
                if year >= climPeriod[0] and year <= climPeriod[1]:
                    ifiles.append(f'{expid}_{syear:d}1101/{expid}_{syear:d}1101_mem{member:0>3d}/{component}/hist/{expid}_{syear:d}1101_mem{member:0>3d}.{outputGroup}.{year:d}-{month:0>2d}.nc')
            print(' - create climatology for single member, single lead month')
            nco.ncra(input=ifiles,output=fpath_clim_member,options=[f'-p {idir:s}/{expid:s}', "-O", "-7", "-L 0"]) 
        print(' - average ensemble')
        nco.ncea(input=fpath_clim_members,output=fpath_clim_leadMonth,options=["-O", "-7", "-L 5"])                
    print(' - concatenate lead months') 
    nco.ncrcat(input=fpath_clim_leadMonths,output=fpath_clim,options=["-O", "-7", "-L 0"])