#!/usr/bin/env python
#
###############################################################################
# $Id: r2v.py
#
# Purpose:  convert a GDAL-readable raster to a OGR-generated vector layer
# Author:   dustymugs, dustymugs@dspiral.net
#
# Lots of this written by referring to the gdal/ogr python examples.
# Some of this (the usage and processing of command args) is from gdal2grd.py
#
###############################################################################
#
###############################################################################
# Copyright (c) 2003, Andrey Kiselev <dron@remotesensing.org>
# Copyright (c) 2007, dustymugs <dustymugs@dspiral.net>
# 
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.
###############################################################################

import gdal, ogr, osr
from gdalconst import *
from os.path import isfile
import sys, struct, time

class Point:
	def __init__(self, x = None, y = None):
		self.x = x
		self.y = y

class Polygon:
	def __init__(self):
		self.points = []
		self.value = None
		self.id = None

	def addPoint(self, x, y):
		self.points.append(Point(x, y))

	def setValue(self, v):
		self.value = v

	def setId(self, v):
		self.id = v

def fieldExist(ldfn, fname):
	for i in range (ldfn.GetFieldCount()):
		if ldfn.GetFieldDefn(i).GetName() == fname:
			return True
	return False

def Usage():
	print 'Usage: r2v.py [-b band] [-n nodata] [-srs srs_string] [-f format_name]'
	print '         [-an attribute_name] [-ln layer_name] [-append] [-quiet] [-group]'
	print '         fileIn fileOut'
	print 'Convert a raster to polygons.  Each pixel becomes a four point polygon.'
	print
	print '  -b band               select a band number to convert (1 based)'
	print '  -n nodata             no data value'
	print '  -srs srs_string       string or filename, string or file contents must be WKT'
	print '  -f format_name        output format name, run \'ogr2ogr\' for possible values'
	print '  -an attribute_name    column name in layer, use "" for multi-word name'
	print '  -ln layer_name        output layer name, use "" for multi-word name'
	print '  -append               append polygons to layer'
	print '  -quiet                do not report any diagnostic information'
	print '  -group                group neighboring polygons with same value (NOT YET IMPLEMENTED)'
	print '  fileIn                name of the input GDAL supported file or stream'
	print '  fileOut               name of the output file or stream'
	print
	sys.exit(1)

# dictionaries for translation
typeRasterData = {'byte' : 'c', 'int16' : 'h', 'int32' : 'i', 'float32' : 'f', 'float64' : 'd'}
typeRasterToVectorData = {'byte' : ogr.OFTBinary, 'int16' : ogr.OFTInteger, 'int32' : ogr.OFTInteger, 'float32' : ogr.OFTReal, 'float64' : ogr.OFTReal}
orderPolyPoint = ['ll', 'ul', 'ur', 'lr']

# required variables from user
fileIn = None
fileOut = None
ogrFormat = 'ESRI Shapefile'

# optional variables from user
layerName = 'New Layer'
layerAppend = 0
srsString = None
attrName = 'value'
bandNum = 1
quiet = 0
dataEmpty = None
groupPoly = 0

# internal variables
layerPoly = []
skipPoly = 0

# object counting
objTotal = 0
objDone = 0
objPDone = 0
objLastPDone = 0

# Parse command line arguments.
i = 1
while i < len(sys.argv):
	arg = sys.argv[i]

	# -b band
	if arg == '-b':
		i = i + 1
		bandNum = int(sys.argv[i])
	# -n nodata
	elif arg == '-n':
		i = i + 1
		dataEmpty = int(sys.argv[i])
	# -f format_name
	elif arg == '-f':
		i = i + 1
		ogrFormat = sys.argv[i]
	# -srs srs_string
	elif arg == '-srs':
		i = i + 1
		srsString = sys.argv[i]
	# -an attribute_name
	elif arg == '-an':
		i = i + 1
		attrName = sys.argv[i]
	# -ln layer_name
	elif arg == '-ln':
		i = i + 1
		layerName = sys.argv[i]
	# -append
	elif arg == '-append':
		layerAppend = 1
	# -quiet
	elif arg == '-quiet':
		quiet = 1
	# -group
	elif arg == '-group':
		groupPoly = 1
	# fileIn
	elif fileIn is None:
		fileIn = arg
	# fileOut
	elif fileOut is None:
		fileOut = arg
	else:
		Usage()

	i = i + 1

# file/stream strings not provided
if fileIn is None or fileOut is None:
	Usage()

# load raster
dsIn = gdal.Open(fileIn, GA_ReadOnly)
if dsIn is None:
	print 'Cannot open: %s' % fileIn
	sys.exit(2)

# get transformation and band info
geotransform = dsIn.GetGeoTransform()
band = dsIn.GetRasterBand(bandNum)

# band not found, fail
if band is None:
	print 'Cannot load band %s from %s' % (bandNum, fileIn)
	sys.exit(2)

# output details of raster
if not quiet:
	print 'Size is %s x %s x %s' % (dsIn.RasterXSize, dsIn.RasterYSize, dsIn.RasterCount)
	print 'Projection is', dsIn.GetProjection()
	print 'Origin = (%s, %s)' % (geotransform[0], geotransform[3])
	print 'Pixel Size = (%s, %s)' % (geotransform[1], geotransform[5])
	print 'Band = %s, Type = %s' % (bandNum, gdal.GetDataTypeName(band.DataType))

# output start of raster processing
if not quiet:
	objTotal = dsIn.RasterYSize * dsIn.RasterXSize
	print 'Processing %s pixels...' % (objTotal)

# process raster by height
for i in range(dsIn.RasterYSize):
	# get line in row
	buf = band.ReadRaster(0, i, dsIn.RasterXSize, 1);
	# datatype unknown, skip additional processing
	if (not typeRasterData.has_key(gdal.GetDataTypeName(band.DataType).lower())):
		if not quiet:
			print
		print 'Unknown DataType:', gdal.GetDataTypeName
		sys.exit(2)

	# split line into elements of a tuple
	data = struct.unpack(str(dsIn.RasterXSize) + typeRasterData[gdal.GetDataTypeName(band.DataType).lower()], buf)

	# empty polygon
	lp = {
		'll' : {'x' : None, 'y' : None},
		'lr' : {'x' : None, 'y' : None},
		'ul' : {'x' : None, 'y' : None},
		'ur' : {'x' : None, 'y' : None},
	}

	# Y Axis corner values
	lp['lr']['y'] = lp['ll']['y'] = geotransform[3] + (i * geotransform[5])
	lp['ur']['y'] = lp['ul']['y'] = lp['ll']['y'] + geotransform[5]

	# process each line element
	for j in range(dsIn.RasterXSize):
		dataVal = data[j]

		# skip processing of those values that are indicated as empty
		if (dataEmpty is not None) and (dataVal == float(dataEmpty)):
			skipPoly += 1
		else:
			# new polygon object
			poly = Polygon()

			# X Axis corner values
			lp['ul']['x'] = lp['ll']['x'] = geotransform[0] + (j * geotransform[1])
			lp['ur']['x'] = lp['lr']['x'] = lp['ll']['x'] + geotransform[1]

			# assign polygon points
			for k in orderPolyPoint:
				poly.addPoint(lp[k]['x'], lp[k]['y'])

			# assign polygon value
			poly.setValue(dataVal)

			# assign polygon id
			poly.setId((i * dsIn.RasterXSize) + (j + 1))

			# add polygon to set of polygons
			layerPoly.append(poly)

		if not quiet:
			objDone = float((i * dsIn.RasterXSize) + (j + 1))
			objPDone = int((objDone / float(objTotal)) * 100)
			if (objPDone > objLastPDone) and ((objPDone % 10) == 0):
				print '%% of Pixels processed: %s' % (objPDone)
				objLastPDone = objPDone

# clean up input
# band is still needed for output, so no delete
del lp, data, dsIn

if not quiet:
	print
	print 'Done processing pixels'
	print 'Pixels skipped: %s' % (skipPoly)
	print

objTotal = len(layerPoly)
objLastPDone = 0

# no polygons from raster, no need to continue
if (objTotal < 1):
	print 'No polygons collected.  Vector dataset cannot be created.'
	sys.exit(2)
elif not quiet:
	print 'Creating vector layer from %s polygons...' % (objTotal)

# output polygons to vector layer
dOut = ogr.GetDriverByName(ogrFormat)
dsOut = dOut.CreateDataSource(fileOut)

# unable to open output
if dsOut is None:
	print 'Cannot open: %s' % fileOut
	sys.exit(2)

# try to create layer with spatial reference
try:
	if srsString is None:
		raise

	# create spatial reference object
	osrOut = osr.SpatialReference()

	# if srsString is a file, read file
	if isfile(srsString):
		f = open(srsString, 'r')
		srsString = f.read()
		f.close()
	# set spatial reference
	osrOut.ImportFromWkt(srsString)

	# create layer
	if not layerAppend:
		# create layer WITH spatial reference
		lOut = dsOut.CreateLayer(layerName, osrOut, ogr.wkbPolygon)
	# append to layer
	else:
		lOut = dsOut.GetLayerByName(layerName)
# failed to load spatial reference
except:
	# create layer
	if not layerAppend:
		# create layer WITHOUT spatial reference
		lOut = dsOut.CreateLayer(layerName, None, ogr.wkbPolygon)
	# append to layer
	else:
		lOut = dsOut.GetLayerByName(layerName)

# attribute column exists?
if not fieldExist(lOut.GetLayerDefn(), attrName):
	# create attribute column in layer
	# this is where the value goes
	colOut = ogr.FieldDefn(attrName, typeRasterToVectorData[gdal.GetDataTypeName(band.DataType).lower()])
	lOut.CreateField(colOut)

# get layer schema
sOut = lOut.GetLayerDefn()

# get index of attribute column
attrIndex = sOut.GetFieldIndex(attrName)

# process each polygon
for i in range(objTotal):
	tp = layerPoly[i]

	f = ogr.Feature(sOut)
	g = ogr.Geometry(ogr.wkbPolygon)
	r = ogr.Geometry(ogr.wkbLinearRing)

	# add points to ring
	for j in tp.points:
		r.AddPoint_2D(j.x, j.y)

	# add ring to geometry
	g.AddGeometry(r)
	g.CloseRings()

	# set geometry of feature
	f.SetGeometry(g)
	# assign feature value
	f.SetField(attrIndex, tp.value)
	# add feature to layer
	lOut.CreateFeature(f)

	if not quiet:
		objDone = float(i + 1)
		objPDone = int((objDone / float(objTotal)) * 100)
		if (objPDone > objLastPDone) and ((objPDone % 10) == 0):
			print '%% of Polygons processed: %s' % (objPDone)
			objLastPDone = objPDone

if not quiet:
	print 'Done processing polygons.'
	print

lOut.SyncToDisk()

if not quiet:
	print 'Raster has been converted to a Vector.'