import sys
import numpy as np
from netCDF4 import Dataset
import os
import time
from pyproj import CRS, Transformer

# C. Pelletier july 2021.

# Example of using PYPROJ to manage projections.
# I highly recommend using this tool, as you're then relying on something documented from projecting (i.e. from lat/lon to stereographic)
# This is pretty useful for ice sheet / ocean coupling and pre/post processing. since most ice sheet models are defined on polar stereographic grids.

# This script just gets the grid corners of the BISICLE grid.
# But it's an illustration of how to use PYPROJ.

# <https://pyproj4.github.io/pyproj/stable/>
# <https://proj.org/>

grid_bisicles = "/afast/pelletie/totten/BISICLES-smaller_grid.nc"
output_file = "/afast/pelletie/totten/BISICLES-smaller_grid_corners.nc"

# description of the projection.
proj_polster = "+proj=stere +lat_0=-90.0 +lat_ts=-71. +lon_0=-90."
# The projection python object
crs_polster = CRS.from_proj4(proj_polster)

tmpdat = Dataset(grid_bisicles, mode='r')

# Center of cells in stereographic projection. No idea why X and Y are interverted.
x_ctn = tmpdat.variables['Y'][:]
y_ctn = tmpdat.variables['X'][:]

mask = tmpdat.variables['MASK'][:].astype(int)

tmpdat.close()

nx = np.size(x_ctn)
ny = np.size(y_ctn)


# Edges of cell in stereographics projection. This is just the middle of two consecutive cells. And a bit of symmetry for the edges.
x_stg = np.ndarray(shape=[nx+1], dtype=float)
y_stg = np.ndarray(shape=[ny+1], dtype=float)

x_stg[1:-1] = .5 * ( x_ctn[:-1] + x_ctn[1:] )
x_stg[0] = 2 * x_ctn[0] - x_stg[1]
x_stg[-1] = 2 * x_ctn[-1] - x_stg[-2]

y_stg[1:-1] = .5 * ( y_ctn[:-1] + y_ctn[1:] )
y_stg[0] = 2 * y_ctn[0] - y_stg[1]
y_stg[-1] = 2 * y_ctn[-1] - y_stg[-2]


# PYPROJ transformation object, polster -> latlon
# <crs_something>.geodetic_crs is always the latlon projection, regardless of <crs_something>
polster_to_latlon = Transformer.from_crs(crs_polster, crs_polster.geodetic_crs)

# Get lon, lat of centers from x,y (which are broadcasted with [None,:] etc.)

lon_ctn, lat_ctn = polster_to_latlon.transform(x_ctn[None,:] * (np.ones(shape=[ny], dtype=float))[:,None], \
                                               y_ctn[:,None] * (np.ones(shape=[nx], dtype=float))[None,:])

# No idea why I'm doing this. I guess everything's corked and dirty.
lon_ctn = - lon_ctn

# Reformat corners in (y,x,4) instead of (y+1,x+1).
x_crn = np.ndarray(shape=[ny, nx, 4], dtype=float)
y_crn = np.ndarray(shape=[ny, nx, 4], dtype=float)

x_crn[:,:,0] = (x_stg[1:])[None,:] * (np.ones(shape=[ny], dtype=float))[:,None]
y_crn[:,:,0] = (y_stg[1:])[:,None] * (np.ones(shape=[nx], dtype=float))[None,:]

x_crn[:,:,1] = (x_stg[1:])[None,:] * (np.ones(shape=[ny], dtype=float))[:,None]
y_crn[:,:,1] = (y_stg[:-1])[:,None] * (np.ones(shape=[nx], dtype=float))[None,:]

x_crn[:,:,2] = (x_stg[:-1])[None,:] * (np.ones(shape=[ny], dtype=float))[:,None]
y_crn[:,:,2] = (y_stg[:-1])[:,None] * (np.ones(shape=[nx], dtype=float))[None,:]

x_crn[:,:,3] = (x_stg[:-1])[None,:] * (np.ones(shape=[ny], dtype=float))[:,None]
y_crn[:,:,3] = (y_stg[1:])[:,None] * (np.ones(shape=[nx], dtype=float))[None,:]



lon_crn = np.ndarray(shape=[ny, nx, 4], dtype=float)
lat_crn = np.ndarray(shape=[ny, nx, 4], dtype=float)

for m in range(0,4):
    lon_crn[:,:,m], lat_crn[:,:,m] = polster_to_latlon.transform(x_crn[:,:,m], y_crn[:,:,m])


lon_crn = - lon_crn

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

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

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

tmp = out_nc.createVariable('lat_corn', datatype='f8', dimensions=['y', 'x','crn'])
tmp[:] = lat_crn

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

tmp = out_nc.createVariable('lon_corn', datatype='f8', dimensions=['y', 'x','crn'])
tmp[:] = lon_crn

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

out_nc.close()