ibgetools

#!/usr/bin/python

import argparse
import fnmatch
import os
import sys

from concurrent.futures import ProcessPoolExecutor

from IBGETools.Geometry import RegionFactory
from IBGETools.Map import Map, MapFactory
from IBGETools.Utils import KMLFileWriter, TileScriptFileWriter


def ScanDirectoryForPDFs(directory):
    pdf_paths = []

    for root, _, filenames in os.walk(directory):
        # All IBGE PDFs we are interested at start with a number,
        # so we can rule out a few by filtering the file name.
        for filename in fnmatch.filter(filenames, "[0-9]*.pdf"):
            pdf_paths.append(os.path.join(root, filename))

    return pdf_paths


def GetMapFromPath(id, map_path):
    ibge_map = MapFactory(map_path)

    if not ibge_map:
        print >> sys.stderr, "%s [Error] PDF not supported: %s" % (id, map_path)
        return ibge_map

    if not ibge_map.IsValid():
        print >> sys.stderr, "%s [Error] Could not parse: %s" % (id, map_path)
        return None

    return ibge_map


def CheckForDirectory(path):
    if not os.path.isdir(path):
        raise argparse.ArgumentTypeError("%s is not a directory." % path)

    return path


def ProcessMap(executor_args):
    map_path, id, save_as_png, verbose = executor_args

    if verbose:
        print "%s [Debug] Processing: %s" % (id, map_path)

    ibge_map = GetMapFromPath(id, map_path)
    if not ibge_map:
        return None

    basename = os.path.splitext(os.path.basename(map_path))[0]

    if (save_as_png):
        ibge_map.SaveMapImageAsPNG("%s_%s" % (id, basename))
    else:
        ibge_map.SaveMapImageAsTIFF("%s_%s" % (id, basename))

    # This will free the resources and make the Map class "picklable".
    ibge_map.Dispose()

    if verbose:
        print "%s [Debug] Processed: %s" % (id, map_path)

    return ibge_map


def ProcessKML(args, regions_list):
    if not regions_list:
        print >> sys.stderr, "%s [Error] Fatal error when parsing." % args.id
        return

    region = reduce(lambda r1, r2: r1.Merge(r2), regions_list)

    output = open("%s.kml" % args.id, "w")
    KMLFileWriter(output, args.id, region)
    output.close()


def ProcessMBTiles(args, regions_list):
    if not regions_list:
        print >> sys.stderr, "%s [Error] Fatal error when parsing." % args.id
        return

    if not args.regions:
        regions_list = [reduce(lambda r1, r2: r1.Merge(r2), regions_list)]

    output = open("%s.sh" % args.id, "w")

    for region in regions_list:
        region_id = "%03d" % regions_list.index(region)
        TileScriptFileWriter(output, args.id, region_id, args.upload, region)

    output.close()

def Main():
    parser = argparse.ArgumentParser(
        description="Generates KML or map tiles from IBGE PDFs.")

    group = parser.add_mutually_exclusive_group(required=True)

    parser.add_argument("-v", action="store_true", help= "Verbose mode.")
    parser.add_argument("--jobs", type=int, default=1, help=
            "Number of jobs for processing PDFs, be aware of memory usage.")
    group.add_argument("--kml", action="store_true", help=
            "Generates PNGs and a .kml file that can be imported to JOSM "
            "(via PicLayer plugin) or opened by Google Earth.")
    group.add_argument("--tiles", action="store_true", help=
            "Generates TIFFs and a shell script that will generate map tiles "
            "using GDAL tools and TileMill when executed.")
    parser.add_argument("--regions", action="store_true", help=
            "Create separated .mbtiles for each district within a "
            "municipality. Might speed up the processing of big areas.")
    parser.add_argument("--upload", action="store_true", help=
            "The generated shell script will also upload the tiles to the "
            "cloud. Make sure to have ~/.tilemill/config.json properly set "
            "with auth keys.")
    parser.add_argument("directory", nargs=1, type=CheckForDirectory, help=
            "Directory to be recursively scanned for IBGE PDFs files.")
    parser.add_argument("--id", required=True, help=
            "An identifier to be used as prefix for the generated files.")

    args = parser.parse_args()

    directory = args.directory[0]

    map_paths = ScanDirectoryForPDFs(directory)

    maps_list = []
    params = [(map_path, args.id, args.kml, args.v) for map_path in map_paths]

    if args.jobs == 1:
        for param in params:
            ibge_map = ProcessMap(param)
            if ibge_map:
                maps_list.append(ibge_map)
    else:
        executor = ProcessPoolExecutor(max_workers=args.jobs)
        for ibge_map in executor.map(ProcessMap, params):
            if ibge_map:
                maps_list.append(ibge_map)

    if args.v:
        ratio = len(maps_list) / float(len(map_paths)) * 100.
        print "%s [Debug] Maps processed: %.2f%%" % (args.id, ratio)

    regions_list = RegionFactory(maps_list)
    if args.kml:
        ProcessKML(args, regions_list)
    else:
        ProcessMBTiles(args, regions_list)

if __name__ == "__main__":
    Main()