"""Plot heatmaps using ArcGIS.

This script must be run from the version of python installed with ArcGIS, e.g.
C:/Python27/ArcGIS10.4/python plot_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'])
    lyr_name = os.path.abspath(params['SYMBOLOGY LAYER'])
    input_tif_dir = os.path.abspath(params['TIF HEATMAP FOLDER'])
    output_jpg_dir = os.path.abspath(params['JPG 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
    heatmap_raster = os.path.join(input_tif_dir, base_name + '.tif')

    print('processing {}'.format(beach))
    mxd = arcpy.mapping.MapDocument(mxd_name)
    df = arcpy.mapping.ListDataFrames(mxd, 'Layers')[0]

    # Load erosion heatmap symbology
    symbology_layer = arcpy.mapping.Layer(lyr_name)

    # Add heatmap for current beach
    heatmap_layer = arcpy.mapping.Layer(heatmap_raster)
    arcpy.mapping.AddLayer(df, heatmap_layer, 'AUTO_ARRANGE')

    # Apply symbology to heatmap
    update_layer = arcpy.mapping.ListLayers(mxd, heatmap_layer, df)[0]
    arcpy.mapping.UpdateLayer(df, update_layer, symbology_layer, 'TRUE')
    arcpy.RefreshTOC()

    # Prepare output directories
    datestr = re.search('\d+$', base_name).group()
    current_date = '{}-{}-{}'.format(datestr[:4], datestr[4:6], datestr[6:])
    output_dir_current = os.path.join(output_jpg_dir, current_date)
    output_dir_latest = os.path.join(output_jpg_dir, 'latest')
    output_dirs = [output_dir_current]

    # Create output directories
    for dir_name in [output_dir_current, output_dir_latest]:
        if not os.path.isdir(dir_name):
            os.makedirs(dir_name)

    # Check dates of all output folders that exist so far
    all_dates = [
        d for d in os.listdir(output_jpg_dir)
        if re.match('\d{4}-\d{2}-\d{2}', d)
    ]

    # Export to 'latest' if current survey is most recent
    is_latest = current_date == sorted(all_dates)[-1]
    if is_latest:
        output_dirs.append(output_dir_latest)

    # Activate data driven pages
    ddp = mxd.dataDrivenPages

    for page_num in range(1, ddp.pageCount + 1):
        ddp.currentPageID = page_num

        # Export beaches that are covered by current survey
        if ddp.pageRow.Parent == beach.lower():
            page_name = ddp.pageRow.Beach
            for output_dir in output_dirs:
                jpg_name = os.path.join(output_dir, page_name + '.jpg')
                # Export to jpg
                arcpy.mapping.ExportToJPEG(
                    mxd, jpg_name, resolution=200, jpeg_quality=80)


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

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

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

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

    # Process all beaches for specific survey date
    C:/Python27/ArcGIS10.4/python plot_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()