from geopandas import *
from pandas import *
from shapely.geometry import Point
import rasterio as rio
import rasterio.mask
from rasterio.warp import reproject, Resampling
import os
import glob

DATA_BASE_PATH = '/home/g214/data_from_chenhao/data_analyse/'
TEMP_PATH = '/home/g214/temp/'
BJ_REGION = '/home/g214/data_from_chenhao/data_analyse/boundary/bj_region/bj_region.shp'


# DATA_BASE_PATH = r"D:\7. business\7.baoji\code\drought_analyse\data/"

def get_buffer_data(gdf, crs={'init': 'epsg:4326', 'no_defs': True}, dis=1):
	data = gdf
	if not data.crs:
		data.crs = crs
	else:
		crs = data.crs

	data = data.to_crs(epsg=3857)
	buffer = data.buffer(dis)
	data.set_geometry(buffer, inplace=True)
	data.to_crs(crs=crs, inplace=True)
	return data


def read_xzqh(path):
	shp_bjqh = GeoDataFrame.from_file(path, encoding='utf-8')
	shp_bjqh = shp_bjqh[['geometry', 'county']]
	shp_bjqh.to_crs(epsg=4326, inplace=True)
	return shp_bjqh


def intersection(gdf1, gdf2):
	geo_list = []
	county = []
	index_gdf1 = 0
	for geo in gdf1.geometry:
		index_gdf2 = 0
		for geo2 in gdf2.geometry:
			if geo.intersects(geo2):
				intersection_geo = geo.intersection(geo2)
				geo_list.append(intersection_geo)
				if 'county' in gdf1:
					county.append(gdf1.iloc[index_gdf1].county)
				elif 'county' in gdf2:
					county.append(gdf2.iloc[index_gdf2].county)
			index_gdf2 += 1
		index_gdf1 += 1

	gdf = GeoDataFrame(geometry=geo_list, crs=gdf1.crs)
	if county and len(county) > 0:
		gdf['county'] = county
	# print(gdf.head())
	return gdf


# def intersection(gdf1, gdf2):
# 	try:
# 		from geopandas.tools import overlay
# 		data = overlay(gdf1, gdf2, how='intersection', use_sindex=False)
# 		data.crs = gdf1.crs
# 		return data
# 	except Exception as e:
# 		print(str(e))
# 		return None


def intersection_xzqh(gdf):
	bjqh = read_xzqh(DATA_BASE_PATH + 'boundary/bj.shp')
	intersection_data = intersection(gdf, bjqh)
	return intersection_data


def join(gdf1, gdf2):
	return geopandas.sjoin(gdf1, gdf2, how="inner", op='intersects')


def cut_raster_by_geo(path, new_path, geo):
	src = rio.open(path)
	# gdf = geo.to_crs(epsg=src.crs)
	feature = [geo.__geo_interface__]
	out_image, out_transform = rio.mask.mask(src, feature, crop=True, nodata=src.nodata)
	out_meta = src.meta.copy()
	src.close()
	out_meta.update({"driver": "GTiff",
	                 "height": out_image.shape[1],
	                 "width": out_image.shape[2],
	                 "transform": out_transform})
	dest = rio.open(new_path, "w", **out_meta)
	dest.write(out_image)
	dest.close()


def get_new_raster_name(path):
	return os.path.join(os.path.dirname(path), "region", os.path.basename(path))


def get_region_geo():
	shp_data = GeoDataFrame.from_file(BJ_REGION)
	region = shp_data.loc[0].geometry
	return region


def cut_files_by_shp(files, shp_path):
	region = get_region_geo()
	for file in files:
		cut_raster_by_geo(file, get_new_raster_name(file), region)


def cut_dir_by_shp(dir, shp_path):
	files = filter(os.path.isfile, glob.glob(dir + "*"))
	cut_files_by_shp(files, shp_path)


def agg_raster_by_raster(path1, path2=os.path.join(TEMP_PATH, 'temp_3857.tif'), type='area'):
	raster1 = rio.open(path1)
	raster2 = rio.open(path2)

	bjqh = read_xzqh(DATA_BASE_PATH + 'boundary/bj.shp').to_crs(epsg=3857)
	print(bjqh.head())
	res_county = {}

	for i in range(0, len(bjqh)):
		geo = bjqh.iloc[i].geometry
		print(geo)
		county_name = bjqh.iloc[i].county
		print(county_name)
		feature = [geo.__geo_interface__]

		try:
			print('first')
			raster1_image, raster1_transform = rio.mask.mask(raster1, feature, crop=True, nodata=raster1.nodata)
			print('second')


			raster2_image, raster2_transform = rio.mask.mask(raster2, feature, crop=True, nodata=raster2.nodata)
			print('reproject')
			raster2_reproject = np.empty(shape=(raster1_image.shape[0],  # same number of bands
	                                             round(raster1_image.shape[1]),
	                                             round(raster1_image.shape[2])))
			reproject(
				raster2_image, raster2_reproject,
				src_transform=raster2_transform,
				dst_transform=raster1_transform,
				src_crs='+proj=latlong',
				dst_crs='+proj=latlong',
				resampling=Resampling.bilinear)

			# out_image_slope_reproject = np.where(out_image_slope_reproject == raster2.nodata, 1000, out_image_slope_reproject)
			# out_image_slope_reproject = out_image_slope_reproject.astype(int)
			for i in range(0, raster1_image.shape[1]):
				for j in range(0, raster1_image.shape[2]):
					data1 = raster1_image[0, i, j]
					data2 = raster2_reproject[0, i, j]
					print(data1, data2)
					if data2 != raster2.nodata and data2 < 0:
						if type == 'area':
							if data1 != raster2.nodata and data1 > 0:
								if county_name in res_county:
									res_county[county_name] = res_county[county_name] + 1
								else:
									res_county[county_name] = 1
						else:
							if county_name in res_county:
								res_county[county_name] = res_county[county_name] + data1
							else:
								res_county[county_name] = data1
		except Exception as e:
			print(str(e))
	return res_county


def agg_raster_by_gdf_value(path, gdf):
	# crs = '+proj=aea +lat_1=25 +lat_2=47 +lat_0=0 +lon_0=105 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs'
	with rio.open(path) as src:
		gdf = gdf.to_crs(epsg=3857)
		sum = {}
		for i in range(0, len(gdf)):
			geo = gdf.iloc[i].geometry
			try:
				feature = [GeoSeries(geo).__geo_interface__['features'][0]['geometry']]
				out_image, out_transform = rio.mask.mask(src, feature, crop=True, nodata=src.nodata)
				res = np.nansum(np.where(out_image == src.nodata, np.nan, out_image))
				if 'county' in gdf:
					county = gdf.iloc[i].county
					print(county)
					if county in sum.keys():
						sum[county] = sum.get(county) + res
					else:
						sum[county] = res
				elif 'all' in sum.keys():
					sum['all'] = sum.get('all') + res
				else:
					sum['all'] = res
			# for band in out_image:
			#     for row in band:
			#         for column in row:
			#             if column != src.nodata:
			#                 sum += column
			# print(sum)
			except Exception as e:
				print(str(e))
		return sum


def agg_raster_by_gdf_area(path, gdf):
	# crs = '+proj=aea +lat_1=25 +lat_2=47 +lat_0=0 +lon_0=105 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs'
	with rio.open(path) as src:
		gdf = gdf.to_crs(epsg=3857)
		sum = {}
		for i in range(0, len(gdf)):
			geo = gdf.iloc[i].geometry
			try:
				feature = [GeoSeries(geo).__geo_interface__['features'][0]['geometry']]
				out_image, out_transform = rio.mask.mask(src, feature, crop=True, nodata=src.nodata)
				out_image = np.where(out_image == 1, 1, np.nan)
				print(out_image)
				res = np.nansum(out_image) * src.width * src.height
				if 'county' in gdf:
					county = gdf.iloc[i].county
					if county in sum.keys():
						sum[county] = sum.get(county) + res
					else:
						sum[county] = res
				elif 'all' in sum.keys():
					sum['all'] = sum.get('all') + res
				else:
					sum['all'] = res
			except:
				pass
		return sum


def agg_prevent_disater_by_gdf_area(path, gdf):
	# crs = '+proj=aea +lat_1=25 +lat_2=47 +lat_0=0 +lon_0=105 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs'
	with rio.open(path) as src:
		gdf = gdf.to_crs(epsg=3857)
		sum = {}
		for i in range(0, len(gdf)):
			geo = gdf.iloc[i].geometry
			try:
				feature = [GeoSeries(geo).__geo_interface__['features'][0]['geometry']]
				out_image, out_transform = rio.mask.mask(src, feature, crop=True, nodata=src.nodata)
				out_image = np.where(out_image == 1, 1, np.nan)
				print(out_image)
				res = np.nansum(out_image) * src.width * src.height
				name = str(gdf.iloc[i].name)
				if name in sum.keys():
					sum[name] = sum.get(name) + res
				else:
					sum[name] = res
			except:
				pass
		return sum


def agg_gdf_by_gdf_area(gdf1, gdf2):
	intersection_data = intersection(gdf1, gdf2)
	area = intersection_data.to_crs(epsg=3857).area
	intersection_data['area'] = area
	intersection_data = intersection_data[['name', 'area']]
	# print(intersection_data.head())
	return intersection_data.groupby('name')['area'].sum()


def agg_raster_by_gdf_area_and_aspect_slope(path, path_aspect, path_slope, gdf):
	with rio.open(path) as src:
		with rio.open(path_aspect) as src1:
			with rio.open(path_slope) as src2:
				gdf = gdf.to_crs(epsg=3857)
				res_flat = 0
				res_30 = 0
				res_more_than_30 = 0
				for i in range(0, len(gdf)):
					geo = gdf.iloc[i].geometry
					try:
						feature = [GeoSeries(geo).__geo_interface__['features'][0]['geometry']]
						out_image, out_transform = rio.mask.mask(src, feature, crop=True, nodata=src.nodata)
						out_image_aspect, out_transform_aspect = rio.mask.mask(src1, feature, crop=True,
						                                                       nodata=src1.nodata)
						out_image_slope, _ = rio.mask.mask(src2, feature, crop=True, nodata=src2.nodata)
						out_image_aspect_reproject = np.empty(shape=(out_image.shape[0],  # same number of bands
						                                             round(out_image.shape[1]),
						                                             round(out_image.shape[2])))

						out_image_slope_reproject = np.empty(shape=(out_image.shape[0],  # same number of bands
						                                            round(out_image.shape[1]),
						                                            round(out_image.shape[2])))
						reproject(
							out_image_aspect, out_image_aspect_reproject,
							src_transform=out_transform_aspect,
							dst_transform=out_transform,
							src_crs='+proj=latlong',
							dst_crs='+proj=latlong',
							resampling=Resampling.bilinear)
						reproject(
							out_image_slope, out_image_slope_reproject,
							src_transform=out_transform_aspect,
							dst_transform=out_transform,
							# src_crs={'init': 'EPSG:3857'},
							# dst_crs={'init': 'EPSG:3857'},
							resampling=Resampling.bilinear)
						out_image_slope_reproject = np.where(out_image_slope_reproject == src1.nodata, 1000,
						                                     out_image_slope_reproject)
						out_image_slope_reproject = out_image_slope_reproject.astype(int)
						out_image_aspect_reproject = np.where(out_image_aspect_reproject == src2.nodata, 1000,
						                                      out_image_aspect_reproject)
						out_image_aspect_reproject = out_image_aspect_reproject.astype(int)
						for i in range(0, out_image.shape[1]):
							for j in range(0, out_image.shape[2]):
								aspect = out_image_aspect_reproject[0, i, j]
								slope = out_image_slope_reproject[0, i, j]
								print(aspect, slope, out_image[0, i, j])
								if out_image[0, i, j] > 0:
									if slope <= 0:
										print('slope')
										if aspect < 0:
											res_flat += 1
									elif slope < 30:
										res_30 += 1
									else:
										res_more_than_30 += 1
						print(res_flat, res_30, res_more_than_30)
					except Exception as e:
						print(str(e))
				return (res_flat, res_30, res_more_than_30)


def merge_gdf_geo_to_geojson(gdf):
	return gdf.to_json()


def reproject_by_gdal(src_path, dst_path, dst_src='EPSG:3857'):
	import gdal
	print('begin warp')
	ds = gdal.Warp(dst_path, src_path, dstSRS=dst_src)
	print('warp')
	ds = None
	# print(ds)


if __name__ == '__main__':
	# cut_dir_by_shp('/home/g214/data_from_chenhao/data_analyse/station/interpolation/rain/', '/home/g214/data_from_chenhao/data_analyse/boundary/bj_region/bj_region.shp')
	cut_dir_by_shp('/home/g214/data_from_chenhao/data_analyse/station/interpolation/temp/',
	               '/home/g214/data_from_chenhao/data_analyse/boundary/bj_region/bj_region.shp')
# cut_dir_by_shp('../data/', '../bj_region/bj_region.shp')