# -*- coding: utf-8 -*-
"""
Created on Fri May 15 13:35:08 2020

@author: alvareza
"""

from netCDF4 import Dataset
import os 

#%% 0.25 degree from -180 to 179.5 (left-right) and 0.25 degree from 90 to -90 (top-down)
pathslp = r'p:\metocean-data\open\ERA5\data\Global\msl_p\era5_Global_msl_p_1981.nc'
subpath = r'chunked_data'

pathchunk=pathslp

#%% Chunk data
dimedict = {'time':1000,'longitude':19,'latitude':19}  # NORTH SEA
dimestart = {'time':0,'longitude':732,'latitude':151} # NORTH SEA

toexclude = [] # in case there are more variables we want to skip
with Dataset(os.path.join(pathchunk),'r') as src, Dataset(os.path.join(os.getcwd(),subpath,os.path.basename(pathchunk)), "w") as chunk:
    # copy global attributes all at once via dictionary
    chunk.setncatts(src.__dict__)
    # copy dimensions
    for name, dimension in src.dimensions.items():
        chunk_dimension = None
        if dimension.isunlimited():
            chunk_dimension = dimedict[dimension]
        chunk.createDimension(name, chunk_dimension)
    # copy all file data except for the excluded
    for name, variable in src.variables.items():
        if name not in toexclude:
            x = chunk.createVariable(name, variable.datatype, variable.dimensions)
            # copy variable attributes all at once via dictionary
            chunk[name].setncatts(src[name].__dict__)
            print('attributes copied', flush=True)
            if len(x.dimensions) > 1:
                chunk[name][:] = src[name][dimestart['time']:dimestart['time']+x.shape[0], dimestart['latitude']-x.shape[1]:dimestart['latitude'], dimestart['longitude']:dimestart['longitude']+x.shape[2]]
                print(x.shape[0])
                print(x.shape[1])
                print(x.shape[2])
            else:
                if name=='latitude':
                    chunk[name][:] = src[name][dimestart[name]-dimedict[name]:dimestart[name]]                     
                else:
                    chunk[name][:] = src[name][dimestart[name]:dimestart[name]+dimedict[name]]