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

"""
datafile = gdal.Open('NN_wgs84.tif')
print( "Driver: ",datafile.GetDriver().ShortName, datafile.GetDriver().LongName)
print( "Size is ", datafile.RasterXSize, datafile.RasterYSize)
print( "Bands = ", datafile.RasterCount)
print( "Coordinate System is:", datafile.GetProjectionRef ())
print( "GetGeoTransform() = ", datafile.GetGeoTransform ())
print( "GetMetadata() = ", datafile.GetMetadata ())

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

bnd1 = datafile.GetRasterBand(1).ReadAsArray()

nx = datafile.RasterXSize # Raster xsize
ny = datafile.RasterYSize # Raster ysize

print("Geotransform",gt)
print("proj=", proj)
xres = gt[1]
yres = gt[5]
print("res=", xres, yres)

# 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)

xlons_tiff = np.arange(xmin,xmax,xres)
xlats_tiff = np.arange(ymin,ymax,abs(yres))
print("xlons_tiff", xlons_tiff)
print("xlats_tiff", xlats_tiff)

"""
infile = Dataset("vegetation_percentage.nc", "r")
inlats = infile.variables['LAT'][:]
inlons = infile.variables['LON'][:]
inpct_pft=infile.variables['PCT_PFT'][:]
infile.close()
#print(inpct_pft[17,500,500:600])
#print(inlats, inlons)

outfile = Dataset('surfdata_20x20_FINNMARK_simyr2000_cesm1_2_x_clm4_5.test.nc', mode='r+')
outlats = outfile.variables['LATIXY'][:]
outlons = outfile.variables['LONGXY'][:]
outpct_pft=outfile.variables['PCT_PFT'][:]
outpct_lake = outfile.variables['PCT_LAKE'][:]

#print(outlats, inlats)

print(inpct_pft.shape)
print(inlats.shape)
print(inlons.shape)

# 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)

plt.imshow[inpct_pft]
plt.colorbar()
plt.show()


tree = inpct_pft[1,:,:] + inpct_pft[2,:,:] + inpct_pft[3,:,:] + inpct_pft[4,:,:] + inpct_pft[5,:,:] + inpct_pft[6,:,:] + inpct_pft[7,:,:] + inpct_pft[8,:,:]
shrub = inpct_pft[9,:,:] + inpct_pft[10,:,:] + inpct_pft[11,:,:] + inpct_pft[14,:,:] + inpct_pft[17,:,:]
#ground_pft = inpct_pft[12,:,:] + inpct_pft[21,:,:]
gras = inpct_pft[13,:,:] + inpct_pft[16,:,:] + inpct_pft[18,:,:] + inpct_pft[19,:,:] + inpct_pft[20,:,:] + inpct_pft[23,:,:]
water = inpct_pft[22,:,:]
lichen = inpct_pft[15,:,:]

#interp = interpolate.interp2d(box_lons, box_lats, box_pct_pft[17,:,:], kind='linear')
#interp = interpolate.RectBivariateSpline(inlats, inlons, inpct_pft[17,:,:])
#interp = interpolate.interp2d(xlons_tiff, xlats_tiff, bnd1, kind='linear') # Må fikse bnd1 for dette

interp_tree = interpolate.RectBivariateSpline(inlats, inlons, tree, kx=1, ky=1)
interp_gras = interpolate.RectBivariateSpline(inlats, inlons, gras, kx=1, ky=1)
interp_shrub = interpolate.RectBivariateSpline(inlats, inlons, shrub, kx=1, ky=1)
interp_water = interpolate.RectBivariateSpline(inlats, inlons, water, kx=1, ky=1)
interp_lichen = interpolate.RectBivariateSpline(inlats, inlons, lichen, kx=1, ky=1)

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()