#!/usr/bin/env python
###############################################################################
# $Id$
#
# Project:  OGR Python samples
# Purpose:  Assemble TIGER area landmarks.
# Author:   Schuyler Erle <schuyler@geocoder.us>
#
###############################################################################
# Copyright (c) 2005, Frank Warmerdam <warmerdam@pobox.com>
# 
# 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 osr
import ogr
import string
import sys, os

#############################################################################
def Usage():
    print 'Usage: tigerpoly.py infile [outfile].shp'
    print
    sys.exit(1)

#############################################################################
# Argument processing.

infile = None
outfile = None

i = 1
while i < len(sys.argv):
    arg = sys.argv[i]
    if infile is None:
	infile = arg
    else:
	Usage()
    i = i + 1

if infile is None:
    Usage()

#############################################################################
# Open the datasource to operate on.

ds = ogr.Open( infile, update = 0 )

area_layer = ds.GetLayerByName( 'AreaLandmarks' )
land_layer = ds.GetLayerByName( 'Landmarks' )

#############################################################################
#	Create output file for the composed polygons.

shp_driver = ogr.GetDriverByName( 'ESRI Shapefile' )

def copy_layer (layer, geom_type, outfile):
    if os.access(outfile, os.R_OK):
       shp_driver.DeleteDataSource( outfile )

    shp_ds = shp_driver.CreateDataSource( outfile )
    shp_layer = shp_ds.CreateLayer( 'out', geom_type = geom_type )

    src_defn = layer.GetLayerDefn()
    field_count = src_defn.GetFieldCount()

    for fld_index in range(field_count):
	src_fd = src_defn.GetFieldDefn( fld_index )
	
	fd = ogr.FieldDefn( src_fd.GetName(), src_fd.GetType() )
	fd.SetWidth( src_fd.GetWidth() )
	fd.SetPrecision( src_fd.GetPrecision() )
	shp_layer.CreateField( fd )

    return shp_layer

def copy_feature (feat, geom, layer):
    defn = layer.GetLayerDefn()
    feat2 = ogr.Feature(feature_def = defn)

    for fld_index in range(defn.GetFieldCount()):
       feat2.SetField( fld_index, feat.GetField( fld_index ) )

    feat2.SetGeometryDirectly( geom )

    layer.CreateFeature( feat2 )
    feat2.Destroy()

area_out_layer = copy_layer(land_layer, ogr.wkbPolygon, "arealandmarks.shp")
land_out_layer = copy_layer(land_layer, ogr.wkbPoint,   "landmarks.shp")

polygon  = {}
landmark = {}
poly_line_link = {}
area_defn = area_layer.GetLayerDefn()
land_id_field = area_defn.GetFieldIndex( 'LAND' )
poly_id_field = area_defn.GetFieldIndex( 'POLYID' )

feat = area_layer.GetNextFeature()
while feat is not None:
    land_id = feat.GetField( land_id_field )
    poly_id = feat.GetField( poly_id_field )

    if not landmark.has_key(poly_id):
	landmark[poly_id] = []
    landmark[poly_id].append( land_id )

    if not polygon.has_key(land_id):
	polygon[land_id] = []
    polygon[land_id].append( poly_id )

    poly_line_link[land_id] = {}
    feat = area_layer.GetNextFeature()

print "Identified %d polygons in %d area landmarks." % (
    len(landmark), len(polygon))

#############################################################################
# Read all features in the line layer, holding just the geometry in a hash
# for fast lookup by TLID.

line_layer = ds.GetLayerByName( 'CompleteChain' )
lines = {}

print "Reading lines..."

feat = line_layer.GetNextFeature()
geom_id_field = feat.GetFieldIndex( 'TLID' )
tile_ref_field = feat.GetFieldIndex( 'MODULE' )
while feat is not None:
    geom_id = feat.GetField( geom_id_field )
    lines[geom_id] = feat.GetGeometryRef().Clone()
    feat = line_layer.GetNextFeature()

print 'Got %d lines.' % (len(lines))

#############################################################################
# Read all polygon/chain links and build a hash keyed by POLY_ID listing
# the chains (by TLID) attached to it. 

print "Reading polygon links..."

link_layer = ds.GetLayerByName( 'PolyChainLink' )

link_defn = link_layer.GetLayerDefn()

geom_id_field = link_defn.GetFieldIndex( 'TLID' )
lpoly_field = link_defn.GetFieldIndex( 'POLYIDL' )
rpoly_field = link_defn.GetFieldIndex( 'POLYIDR' )

polyid_field   = area_defn.GetFieldIndex( 'POLYID' )

link_count = 0

link_layer.ResetReading()
link = link_layer.GetNextFeature()
while link is not None:
    tlid = link.GetField( geom_id_field )
    lpoly_id = link.GetField( lpoly_field )
    rpoly_id = link.GetField( rpoly_field )

    # if lpoly_id == rpoly_id:
	# link = link_layer.GetNextFeature()
	# continue

    for poly_id, side in ((rpoly_id, "right"), (lpoly_id, "left")):
	if not landmark.has_key(poly_id):
	    continue

	# print "Checking %s poly_id %d against landmarks %s" % (
	#    side, poly_id, str(landmark[poly_id]))
	for land_id in landmark[poly_id]:
	    # print "Adding TLID %d to landmark %d" % (tlid, land_id)
	    if poly_line_link[land_id].has_key(tlid):
		poly_line_link[land_id][tlid] += 1
	    else:
		poly_line_link[land_id][tlid] = 1

    link_count = link_count + 1
    link = link_layer.GetNextFeature()

print 'Got %d relevant links.' % link_count

#############################################################################
# Process all polygon features.

poly_count = 0
point_count = 0

print 'Collecting landmark edges...'
defn = land_layer.GetLayerDefn()
fields = {}
for f in ( 'LAND', 'LALAT', 'LALONG' ):
    fields[f] = defn.GetFieldIndex( f )

feat = land_layer.GetNextFeature()
out_of_bounds = 0
while feat is not None:
    data = {}
    for key, field in fields.iteritems():
	data[key] = feat.GetField(field)

    land_id = data['LAND']

    if poly_line_link.has_key(land_id):
	link_coll = ogr.Geometry( type = ogr.wkbGeometryCollection )

	if len( poly_line_link[land_id] ) < 3:
	    out_of_bounds += 1
	    feat = land_layer.GetNextFeature()
	    continue

	for tlid, count in poly_line_link[land_id].iteritems():
	    if count == 1:
		link_coll.AddGeometry(lines[tlid])

	try:
	    geom = ogr.BuildPolygonFromEdges( link_coll, 1, 1 )
	    copy_feature( feat, geom, area_out_layer )
	    poly_count += 1
	except:
	    raise
    elif data['LALAT'] and data['LALONG']:
	geom = ogr.CreateGeometryFromWkt( 'POINT(%s %s)' % (data['LALONG'], data['LALAT']) )
	copy_feature( feat, geom, land_out_layer )
	point_count += 1

    feat = land_layer.GetNextFeature()

if out_of_bounds:
    print 'Discarded %d incomplete polygons.' % out_of_bounds

print 'Built %d polygons and %d points.' % (poly_count, point_count)

#############################################################################
# Cleanup

ds.Destroy()