libpostal/scripts/geodata/polygons/reverse_geocode.py

import argparse
import os
import sys

this_dir = os.path.realpath(os.path.dirname(__file__))
sys.path.append(os.path.realpath(os.path.join(os.pardir, os.pardir)))

from geodata.polygons.index import *
from geodata.encoding import safe_decode


class ReverseGeocoder(RTreePolygonIndex):
    COUNTRIES_FILENAME = 'qs_adm0.shp'
    ADMIN1_FILENAME = 'qs_adm1.shp'
    ADMIN1_REGION_FILENAME = 'qs_adm1_region.shp'
    ADMIN2_FILENAME = 'qs_adm2.shp'
    ADMIN2_REGION_FILENAME = 'qs_adm2_region.shp'
    LOCAL_ADMIN_FILENAME = 'qs_localadmin.shp'
    LOCALITIES_FILENAME = 'qs_localities.shp'
    NEIGHBORHOODS_FILENAME = 'qs_neighborhoods.shp'

    sorted_levels = ('adm0',
                     'adm1',
                     'adm1_region',
                     'adm2',
                     'adm2_region',
                     'localadmin',
                     'locality',
                     'neighborhood',
                     )

    sort_levels = {k: i for i, k in enumerate(sorted_levels)}

    polygon_properties = {
        COUNTRIES_FILENAME: {
            'qs_a0': safe_decode,
            'qs_iso_cc': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        ADMIN1_FILENAME: {
            'qs_a1': safe_decode,
            'qs_a1_lc': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        ADMIN1_REGION_FILENAME: {
            'qs_a1r': safe_decode,
            'qs_a1r_lc': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        ADMIN2_FILENAME: {
            'qs_a2': safe_decode,
            'qs_a2_lc': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        ADMIN2_REGION_FILENAME: {
            'qs_a2r': safe_decode,
            'qs_a2r_lc': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        LOCAL_ADMIN_FILENAME: {
            'qs_la': safe_decode,
            'qs_la_lc': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        LOCALITIES_FILENAME: {
            'qs_loc': safe_decode,
            'qs_loc_alt': safe_decode,
            'qs_level': safe_decode,
            'qs_gn_id': int,
            'qs_woe_id': int,
        },
        NEIGHBORHOODS_FILENAME: {
            'name': safe_decode,
            'name_en': safe_decode,
            'qs_level': safe_decode,
            'woe_id': int,
            'gn_id': int,
        }
    }

    @classmethod
    def create_from_shapefiles(cls,
                               input_files,
                               output_dir,
                               index_filename=DEFAULT_INDEX_FILENAME,
                               polys_filename=DEFAULT_POLYS_FILENAME):

        index = cls(save_dir=output_dir, index_filename=index_filename)

        i = 0

        for input_file in input_files:
            f = fiona.open(input_file)

            include_props = cls.polygon_properties.get(input_file)

            filename = os.path.split(input_file)[-1]

            for rec in f:
                if not rec or not rec.get('geometry') or 'type' not in rec['geometry']:
                    continue

                properties = rec['properties']

                if filename == cls.NEIGHBORHOODS_FILENAME:
                    properties['qs_level'] = 'neighborhood'

                if include_props:
                    for k, func in include_props.iteritems():
                        v = properties.get(k, None)
                        if v is not None:
                            properties[k] = func(v)

                include_props = set(include_props.keys())

                poly_type = rec['geometry']['type']
                if poly_type == 'Polygon':
                    poly = Polygon(rec['geometry']['coordinates'][0])
                    index.index_polygon(i, poly)
                    poly = index.simplify_polygon(poly)
                    index.add_polygon(poly, dict(rec['properties']), include_only_properties=include_props)
                elif poly_type == 'MultiPolygon':
                    polys = []
                    for coords in rec['geometry']['coordinates']:
                        poly = Polygon(coords[0])
                        polys.append(poly)
                        index.index_polygon(i, poly)

                    multi_poly = index.simplify_polygon(MultiPolygon(polys))
                    index.add_polygon(multi_poly, dict(rec['properties']), include_only_properties=include_props)
                else:
                    continue

                i += 1
        return index

    @classmethod
    def create_with_quattroshapes(cls, quattroshapes_dir,
                                  output_dir,
                                  index_filename=DEFAULT_INDEX_FILENAME,
                                  polys_filename=DEFAULT_POLYS_FILENAME):

        admin0_filename = os.path.join(quattroshapes_dir, cls.COUNTRIES_FILENAME)
        admin1_filename = os.path.join(quattroshapes_dir, cls.ADMIN1_FILENAME)
        admin1r_filename = os.path.join(quattroshapes_dir, cls.ADMIN1_REGION_FILENAME)
        admin2_filename = os.path.join(quattroshapes_dir, cls.ADMIN2_FILENAME)
        admin2r_filename = os.path.join(quattroshapes_dir, cls.ADMIN2_REGION_FILENAME)
        local_admin_filename = os.path.join(quattroshapes_dir, cls.LOCAL_ADMIN_FILENAME)
        localities_filename = os.path.join(quattroshapes_dir, cls.LOCALITIES_FILENAME)
        neighborhoods_filename = os.path.join(quattroshapes_dir, cls.NEIGHBORHOODS_FILENAME)

        return cls.create_from_shapefiles([admin0_filename, admin1_filename, admin1r_filename,
                                          admin2_filename, admin2r_filename, local_admin_filename,
                                          localities_filename, neighborhoods_filename],
                                          output_dir, index_filename=index_filename,
                                          polys_filename=polys_filename)

    def sort_level(self, i):
        props, p = self.polygons[i]
        return self.sort_levels.get(props['qs_level'], 0)

    def get_candidate_polygons(self, lat, lon):
        candidates = OrderedDict.fromkeys(self.index.intersection((lon, lat, lon, lat))).keys()
        return sorted(candidates, key=self.sort_level, reverse=True)


if __name__ == '__main__':
    # Handle argument parsing here
    parser = argparse.ArgumentParser()

    parser.add_argument('-q', '--quattroshapes-dir',
                        help='Path to quattroshapes dir')

    parser.add_argument('-o', '--out-dir',
                        default=os.getcwd(),
                        help='Output directory')

    args = parser.parse_args()
    index = ReverseGeocoder.create_with_quattroshapes(args.quattroshapes_dir, args.out_dir)
    index.save()