import numpy as np
import pykrige.kriging_tools as kt
from pykrige import OrdinaryKriging
import rasterio as rio


class KrigingSpiData(object):

	def __init__(self, data, save_path, XY):
		self.save_path = save_path
		self.data = data
		self.XY = XY
		self.get_conf()

	def get_minmax(self):
		delta = 0.1
		miny = self.XY['X'].min() - delta
		minx = self.XY['Y'].min() - delta
		maxy = self.XY['X'].max() + delta
		maxx = self.XY['Y'].max() + delta
		print(miny, minx, maxy, maxx)
		return (minx, miny, maxx, maxy)

	def get_conf(self):
		(minx, miny, maxx, maxy) = self.get_minmax()
		res = 0.003
		self.gridx = np.arange(minx, maxx, res)
		self.gridy = np.arange(miny, maxy, res)
		self.height = len(self.gridx)
		self.width = len(self.gridy)
		from rasterio.transform import from_origin
		self.transform = from_origin(miny - res / 2, maxx + res / 2, res, res)

	def write_tif(self, kriging_data):
		out_meta = {"driver": "GTiff",
		            "height": self.height,
		            "width": self.width,
		            "transform": self.transform}
		with rio.open(self.save_path, "w", count=1, dtype='float64', crs='+proj=latlong', **out_meta) as dest:
			dest.write_band(1, kriging_data)

	def kriging(self):
		print(np.asarray(self.XY['Y']))
		# self.data = np.random.rand(179)
		print(self.data)

		OK = OrdinaryKriging(np.asarray(self.XY['X']), np.asarray(self.XY['Y']), self.data, variogram_model='spherical', verbose=False,
		                     enable_plotting=False)
		z, ss = OK.execute('grid', self.gridx, self.gridy)
		self.write_tif(z)


if __name__ == '__main__':
	kd = KrigingSpiData("D:/work/3.scala-gis/geotrellis-through/geotrellis/bj_197901.txt", "D://")
	kd.kriging()