#!/usr/bin/python
# -*- coding: utf-8 -*-
#"POLYGON((-89.65513229370117 21.048938357786966,-89.62852478027342 21.047816903874505,-89.6261215209961 21.03339745836918,-89.6465492248535 21.022822311328852,-89.65873718261717 21.039325621609095,-89.65513229370117 21.048938357786966))"
import osgeo.ogr as ogr
import osgeo.osr as osr
import numpy as np
import gdal
import sys
from collections import Counter
#-------------------------------------------------------------------------------
def getPolygonArray(image, shp_filename):
	rows = image.RasterYSize
	cols = image.RasterXSize
	geotransform =  image.GetGeoTransform()
	projection = image.GetProjection()
	sr = osr.SpatialReference(wkt=projection)

	sr4326 = osr.SpatialReference()
	sr4326.ImportFromEPSG(4326)


	driverMEM	= gdal.GetDriverByName('MEM')
	driverMemory	= ogr.GetDriverByName('Memory')
	driverTiff	= gdal.GetDriverByName('GTiff')
	driverShape	= ogr.GetDriverByName("ESRI Shapefile")

	#RASTER
	target_ds =  driverMEM.Create('raster.tif', cols, rows, 1, gdal.GDT_Byte) #MEM
	target_ds.GetRasterBand(1).SetNoDataValue(0)
	target_ds.GetRasterBand(1).WriteArray(np.zeros((rows, cols), dtype=np.uint8))
	target_ds.SetProjection(projection)
	target_ds.SetGeoTransform(geotransform)
	target_ds.GetRasterBand(1).FlushCache()

	#VECTOR
	dataSource = driverShape.Open(shp_filename, 0)
	# Check to see if shapefile is found.
	if dataSource is None:
	    print('Could not open', shp_filename)
	else:
	    print('Opened %s', shp_filename)
	    layer = dataSource.GetLayer()
	    featureCount = layer.GetFeatureCount()
	    print("Number of features in", shp_filename, featureCount)
#	rast_ogr_ds = driverMemory.CreateDataSource( 'myshape.shp' )
#	rast_mem_lyr = rast_ogr_ds.CreateLayer( 'POLYGON', sr4326,  ogr.wkbPolygon)
#	feat = ogr.Feature( rast_mem_lyr.GetLayerDefn() )
#	feat.SetGeometry( ogr.CreateGeometryFromWkt(wkt_geom) )
#	rast_mem_lyr.CreateFeature( feat )


	err = gdal.RasterizeLayer(target_ds, [1], layer, burn_values=[255])

	if err != 0:
		print(err)
		gdaltest.post_reason( 'got non-zero result code from RasterizeLayer' )
		return 'fail'

	target_ds.GetRasterBand(1).FlushCache()
	np_polygon = target_ds.GetRasterBand(1).ReadAsArray(0,0,cols,rows)
	count = Counter([])
#	for r in np_polygon:
#		count = count + Counter(r)
#	colorDict = {str(k):v for k,v in dict(count).items()}
#	print(colorDict)
	target_ds = None

	return np_polygon
#-------------------------------------------------------------------------------
def rasterWkt(shpfile, tifffile, outputfile):
	''' 
	Draw WKT Polygon into a Image
	'''

	inputImage = gdal.Open(tifffile)
	rows = inputImage.RasterYSize
	cols = inputImage.RasterXSize
	geotransform =  inputImage.GetGeoTransform()
	projection = inputImage.GetProjection()

	sr = osr.SpatialReference()
	sr.ImportFromWkt(projection)

	driverTiff = gdal.GetDriverByName('GTiff')
	output = driverTiff.Create(outputfile, cols, rows, inputImage.RasterCount, gdal.GDT_Byte)

	np_polygon = getPolygonArray(inputImage, shpfile)

	for b in range(1, output.RasterCount+1):
		np_input = inputImage.GetRasterBand(b).ReadAsArray(0,0,cols,rows)
		output.GetRasterBand(b).SetNoDataValue(0)
		output.GetRasterBand(b).WriteArray(np_polygon)#np.bitwise_and(np_polygon, np_input))
		output.GetRasterBand(b).FlushCache()

	output.SetGeoTransform(geotransform)
	output.SetProjection(projection)

	#Close files
	polygonRaster = None
	output = None
	inputImage = None

def main(argv):
	rasterWkt(argv[1], argv[2], argv[3])

if __name__ == "__main__":
	main(sys.argv)
