"""Calculates and compares area within any polygonal dataset.

If no comparison field is given, total area of dataset will be printed to
console. Specifying a comparison field will calculate and print, for each value
in the field, the total area of all elements with that value in the field.

Can be run as a geoprocessing tool within ArcCatalog using calcarea.tbx,
or run from command line with the following arguments:
path to dataset,
(optional) name of field containing categories for area comparison

"""

import arcpy
import os

def message(*args):
    """Output messages to user in a standardized format.
    """
    args = " ".join([str(arg) for arg in args])
    arcpy.AddMessage("\t * %s" % args)

def error(*args):
    """Output error messages in a standardized format.
    """
    args = " ".join([str(arg) for arg in args])
    arcpy.AddError("\t * %s" % args)

class WrongShapeType(Exception):
    """Operation requires different shape type than that found in the dataset.
    """

    def __init__(self, path, shape_type):
        self.path = path
        self.shape_type = shape_type

def get_unique_values(dataset_path, comparison_field):
    """Creates a set of all unique values in user's input comparison field.
    """
    rows = arcpy.SearchCursor(dataset_path)
    value_set = set([])
    for row in rows:
        value = row.getValue(comparison_field)
        value_set.add(value)
    # Make sure reference count can reach zero
    del row, rows
    return value_set

def get_unique_areas(dataset_path, comparison_field, area_field="Calc_Area"):
    """Get the total area of all polygons associated with each field type.

    Returns a dictionary mapping each unique value found in the comparison field
    to the total area of all polygons associated with that value.
    """
    rows = arcpy.SearchCursor(dataset_path)
    dic = {}
    for row in rows:
        value = row.getValue(comparison_field)
        area = row.getValue(area_field)
        if value in dic:
            dic[value] += area
        else:
            dic[value] = area
    # Make sure reference count can reach zero
    del row, rows
    return dic

def total_area(dataset_path, area_field="Calc_Area"):
    """Computes the total area of all polygons in the file.
    """
    rows = arcpy.SearchCursor(dataset_path)
    total_area = 0

    for row in rows:
        area = row.getValue(area_field)
        total_area += area
    # Make sure reference count can reach zero
    del row, rows
    return total_area

def calc_area(workspace_path, dataset_filename):
    """Adds a calculated field containing the area of each polygon in the
       dataset.
    """
    dataset_description = arcpy.Describe(dataset_filename)
    # Make sure that the dataset contains Polygons,
    # and if not handle that gracefully rather than letting arc choke on it
    dataset_shape = dataset_description.shapeType
    if dataset_shape == "Polygon":
        # delete Calc_Area field if it already exists to avoid using stale data
        arcpy.DeleteField_management(dataset_filename, "Calc_Area")
        # precision & scale set to maximum to account for giant polygons
        # warning: REALLY big polygons might still cause problems
        #   (e.g., rounding error)
        arcpy.AddField_management (dataset_filename, "Calc_Area", "DOUBLE", 16, 8)
        message("Calculating area...")
        arcpy.CalculateField_management(dataset_filename, "Calc_Area", "!Shape.Area!", "PYTHON")
    else:
        raise WrongShapeType(workspace_path, dataset_shape)

def main():
    """Main entry point for when this module is run standalone (as opposed to as a library).
    """

    # get user input & set environment
    dataset_path = arcpy.GetParameterAsText(0)
    comparison_field = arcpy.GetParameterAsText(1)
    workspace_path = (os.path.split(dataset_path))[0]
    arcpy.env.workspace = workspace_path
    dataset_filename = (os.path.split(dataset_path))[1]

    # Add calculated area field, fail gracefully
    try:
        calc_area(workspace_path, dataset_filename)
    except WrongShapeType, exception:
        error("Shapefile has wrong type: cannot calculate area for %s features in %s" % (exception.shape_type, exception.path))
        return

    # If user provided comparison field, break down areas by field values
    if comparison_field:
        unique_areas = get_unique_areas(dataset_path, comparison_field)
        # Make a pretty string out of the dictionary
        buf = ""
        keys = unique_areas.keys()
        keys.sort()
        for key in keys:
            pretty_key = "<No Value>" if (str(key)).strip() == "" else key
            buf += "\t\t%s: %f\n" % (pretty_key, unique_areas[key])

        message("Total areas by %(field)s\n%(buf)s" % {'field':comparison_field, 'buf':buf})

    # Otherwise, just calculate overall area for dataset
    else:
        area = total_area(dataset_path)
        message("Total area: %s" % (area))

if __name__ == "__main__":
    main()