"""
 Tool Name:  PacakageData
 Description: Pacakges preexisting data for WBD work.
 Author: Patrick Longley (plongley@usgs.gov)
 Created: 08/11/2020
 Language: Written in python3 (arcpro). Modified to also work in python2 (arcmap).
 History:
 TODO: units for contours?
"""

# IMPORTS
import os
import arcpy
import sys
import wbd_f
import wbd_params

#PYTHON 2/3 issues
python_version = sys.version_info.major
template_dict = wbd_f.template_dict

class PackageData(object):
    """
    Pacakges preexisting data for WBD work.

    Args:
        single_huc (str): HUC code.
        wbd_fd (feature dataset): Reprojected WBD feature dataset
        nhd_fd (feature dataset): Reprojected NHD feature dataset

        out_folder (str):  Output folder where data will be saved.
        contour_spacing (int): Contour spacing
        include_artificial (boolean): Include ditches and canals in NHD flowline featureclass?


    Outputs:
        returns: output feature dataset

    """
    def __init__(self):
        """
        Initialize variables
        """
        self.label       = "Package data"
        self.description = "This tool packages WBD, NHD, and 3DEP data for wbd work."
        self.callfrom_pyt = True
        self.category = 'Download data'

    def getParameterInfo(self):
        """
        Define the parameters for use in arcmap/pro.
        """
        params = [wbd_params.single_huc,
                  wbd_params.wbd_fd,
                  wbd_params.nhd_fd,
                  wbd_params.out_folder,
                  wbd_params.include_artificial]
        return params

    def updateMessages(self, params):
        """
        Modify the messages created by internal validation for each tool
        parameter.This method is called after internal validation.
        """
        if params[0].value and params[0].altered and params[0].value!='':
            huclist = params[0].valueAsText.replace(" ", "").split(';')
            c1 = wbd_f.check_hucre(huclist)
            if not c1:
                message = 'Invalid HUC code(s).'
                params[0].setErrorMessage(message)
        if params[1].altered and params[2].altered and params[3].altered:
            datasets = [params[1].valueAsText, params[2].valueAsText, params[3].valueAsText]
            if not wbd_f.check_srmatches(datasets):
                message = "Spatial references do not match for input datasets."
                params[1].setErrorMessage(message)
                params[2].setErrorMessage(message)
                params[3].setErrorMessage(message)

    def create_folders(self):
        """
        Creates folders where data will be saved.
        """
        self.prep_folder = os.path.join(self.out_folder, self.huc + '_prep')
        if not os.path.exists(self.prep_folder):
            os.mkdir(self.prep_folder)
        self.prep_gdb_name = self.huc + '_data.gdb'
        self.prep_gdb = os.path.join(self.prep_folder, self.prep_gdb_name)
        self.prep_fd = os.path.join(self.prep_gdb, 'initial_data')
        if not os.path.exists(self.prep_gdb):
            arcpy.CreateFileGDB_management(self.prep_folder, self.prep_gdb_name)
            fd = arcpy.CreateFeatureDataset_management(self.prep_gdb,
                                                       'initial_data',
                                                       self.nhd_fd)
            arcpy.CreateFeatureclass_management(fd, 'batchpoints_' + self.huc, "POINT")
            arcpy.CreateFeatureclass_management(fd, 'walls_' + self.huc, "POLYLINE")

    def copy_wbd(self):
        """
        Copies WBD data into the correct folder.
        """
        hudigit_list = ['16', '14', '12', '10', '8', '6', '4', '2']
        hudigit_list = hudigit_list[:hudigit_list.index(self.hudigit) + 1]
        # suffix
        if self.wbd_fd.endswith('_rpj'):
            suffix = '_1'
        else:
            suffix = ''
        for hu in hudigit_list:
            # copy polygons
            original_fc = os.path.join(self.wbd_fd, ''.join(['WBDHU', hu + suffix]))
            new_fc_fname = ''.join(['HU', hu, '_', self.huc])
            field = 'huc' + hu
            where = """ {} like {} """.format(arcpy.AddFieldDelimiters(original_fc, field),
                                        "".join(["'", self.huc, "%'"]))
            with arcpy.EnvManager(overwriteOutput=True):
                fl_polygon = arcpy.MakeFeatureLayer_management(original_fc, 'fl_temp', where)
            if int(arcpy.GetCount_management(fl_polygon)[0]) > 0:
                # TODO domains: 
                new_polygons = arcpy.CreateFeatureclass_management(                                  # create empty feature class using template
                    self.prep_fd,
                    new_fc_fname,
                    template=template_dict[hu]
                )
                arcpy.Append_management(fl_polygon, new_polygons, "NO_TEST")                           # append from feature layer
                # copy label points
                point_fname = ''.join(['HU', hu, '_LabelPt_', self.huc])
                point_fc = os.path.join(self.prep_fd, point_fname)
                arcpy.FeatureToPoint_management(new_polygons, point_fc, "INSIDE")
        self.wbd_fl = fl_polygon
        # copy WBD line
        oldline_fc = os.path.join(self.wbd_fd, ''.join(['WBDLine', suffix]))
        newline_fname = ''.join(['WBDLine', '_', self.huc])
        with arcpy.EnvManager(overwriteOutput=True):
            fl_line = arcpy.MakeFeatureLayer_management(oldline_fc, 'fl_line')
        arcpy.SelectLayerByLocation_management(fl_line, 'within', self.wbd_fl)
        new_lines = arcpy.CreateFeatureclass_management(                                        # create empty feature class using template
            self.prep_fd,
            newline_fname,
            template=template_dict['line']
        )
        arcpy.Append_management(fl_line, new_lines, "NO_TEST")                                 # append from feature layer
        # wbdline_fc = arcpy.CopyFeatures_management('fl_line', os.path.join(self.prep_fd, 'WBDLine'))
        # return wbdline_fc[0]
        return new_lines

    def copy_nhd(self):
        """
        Copies NHD data into the correct folder.
        """
        if self.nhd_fd.endswith('_rpj'):
            self.flowline_fc = os.path.join(self.nhd_fd, 'NHDFlowline_1')
            self.area_fc = os.path.join(self.nhd_fd, 'NHDArea_1')
            self.waterbody_fc = os.path.join(self.nhd_fd, 'NHDWaterbody_1')
        else:
            self.flowline_fc = os.path.join(self.nhd_fd, 'NHDFlowline')
            self.area_fc = os.path.join(self.nhd_fd, 'NHDArea')
            self.waterbody_fc = os.path.join(self.nhd_fd, 'NHDWaterbody')
        with arcpy.EnvManager(overwriteOutput=True):
            arcpy.MakeFeatureLayer_management(self.flowline_fc, 'flowline_fl')
            arcpy.MakeFeatureLayer_management(self.flowline_fc, 'coastline_fl')
            arcpy.MakeFeatureLayer_management(self.area_fc, 'area_fl')
            arcpy.MakeFeatureLayer_management(self.waterbody_fc, 'waterbody_fl')
        fl_list = ['flowline_fl', 'coastline_fl', 'area_fl', 'waterbody_fl']
        # select nhd features that intersect WBD feature layer
        for fl in fl_list:
            arcpy.SelectLayerByLocation_management(fl, select_features=self.wbd_fl)
        # select by attribute for NHD flowline
        if self.include_artificial:
            where_flow = """{} <> 566""".format(arcpy.AddFieldDelimiters('flowline_fl', 'FType'))
        else:
            where_flow = """{} not in (566, 336, 334, 428, 420)""".format(arcpy.AddFieldDelimiters('flowline_fl', 'FType'))
        arcpy.SelectLayerByAttribute_management('flowline_fl', "SUBSET_SELECTION", where_clause=where_flow)
        # select by attribute for NHD coastline
        where_coast = """{} = 566""".format(arcpy.AddFieldDelimiters('coastline_fl', 'FType'))
        arcpy.SelectLayerByAttribute_management('coastline_fl', "SUBSET_SELECTION", where_clause=where_coast)
        # copy feature layers to new feature classes
        flowline = arcpy.CopyFeatures_management('flowline_fl', os.path.join(self.prep_fd, 'NHDFlowline_' + self.huc))
        area = arcpy.CopyFeatures_management('area_fl', os.path.join(self.prep_fd, 'NHDArea_' + self.huc))
        waterbody = arcpy.CopyFeatures_management('waterbody_fl', os.path.join(self.prep_fd, 'NHDWaterbody_' + self.huc))
        if int(arcpy.GetCount_management('coastline_fl')[0]) > 0:
            arcpy.CopyFeatures_management('coastline_fl', os.path.join(self.prep_fd, 'NHDCoastline_' + self.huc))
        return flowline[0], area[0], waterbody[0]

    def execute(self, params, messages):
        """
        Executes above functions to package data for WBD work.
        """
        # parameters
        if self.callfrom_pyt:
            self.huc = params[0].valueAsText
            self.wbd_fd = params[1].valueAsText
            self.nhd_fd = params[2].valueAsText
            self.out_folder = params[3].valueAsText
            self.contour_dist = params[4].valueAsText
            self.include_artificial = params[5].value
        else:
            self.huc = params[0]
            self.wbd_fd = params[1]
            self.nhd_fd = params[2]
            self.out_folder = params[3]
            self.contour_dist = params[4]
            self.include_artificial = params[5]
        self.hudigit = str(len(self.huc))
        # copy data
        self.create_folders()
        wbdline = self.copy_wbd()
        flowline, area, waterbody = self.copy_nhd()
        # add outputs to map
        if self.callfrom_pyt:
            arcpy.SetParameter(8, self.prep_fd)
        return self.prep_folder


if __name__ == '__main__':
    """
    Execute as standalone script.
    """
    hucs = '180101010101'
    wbd_fd = r'C:\Users\plongley\Downloads\WBD_18_HU2_GDB.gdb\WBD_rpj'
    nhd_fd = r'C:\Users\plongley\Downloads\NHD_H_1801_HU4_GDB.gdb\Hydrography_rpj'
    out_folder = r'C:\Users\plongley\Downloads'
    include_artificial = False
    params = (hucs, wbd_fd, nhd_fd, out_folder, include_artificial)
    package_data = PackageData()
    package_data.callfrom_pyt = False
    package_data.execute(params, None)
