from __future__ import division
import numpy as np
from netCDF4 import Dataset
import math
from scipy.interpolate import interp1d
import netCDF4 as cdf

dw1 = Dataset('../Experiments/LOWRES_DCMIP41/netcdf/output_dcmip2016_regular.nc','r')
dtfull = dw1.variables['T'][:,:,:,:]
dpfull = dw1.variables['P'][:,:,:,:]
dufull = dw1.variables['U'][:,:,:,:]
dvfull = dw1.variables['V'][:,:,:,:]
dPSfull = dw1.variables['PS'][:,:,:]
time_counterfull = dw1.variables['time_counter'][:]
lat = dw1.variables['lat'][:]
lon = dw1.variables['lon'][:]
levels = [100., 200., 300., 500., 700., 1000., 2000., 3000., 5000., 7000., 10000., 12500., 15000., 17500., 20000., 22500.,25000., 30000., 35000., 40000., 
    45000., 50000., 55000., 60000., 65000., 70000., 75000., 77500., 80000., 82500., 85000., 87500., 90000., 92500., 95000., 97500., 100000.]

levels = levels[::-1]
print(levels)
out3 = np.zeros((((dpfull.shape[0],len(levels),dpfull.shape[2],dpfull.shape[3]))))
out4 = np.zeros((((dpfull.shape[0],len(levels),dpfull.shape[2],dpfull.shape[3]))))
out5 = np.zeros((((dpfull.shape[0],len(levels),dpfull.shape[2],dpfull.shape[3]))))

def vertical_int2(levels, pfull, dataset):
    nl, ni, nj = len(levels), dataset.shape[1], dataset.shape[2]
    out = np.zeros((nl,ni,nj))
    for i in range(dataset.shape[1]):
        for j in range(dataset.shape[2]):
            f = interp1d(pfull[:,i,j],dataset[:,i,j],kind='linear',fill_value="extrapolate")
            out[:,i,j] = f(levels)
    return out


for t in range(dtfull.shape[0]):
    print("time step=",t)
    tfull = dtfull[t,:,:,:]
    pfull = dpfull[t,:,:,:]
    ufull = dufull[t,:,:,:]
    vfull = dvfull[t,:,:,:]
    out2 = vertical_int2(levels, pfull, tfull)
    out3[t,:,:,:]=out2[:,:,:]
    out2 = vertical_int2(levels, pfull, ufull)
    out4[t,:,:,:]=out2[:,:,:]
    out2 = vertical_int2(levels, pfull, vfull)
    out5[t,:,:,:]=out2[:,:,:]
    print("interpolated field:",out2)

print("-----------------------------writing NetCDF--------------------------------")
print(out3)
#sys.exit(0)
try:
    f = cdf.Dataset('DYNAMICO.nc', 'w', format='NETCDF4')
except:
    print("Error occurred while opening new netCDF file, Error: ", sys.exc_info()[0])
levels2 = [levels[i] / 100. for i in range(len(levels))]
#levels2 = [int(i) for i in levels2]
#print(levels2)
f.createDimension('lon', len(lon))
f.createDimension('lat', len(lat))
f.createDimension('lev', len(levels2))
f.createDimension('time_counter', len(time_counterfull))
vlon = f.createVariable('lon', 'f4', 'lon')
vlat = f.createVariable('lat', 'f4', 'lat')  
vlev = f.createVariable('lev', 'f4', 'lev')
#vlev = f.createVariable('lev', 'i4', 'lev')
vtime = f.createVariable('time_counter', 'f8', 'time_counter')
vt = f.createVariable('t', 'f4',('time_counter', 'lev', 'lat', 'lon'))
vu = f.createVariable('u', 'f4',('time_counter', 'lev', 'lat', 'lon'))
vv = f.createVariable('v', 'f4',('time_counter', 'lev', 'lat', 'lon'))
vPS = f.createVariable('PS', 'f4',('time_counter', 'lat', 'lon'))
vlon[:] = lon
vlat[:] = lat
vlev[:] = levels2
vtime[:] = time_counterfull
vt[:] = out3
vu[:] = out4
vv[:] = out5
vPS[:] = dPSfull
f.close()
print("-----------------------done--------------------")