from osgeo import gdal
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import numpy as np
from scipy import interpolate

# Read the data and metadata
datafile = gdal.Open(r'2c_lavhei_68N_Jenny_wgs84.tif')
bnd1 = datafile.GetRasterBand(1).ReadAsArray()
nx = datafile.RasterXSize # Raster xsize
ny = datafile.RasterYSize # Raster ysize

gt = datafile.GetGeoTransform()
proj = datafile.GetProjection()

xres = gt[1]
yres = gt[5]

# get the edge coordinates and add half the resolution
# to go to center coordinates
xmin = gt[0] + xres * 0.5
xmax = gt[0] + (xres * nx) - xres * 0.5
ymin = gt[3] + (yres * ny) + yres * 0.5
ymax = gt[3] - yres * 0.5
print("xmin=", xmin,"xmax=", xmax,"ymin=",ymin, "ymax=", ymax)

# create a grid of lat/lon coordinates in the original projection
(lon_source,lat_source) = np.mgrid[xmin:xmax+xres:xres, ymax+yres:ymin:yres]
print(xmin,xmax+xres,xres, ymax+yres,ymin,yres)
print(lon_source[:,0], np.flip(lat_source[0,:],0))

# Read in coordinates and pfts for output file:
outfile = Dataset('surfdata_20x20_FINNMARK_simyr2000_cesm1_2_x_clm4_5.2c_lavhei.nc', mode='r+')
outlats = outfile.variables['LATIXY'][:]
outlons = outfile.variables['LONGXY'][:]
outpct_pft=outfile.variables['PCT_PFT'][:]
outpct_lake = outfile.variables['PCT_LAKE'][:]

outlats = outlats[:,0]
outlons = outlons[0,:]

#print(outlats, outlons)
print(bnd1.shape, lon_source.shape)

# Remove al negative values:
bnd1[bnd1 < 0] = 0

print(np.min(bnd1))

"""
# Edges of model box:
edgen = 69.9
edges = 68.8
edgew = 22.2
edgee = 25.9

# Find indices of wanted box:
indexs = 0
while inlats[indexs] <= edges:
	indexs += 1

indexn = indexs
while inlats[indexn] <= edgen:
	indexn += 1

indexw = 0
while inlons[indexw] <= edgew:
	indexw += 1

indexe = indexw
while inlons[indexe] <= edgee:
	indexe += 1

# Make new box containing only wanted box:
box_lats = inlats[indexs:indexn]
box_lons = inlons[indexw:indexe]
box_pct_pft = inpct_pft[:,indexs:indexn,indexw:indexe]
#print(box_pct_pft[17,:,:])

inpct_pft = np.nan_to_num(inpct_pft)
"""



interp_lavhei = interpolate.RectBivariateSpline(np.flip(lat_source[0,:-1],0), lon_source[:-1,0], bnd1, kx=1, ky=1)

new_lavhei = interp_lavhei(outlats, outlons)
print(np.max(new_lavhei))

plt.imshow(new_lavhei)
plt.colorbar()
plt.show()

"""
outlats = outlats[:,0]
outlons = outlons[0,:]

newtree = interp_tree(outlats, outlons)
newgras = interp_gras(outlats, outlons)
newshrub = interp_shrub(outlats, outlons)
newwater = interp_water(outlats, outlons)
newlichen = interp_lichen(outlats, outlons)

newground = np.zeros((len(outlats), len(outlons)))

for i in range(len(outlats)):
	for j in range(len(outlons)):
		pct = newtree[i,j] + newgras[i,j] + newshrub[i,j] + newwater[i,j] + newlichen[i,j]
		if pct > 100:
			if newtree[i,j] >= (pct-100):
				newtree[i,j] = newtree[i,j] - (pct-100)
			elif newshrub[i,j] >= (pct-100):
                                newshrub[i,j] = newshrub[i,j] - (pct-100)
			else:
				print('pct blir negativ for i=', i, ' j=', j)
		elif pct < 100:
			newground[i,j] = 100-pct

pct_pft = np.zeros((17,len(outlats),len(outlons)))
pct_pft[0,:,:] = newground[:,:]
pct_pft[2,:,:] = newtree[:,:]
pct_pft[11,:,:] = newshrub[:,:]
pct_pft[12,:,:] = newgras[:,:]
pct_pft[4,:,:] = newlichen[:,:]

print(outlats)
print(outlons)
#print(newground)
#print(sum(pct_pft[:,0,0]))
i,j = np.unravel_index(newlichen.argmax(), newlichen.shape)
print(newlichen[i,j], i, j, outlats[i], outlons[j])

outfile.variables['PCT_PFT'][:] = pct_pft
outfile.variables['PCT_LAKE'][:] = newwater

#lichen_pct_pft = interp(outlats, outlons)
#print(lichen_pct_pft)

#a = np.arange(6).reshape(2,3)
#print(a)
#plt.imshow(a)
#plt.colorbar()
#plt.show()

#plt.subplot(2,1,1)
#plt.imshow(box_pct_pft[15,:,:])
#plt.subplot(2,1,2)
plt.imshow(newlichen)
plt.show()

#for i in range(len(outlats)):
#        for j in range(len(outlons)):
#                print(newtree[i,j] + newgras[i,j] + newshrub[i,j] + newwater[i,j] + newlichen[i,j] + newground[i,j])

outfile.close()
"""