import sys

import numpy as np
from netCDF4 import Dataset

import os
import time


org_grid = '/afast/pelletie/orca025/eORCA025-SO_BMbathy_surf.nc'
dest_grid = '/afast/pelletie/orca025/eORCA025-SO_corners_new.nc'

data_org = Dataset(org_grid, mode='r')

glamt = data_org.variables['glamt'][:]
gphit = data_org.variables['gphit'][:]

glamu = data_org.variables['glamu'][:]
gphiu = data_org.variables['gphiu'][:]

glamv = data_org.variables['glamv'][:]
gphiv = data_org.variables['gphiv'][:]

glamf = data_org.variables['glamf'][:]
gphif = data_org.variables['gphif'][:]

mask = data_org.variables['tmask'][:]

data_org.close()

ny, nx = np.shape(glamt)

out_nc = Dataset(dest_grid, mode='w')

out_nc.createDimension('y', ny)
out_nc.createDimension('x', nx)

tmp = out_nc.createVariable(varname='lon', datatype='f8', dimensions=['y', 'x'])
tmp[:] = glamt[:]
tmp.standard_name = "longitude"
tmp.units = "degrees East"
tmp.bounds = "lon_corn"

tmp = out_nc.createVariable(varname='lat', datatype='f8', dimensions=['y', 'x'])
tmp[:] = gphit[:]
tmp.standard_name = "latitude"
tmp.units = "degrees North"
tmp.bounds = "lat_corn"

tmp = out_nc.createVariable(varname='mask', datatype='b', dimensions=['y', 'x'])
tmp[:] = mask[:]
tmp.coordinates = "lat lon"


out_nc.createDimension('corners', 4)

corn_tmp = np.ndarray(shape=[ny+1,nx+1], dtype=float)

corn_tmp[1:, 1:] = glamf
corn_tmp[0,1:] = 2*glamu[0,:] - glamf[0,:]
corn_tmp[:,0] = corn_tmp[:,-3] #- 360.

lon_c = out_nc.createVariable(varname='lon_corn', datatype='f8', dimensions=['y', 'x', 'corners'])

lon_c[:,:,0] = corn_tmp[:-1,:-1]
lon_c[:,:,1] = corn_tmp[1:,:-1]
lon_c[:,:,2] = corn_tmp[:-1,1:]
lon_c[:,:,3] = corn_tmp[1:,1:]

res = .25

corr_lonc = np.logical_or(np.abs(lon_c[:,:,0] - lon_c[:,:,2]) < res/100.,
                          np.abs(lon_c[:,:,1] - lon_c[:,:,3]) < res/100.)

lon_c[:,:,0] = np.where(corr_lonc,\
                        glamt[:] - .5 * res, lon_c[:,:,0])
lon_c[:,:,1] = np.where(corr_lonc,\
                        glamt[:] - .5 * res, lon_c[:,:,1])
lon_c[:,:,2] = np.where(corr_lonc,\
                        glamt[:] + .5 * res, lon_c[:,:,2])
lon_c[:,:,3] = np.where(corr_lonc,\
                        glamt[:] + .5 * res, lon_c[:,:,3])


del lon_c

corn_tmp[1:, 1:] = gphif
corn_tmp[0,1:] = 2*gphiu[0,:] - gphif[0,:]
corn_tmp[:,0] = corn_tmp[:,-1]

lat_c = out_nc.createVariable(varname='lat_corn', datatype='f8', dimensions=['y', 'x', 'corners'])

lat_c[:,:,0] = corn_tmp[:-1,:-1]
lat_c[:,:,1] = corn_tmp[1:,:-1]
lat_c[:,:,2] = corn_tmp[:-1,1:]
lat_c[:,:,3] = corn_tmp[1:,1:]

corr_latc = np.logical_or(np.abs(lat_c[:,:,0] - lat_c[:,:,1]) < res/100.,
                          np.abs(lat_c[:,:,2] - lat_c[:,:,3]) < res/100.)

lat_c[:,:,0] = np.where(corr_latc,\
                        gphit[:] - .5 * res, lat_c[:,:,0])
lat_c[:,:,1] = np.where(corr_latc,\
                        gphit[:] + .5 * res, lat_c[:,:,1])
lat_c[:,:,2] = np.where(corr_latc,\
                        gphit[:] - .5 * res, lat_c[:,:,2])
lat_c[:,:,3] = np.where(corr_latc,\
                        gphit[:] + .5 * res, lat_c[:,:,3])



out_nc.close()