"""Create a csv matrix of distances between shapefile geometry objects.

Requirements: fiona, shapely

Written by: Taylor Denouden
Date: November 25, 2015
"""

import fiona
from shapely.geometry import shape
from multiprocessing import Pool


def extract_ids(input_file):
    """Extract all polygon ids from input shapefile."""
    with fiona.open(input_file, 'r') as source:
        return [shp['id'] for shp in source]


def calculate_distances(shp_id):
    """Calculate distances between `shp_id` and all other features in shapefile.

    `shp_id` is the feature id in the shapefile `infile` (global). `shp_id` and
    each id in `ids` (global) are cast to shapely features and then passed to a
    distance function contained in the shapely library. An array is returned
    with all these distances.
    """
    with fiona.open(infile) as source:
        source = list(source)
        result = []

        # Calculate distances and store in result
        i_shp = shape(source[int(shp_id)]['geometry'])
        for j in ids:
            if int(j) < int(shp_id):
                result.append(-1)
                continue

            j_shp = shape(source[int(j)]['geometry'])
            if i_shp.is_valid and j_shp.is_valid:
                result.append(i_shp.distance(j_shp))
            else:
                result.append(-1)

        return result


if __name__ == "__main__":
    """Main execution thread."""
    # infile = "./data/random_points/test_polys.shp"
    infile = "./data/low_water_final/low_water.shp"
    ids = extract_ids(infile)

    print "Calculating distances"
    # Map each id to a list of distances to all other ids
    pool = Pool()  # Spawn a worker process for each cpu
    data = pool.map(calculate_distances, ids, chunksize=100)

    # Write the data to a new csv file
    outfile = open("test.csv", "w")

    # Write header
    print "Writing Header"
    outfile.write("NODE")
    for i in ids:
        outfile.write("," + i)
    outfile.write("\n")

    # Write rows
    print "Writing Rows"
    for i in ids:
        outfile.write(str(i) + ",")
        i = int(i)

        for j in ids:
            j = int(j)
            outfile.write(str(data[i][j]) + ",")
        outfile.write("\n")

    outfile.close()

@@ -440,22 +440,20 @@
 stances(
-shp_id
+args
 ):%0A    %22
@@ -456,32 +456,36 @@
    %22%22%22Calculate 
+the 
 distances betwee
@@ -490,30 +490,80 @@
 een 
-%60shp_id%60 and all other
+all %60ids1%60 to %60ids2%60 in %60infile%60.%0A%0A    infile should be a shapefile with
  fea
@@ -571,289 +571,416 @@
 ures
+,
  i
-n shapefile.%0A%0A    %60shp_id%60 is the feature id in the shapefile %60infile%60 (global). %60shp_id%60 and%0A    each id
+ds1 and ids2 should be%0A    a list of shape ids you wish to have distances calculated for.%0A    Distances between polygons that are both
  in 
-%60
 ids
-%60 (global) are cast to shapely features and then passed to a%0A    distance function contained in the shapely library. An array is returned%0A    with all these distan
+1 will not be calculated%0A    unless they are also in ids2 and vice versa. This allows the number of%0A    distance queries to be limited.%0A    q is a multiprocessing q object to which the returned data is placed so%0A    that it may be accessed from the calling pro
 ces
+s
 .%0A  
@@ -1005,22 +1005,30 @@
 na.open(
+args%5B'
 infile
+'%5D
 ) as sou
@@ -1166,22 +1166,30 @@
 rce%5Bint(
+args%5B'
 shp_id
+'%5D
 )%5D%5B'geom
@@ -1205,35 +1205,43 @@
        for j in 
+args%5B'
 ids
+'%5D
 :%0A            if
@@ -1254,22 +1254,30 @@
 ) %3C int(
+args%5B'
 shp_id
+'%5D
 ):%0A     
@@ -1569,33 +1569,18 @@
 t%0A%0A%0A
-if __name__ == %22__
+def 
 main
-__%22
+()
 :%0A  
@@ -1754,23 +1754,18 @@
     
-print %22
+# 
 Calculat
 ing 
@@ -1764,182 +1764,261 @@
 ulat
-ing distances%22%0A    # Map each id to a list of distances to all other ids%0A    pool = Pool()  # Spawn a worker process for each cpu%0A    data = pool.map(calculate_distances, ids
+e each the distance from each id to ids using a process pool%0A    print %22Calculating distances%22%0A    pool = Pool()%0A    data = pool.map(calculate_distances, %5B%7B%0A        'shp_id': shp_id,%0A        'infile': infile,%0A        'ids': ids,%0A    %7D for shp_id in ids%5D
 , ch
@@ -2125,16 +2125,31 @@
 e header
+ of output file
 %0A    pri
@@ -2527,8 +2527,47 @@
 close()%0A
+%0Aif __name__ == %22__main__%22:%0A    main()%0A

Revert "Add a couple global variables to simplify the map call. Update function docstrings"