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


class KrigingDataToStation(object):

	def __init__(self, path, save_dir):
		self.data = KrigingDataToStation.read_data(path)
		self.save_dir = save_dir
		# self.get_conf()
		self.get_xy()
		self.get_date(path)

	def get_date(self, path):
		name = os.path.basename(path)
		self.year = name[0:4]
		self.month = name[4:name.index('.')]

	@staticmethod
	def read_data(path):
		data = []
		with open(path, 'r') as file:
			for line in file.readlines():
				arr = line.split(',')
				data.extend([float(arr[0]), float(arr[1]), float(arr[4]), float(arr[5])])
		return np.asarray(data).reshape(179, 4)

	def get_minmax(self):
		delta = 0.1
		miny = self.data[:, 1].min() - delta
		minx = self.data[:, 2].min() - delta
		maxy = self.data[:, 1].max() + delta
		maxx = self.data[:, 2].max() + delta
		return (minx, miny, maxx, maxy)

	def get_xy(self):
		path = '../data/station/station_new.csv'
		with open(path, 'r') as file:
			list_x, list_y, list_name = [], [], []
			lines = file.readlines()[1:]
			for line in lines:
				arr = line.split(',')
				x, y, name = arr[1], arr[2], arr[0]
				list_x.append(float(x))
				list_y.append(float(y))
				list_name.append(arr[0])
			self.gridx = np.asarray(list_x)
			self.gridy = np.asarray(list_y)
			self.station_names = list_name

	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 kriging_name(self, value_name):
		if value_name == 'rain':
			value_index = 2
		else:
			value_index = 3
		OK = OrdinaryKriging(self.data[:, 0], self.data[:, 1], self.data[:, value_index], variogram_model='spherical',
		                     verbose=False, enable_plotting=False)
		# print(self.gridx, self.gridy)
		z, ss = OK.execute('grid', self.gridx, self.gridy)
		return z

	def kriging(self):
		data_rain = self.kriging_name(value_name="rain")
		data_temp = self.kriging_name(value_name="temp")
		for i in range(0, len(data_rain)):
			print(data_rain[i, i], data_temp[i, i])
			path = '../data/station/station_new/%s.csv' % (self.station_names[i])
			with open(path, 'a') as file:
				print(self.gridy[i])
				line = str(self.gridx[i]) + ',' + str(self.gridy[i]) + "," + self.year + "," + self.month + "," + str(data_rain[i, i]) + "," + str(data_temp[i, i]) + "\r"
				file.write(line)


if __name__ == '__main__':
	import glob
	import os

	search_dir = "../data/station/value/"
	files = filter(os.path.isfile, glob.glob(search_dir + "*"))
	files = [f for f in files]
	files.sort(key=lambda x: os.path.getmtime(x))

	for file in files:
		print(file)
		kd = KrigingDataToStation(file, "D://")
		kd.kriging()