#Marco Miani#
#21 Oct 2020, 17:30#
#generate netcdf with curvilinear grid on Z axis!#

import time
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
import pandas as pd
import numpy.matlib
from netCDF4 import Dataset




#close open windowws
plt.close()


#Set some params used later on: 
#x domain
xmin=0
xmax=2

#amplitude for a sine-wave perturbation 
Amp = .1

#grid points for domain: 
nx = 500   
nz = 12
nt = 100

#sloping of the ramp (to add to the z coordiante of each zlevel)
tan_ramp = 0.65



#define time vector, the way Paraview likes it: 
times =  pd.date_range(start='2000-01-01',freq='1H',periods=nt)

#define fake data
sal = 15 + 8 * np.random.randn(nt, nz, nx)


#here i define the coordiantes: 
Xp    = np.linspace(xmin,xmax,nx)
noise = .2* np.random.random(nx)
ramp  = tan_ramp * Xp
z     = Amp* (np.sin(np.linspace(0, 6.29,nx))) + ramp + noise



plt.plot(Xp, z)

# SOme (overcomplicated?) method to stack vertical levels of z, over X domain: 
for j in range(1,nz):
    dz = 0.5 *j
    noise = .2* np.random.random(nx)
    ramp  = tan_ramp * Xp
    temp_z = Amp* np.sin(np.linspace(0, 6.29,nx)) + noise +ramp -dz
    z =np.vstack([z,temp_z])
    plt.plot(Xp, z[j,:])
    plt.show()

#print
z

#show
plt.show()



##Crete an array: 
Xp = np.matlib.repmat(Xp, nz, 1)

ds = xr.Dataset(
    {
        "salinity": ([ "time", "y", "x"], sal)
    },
    coords={
         "time": pd.date_range("2014-09-06", periods=nt),
         "xc": ([ "y","x"], Xp),
         "yc": ([ "y","x"], z)  #------- as in CHapter 6.2. of http://cfconventions.org/cf-conventions/cf-conventions.html#discrete-axis
         #"reference_time": pd.Timestamp("2020-10-21"),
    },
)#

ds['salinity'].attrs['coordinates'] = 'yc xc' #--this doesn't seem to work???
ds.attrs['convention'] ='CF-1.4'

#Set attributes: 
ds['xc'].attrs['axis']="X"
ds['xc'].attrs['long_name']="x-coordinate in Cartesian system"
ds['xc'].attrs['units']="m"
ds['yc'].attrs['axis']="Z"
ds['yc'].attrs['long_name']="sigma"
ds['yc'].attrs['positive']="down"
ds['yc'].attrs['units']="m"

#make sure they are applied: 
ds.update(ds, inplace=None)


#print
print("--------------------------------Generated in house----------------")
print(ds)

print("")
print("")



print("--------------------------------imported from xrray----------------")
ds1=xr.tutorial.open_dataset('rasm')
print(ds1)




##modefiy the first field that xarray could not modify: 
#coordinates ='yc xc'


fname= "from_scratch.nc"
ds.to_netcdf(fname)
ds.close()
print("---------------------------------")
f = Dataset(fname,'r+')

#'coordinates' = 'yc xc'
print("original:")
print(f['salinity'].coordinates)

f['salinity'].coordinates ='yc xc' 
print("modified:")
print(f['salinity'].coordinates)
print("---------------------------------")

f.close()