"""Generate heatmaps using ArcGIS.

This script must be run from the version of python installed with ArcGIS, e.g.
C:/Python27/ArcGIS10.4/python generate_heatmaps.py

Note: arcpy does not play nicely with relative paths. This is because arcpy
changes the python working directory when an mxd is loaded, and during many
mapping operations. Relative paths may be used as inputs to this script, but
they will be converted to absolute paths internally.
"""

import os
import re
import sys
import argparse
from glob import glob

import arcpy

# Checkout Spatial Analyst toolbox to prevent licence errors
arcpy.CheckOutExtension('Spatial')

# Overwrite duplicate files
arcpy.env.overwriteOutput = True


def parse_yaml(yaml_name):
    """Parse yaml file manually (when 'yaml' module is unavailable)

    Args:
        yaml_name:  path to input yaml file

    Returns:
        dict of objects from yaml file

    """
    with open(yaml_name, 'r') as f:
        yaml = f.read()

    # Fit wrapped strings on one line
    yaml = re.sub('\n\s', '', yaml)

    # Parse yaml file
    params = {}
    for line in yaml.split('\n'):
        if line:
            key, val = line.split(':', 1)
            params[key] = val.strip()

    return params


def process(yaml_name):
    params = parse_yaml(yaml_name)
    beach = params['BEACH']
    #base_name = os.path.splitext(os.path.basename(params['INPUT LAS']))[0]
    base_name='%s_%s' % (beach.lower().replace(" ","_"), params['SURVEY DATE'])


    # Make all input and output paths absolute
    mxd_name = os.path.abspath(params['MAP DOCUMENT'])
    clip_poly = os.path.abspath(params['HEATMAP CROP POLY'])
    symb_lyr_name = os.path.abspath(params['SYMBOLOGY LAYER'])
    input_tif_dir = os.path.abspath(params['TIF DEM FOLDER'])
    output_tif_dir = os.path.abspath(params['TIF HEATMAP FOLDER'])

    # Check if previous survey date was provided
    previous_date = params['PREVIOUS SURVEY DATE']
    try:
        int(previous_date)
    except ValueError:
        raise ValueError('No previous survey date provided')

    # Set paths for raster files
    current_raster = os.path.join(input_tif_dir, base_name + '.tif')
    previous_base_name = re.sub('\d+', previous_date, base_name)
    previous_raster = os.path.join(input_tif_dir, previous_base_name + '.tif')
    heatmap_raster = os.path.join(output_tif_dir, base_name + '.tif')

    print('processing {}'.format(beach))

    # Save heatmap in memory
    heatmap_memory = 'in_memory/' + base_name

    # Subtract the previous raster from the current raster
    arcpy.Minus_3d(
        in_raster_or_constant1=current_raster,
        in_raster_or_constant2=previous_raster,
        out_raster=heatmap_memory)

    # Clip the heatmap
    arcpy.Clip_management(
        in_raster=heatmap_memory,
        rectangle='350000 6290000 360000 6300000',
        out_raster=heatmap_raster,
        in_template_dataset=clip_poly,
        nodata_value='-9.999000e+003',
        clipping_geometry='ClippingGeometry',
        maintain_clipping_extent='NO_MAINTAIN_EXTENT')


def main():
    example_text = """examples:

    # Process single survey at specific beach
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1-avoca.yaml

    # Process single survey at multiple beaches
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1-avoca.yaml survey-1-pearl.yaml

    # Process all surveys at specific beach
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py *avoca.yaml

    # Process all beaches for specific survey date
    C:/Python27/ArcGIS10.4/python generate_heatmaps.py survey-1*.yaml
    """

    # Set up command line arguments
    parser = argparse.ArgumentParser(
        epilog=example_text,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('input', help='path to yaml file(s)', nargs='*')

    # Print usage if no arguments are provided
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    # Parse arguments
    args = parser.parse_args()
    yaml_files = []
    [yaml_files.extend(glob(f)) for f in args.input]

    for yaml_file in yaml_files:
        process(yaml_file)


if __name__ == '__main__':
    main()