#!/home/CWA_MDPS/.conda/envs/rd/bin/python
import pytools.nctools as nct
import pytools.timetools as tt
from pytools.fillNans2d import fillNans2d
from tqdm import tqdm
import numpy as np
import os
import netCDF4
# !!! run this file on betago !!!

#       ssu     ssv     ssh     
# 2011                   
# 2012                   
# 2013                   
# 2014                   
# 2015                   
# 2016                   
# 2017                   

def main():
    varName = 'ssu'
    for year in [2017, 2016, 2015, 2014, 2013, 2012, 2011]:
        for varName in ['ssu', 'ssv', 'ssh']:
            run(varName, year=year, debug=False)


def run(varName, year, debug):
    numSmoothsNan = 20 * 4 # 20 deg * 4 grid/deg
    overwrite = False
    print(f'{varName=}, {year=}, {overwrite=}, {debug=}')
    if debug:
        nt = 5
    else:
        nt = 1460 + 4 * tt.isleap(tt.ymd2float(year, 1, 1))

    srcPath = getSrcPath(varName, year)
    desPath = getDesPath(varName, year)
    if not srcPath:
        return

    if os.path.exists(desPath) and not overwrite:
        print(f'destination path already exists: {desPath}')
        return

    if not os.path.exists(desPath):
        varShape = nct.getVarShape(srcPath, varName)
        dimNames = nct.getDimNames(srcPath, varName)

        nct.create(
            desPath, varName, varShape, dimNames, use_my_attrs=True,
            significant_digits=None, dtype=np.float32, complevel=0, 
            shuffle=False
        )

        for dimName in dimNames:
            nct.write(
                desPath,
                dimName,
                nct.read(srcPath, dimName)
            )

    with netCDF4.Dataset(desPath, 'a') as hDes, netCDF4.Dataset(srcPath, 'r') as hSrc:
        src = hSrc[varName]
        des = hDes[varName]
        for it in tqdm(range(nt)):
            data = src[it, :]
            data = fillNans2d(data, numSmoothsNan)
            des[it, :] = data
            


def getSrcPath(varName, year):
    srcPath = f'../data/merged/{varName}_{year}.nc'
    if not os.path.exists(srcPath):
        print(f'[fatal] source path not found: {srcPath}')
        return None
    return srcPath


def getDesPath(varName, year):
    desPath = f'../data/merged_interp/{varName}_{year}.nc'
    return desPath



if __name__ == '__main__':
    main()