"""
 Tool Name:  ToHUC_HUType
 Description: Checks/fills the ToHUC and HUType fields.
 Author: Patrick Longley (plongley@usgs.gov)
 Created: 10/08/2020
 Language: Written in python3 (arcpro). Modified to also work in python2 (arcmap).
 History:
    validation 10/13/2020
    constants 10/13/2020
    metadata 10/13/2020
    general review 10/13/2020
    no objectid issues 10/20/2020
    tested with new schema/sychronization 12/09/2020
"""

# IMPORTS
import sys
import os
import re
import collections
import arcpy
import numpy as np
import wbd_f
import wbd_params
import wbd_c

# Constants
PYTHON_VERSION = sys.version_info.major
MEMORY_FPATH = wbd_f.get_memoryfpath(PYTHON_VERSION)
TOHUC_CALC = wbd_c.F_TOHUC + wbd_c.CALC_SUFFIX
TOHUC_FLAG = wbd_c.F_TOHUC + wbd_c.FLAG_SUFFIX
HUTYPE_CALC = wbd_c.F_HUTYPE + wbd_c.CALC_SUFFIX
HUTYPE_FLAG = wbd_c.F_HUTYPE + wbd_c.FLAG_SUFFIX
WHITESPACE = r'^\s*$'

# TOOLS
class TohucHutype(object):
    """
    arcpy tool that fills/checks the ToHUC and HUType field for a WBD polygon feature class

    Args:
        polygonfc_updated (polygon feature class): WBD polygon feature class.
            Must contain a HUC field.
        nhdline_fc  (line feature class): NHD flowline feature class.
        nhdarea_fc  (polygon feature class): NHD area polygon feature class

    Outputs:
    returns: None
    output parameter: Modifies the polygonfc_updated in place.
    """

    def __init__(self):
        """
        Initialize variables
        """
        self.label       = "4) Polygon: ToHUC/HUType Check"
        self.description = "This tool checks/fills the ToHUC and HUType fields."
        self.callfrom_pyt = True
        self.category = 'Attribution'

    def getParameterInfo(self):
        """
        Define the parameters for use in arcmap/pro.
        """
        params = [
            wbd_params.polygonfc_updated,
            wbd_params.nhdline_fc,
            wbd_params.nhdarea_fc,
            wbd_params.outlets,
            wbd_params.frontal_outlets,
            wbd_params.fill_missing,
        ]
        return params

    def updateMessages(self, params):
        """
        Modify the messages created by internal validation for each tool
        parameter. This method is called after internal validation.
        """
        MESSAGE = "Feature class must contain the following fields: {}."
        required_fields = [wbd_c.F_TOHUC, wbd_c.F_HUTYPE]
        if params[0].altered:
            polygon_fc = params[0].valueAsText
            if not params[3].altered:
                params[3].value = os.path.basename(params[0].valueAsText) + '_outlets'
            if not params[4].altered:
                params[4].value = os.path.basename(params[0].valueAsText) + '_frontaloutlets'
            if (not wbd_f.check_fieldsexist(polygon_fc, required_fields) or
                    not wbd_f.get_hucfield(polygon_fc)):
                params[0].setErrorMessage(MESSAGE.format(', '.join(required_fields + ['huc'])))

    def create_fls(self):
        """
        Creates feature layers.
        """
        with arcpy.EnvManager(overwriteOutput=True):
            self.polygon_fl = arcpy.MakeFeatureLayer_management(self.polygonfc_updated, 'polygon_fl')
            self.nhdarea_fl = arcpy.MakeFeatureLayer_management(self.nhdarea_fc, 'nhdarea_fl')
            where_coast = """{} <> 566""".format(arcpy.AddFieldDelimiters(self.nhdline_fc, wbd_c.F_FTYPE))
            try:
                self.flowline_fl = arcpy.MakeFeatureLayer_management(self.nhdline_fc, 'flowline_fl', where_coast)
            except:
                self.flowline_fl = arcpy.MakeFeatureLayer_management(self.nhdline_fc, 'flowline_fl')

    def split_flowlines(self):
        """
        Split flowlines at polygon boundaries.
        """
        with arcpy.EnvManager(overwriteOutput=True):
            self.flowlinessplit = arcpy.Intersect_analysis([self.polygon_fl, self.flowline_fl], 'flowlines_split')
            wbd_f.delete_extrafields(self.flowlinessplit, [self.hucfield])

    def create_points(self):
        """
        Get end points and startpoints of split flowlines.
        """
        with arcpy.EnvManager(overwriteOutput=True):
            self.end_points = arcpy.FeatureVerticesToPoints_management(self.flowlinessplit,
                                                                       os.path.join(MEMORY_FPATH, 'end_points'),
                                                                       "END")
            self.start_points = arcpy.FeatureVerticesToPoints_management(self.flowlinessplit,
                                                                         'start_points',
                                                                         "START")
            arcpy.AlterField_management(self.start_points, field=self.hucfield, new_field_name=wbd_c.F_TOHUC)
            self.startpoints_fl = arcpy.MakeFeatureLayer_management(self.start_points)

    def find_outlets(self):
        """
        Find standard outlets (end point intersects a start point and is on a polygon boundary)
        """
        arcpy.SelectLayerByLocation_management(self.startpoints_fl,
                                               select_features=self.polygon_fl,
                                               overlap_type='BOUNDARY_TOUCHES')
        with arcpy.EnvManager(overwriteOutput=True):
            self.outlets = wbd_f.spatialjoin_singlefield(
                self.end_points,
                self.startpoints_fl,
                wbd_c.F_TOHUC,
                wbd_c.F_TOHUC,
                self.outlets_fpath ,
                join_type='KEEP_Common'
            )

        arcpy.DeleteIdentical_management(self.outlets, ['Shape', self.hucfield, wbd_c.F_TOHUC])
        # delete features where tohuc = huc
        where = """ {} = {} """.format(arcpy.AddFieldDelimiters(self.outlets,self.hucfield),
                                       arcpy.AddFieldDelimiters(self.outlets,wbd_c.F_TOHUC))
        outlets_fl = arcpy.MakeFeatureLayer_management(self.outlets, 'outlets_fl')
        arcpy.SelectLayerByAttribute_management(outlets_fl, where_clause=where)
        arcpy.DeleteFeatures_management(outlets_fl)

    def frontal_outlets(self):
        """
        Find coast outlets (endpoints that are not start points, do intersect nhd area, not on polygon edge).
        """
        self.frontaloutlets_fl =  arcpy.MakeFeatureLayer_management(self.end_points,
                                                                    'frontaloutlets_fl')
        arcpy.SelectLayerByLocation_management(self.frontaloutlets_fl,
                                               select_features=self.nhdarea_fl)
        arcpy.SelectLayerByLocation_management(self.frontaloutlets_fl,
                                               select_features=self.start_points,
                                               selection_type='SUBSET_SELECTION',
                                               invert_spatial_relationship='INVERT')
        arcpy.SelectLayerByLocation_management(self.frontaloutlets_fl,
                                               select_features=self.polygon_fl,
                                               overlap_type='BOUNDARY_TOUCHES',
                                               selection_type='SUBSET_SELECTION',
                                               invert_spatial_relationship='INVERT')
        with arcpy.EnvManager(overwriteOutput=True):
            arcpy.CopyFeatures_management(self.frontaloutlets_fl, self.frontaloutlets_fpath)
        self.frontaloutlets_stats = arcpy.analysis.Statistics(self.frontaloutlets_fl,
                                                              os.path.join(MEMORY_FPATH, 'frontal_outlets_stats'),
                                                              "{} COUNT".format(self.hucfield),
                                                              [self.hucfield])[0]

    def create_dictionaries(self):
        """
        Create dictionaries:
            1) self.outletscount_dict: key = huc, value = number of outlets
            2) self.frontaloutletscount_dict: key = huc, value = number of coastal outlets
            3: self.tohuc_dict: key =huc, value = tohuc
        """
        # convert to numpy array
        outlets_arr = arcpy.da.TableToNumPyArray(self.outlets, [self.hucfield, wbd_c.F_TOHUC])
        outlets_arr = [x for x in outlets_arr if x[0] != x[1]]
        frontaloutlets_arr = arcpy.da.TableToNumPyArray(self.frontaloutlets_stats, [self.hucfield, 'FREQUENCY'])
        # make dictionaries
        self.outletscount_dict = dict([(x, n) for x, n in collections.Counter([x[0] for x in outlets_arr]).items()])
        self.frontaloutletscount_dict = dict(frontaloutlets_arr)
        # tohuc dict
        hucs = [x[0] for x in np.unique(outlets_arr)]
        duplicates = [x for x, n in collections.Counter(hucs).items() if n > 1]
        lst = list(set([(x[0], x[1]) if x[0] not in duplicates else(x[0], 'multiple') for x in outlets_arr]))
        self.tohuc_dict = dict(lst)

    def add_fields(self, fc):
        fields = [f.name for f in arcpy.ListFields(fc)]
        if TOHUC_CALC not in fields:
            arcpy.AddField_management(fc, TOHUC_CALC, 'TEXT', field_length=16)
        if HUTYPE_CALC not in fields:
            arcpy.AddField_management(fc, HUTYPE_CALC, 'TEXT', field_length=1)
        if TOHUC_FLAG not in fields:
            arcpy.AddField_management(fc, TOHUC_FLAG, 'TEXT', field_length=wbd_c.FLAG_LENGTH)
            arcpy.CalculateField_management(fc, TOHUC_FLAG, "'{}'".format(wbd_c.CORRECT_FLAG), 'PYTHON')
        if HUTYPE_FLAG not in fields:
            arcpy.AddField_management(fc, HUTYPE_FLAG, 'TEXT', field_length=wbd_c.FLAG_LENGTH)
            arcpy.CalculateField_management(fc, HUTYPE_FLAG, "'{}'".format(wbd_c.CORRECT_FLAG), 'PYTHON')

    def gethutype_fromdicts(self, outlets_dict, frontaloutlets_dict, key):
        # normal outlets >>> S or M
        if key in outlets_dict and key not in frontaloutlets_dict:
            n_outlets = outlets_dict[key]
            new_hutype = ('S' if n_outlets == 1 else 'M')
        # frontal outlets >>> S or F
        elif key not in outlets_dict and key in frontaloutlets_dict:
            n_outlets = frontaloutlets_dict[key]
            new_hutype = ('S' if n_outlets == 1 else 'F')
        # both outlet types >>> problem
        elif key in outlets_dict and key in frontaloutlets_dict:
            new_hutype = None
        # no outlets >>> C
        else:
            new_hutype = 'C'
        return new_hutype

    def checkfill_data(self):
        """
        Checks/fills HUType and ToHUC.
        """
        self.add_fields(self.polygonfc_updated)
        workspace = os.path.dirname(wbd_f.get_fpath(self.polygonfc_updated))
        if '.gdb' in workspace and not workspace.endswith('.gdb'):
            workspace = os.path.dirname(workspace)
        with arcpy.da.Editor(workspace) as _:
            # TOHUC
            fields = [self.hucfield, wbd_c.F_TOHUC, TOHUC_CALC, TOHUC_FLAG]
            with arcpy.da.UpdateCursor(self.polygonfc_updated, fields) as cursor:
                for row in cursor:
                    row = wbd_f.updaterow_fromdict(
                        row,
                        1,
                        2,
                        3,
                        self.tohuc_dict,
                        row[0],
                        self.fill_missing
                    )
                    cursor.updateRow(row)
            # HUTYPE
            fields = [self.hucfield, wbd_c.F_HUTYPE, HUTYPE_CALC, HUTYPE_FLAG]
            with arcpy.da.UpdateCursor(self.polygonfc_updated, fields) as cursor:
                for row in cursor:
                    if row[1] not in ('M', 'S', 'F', 'C'):  # don't check other types
                        row[2] = None
                        new_flag = wbd_c.NOTCHECKED_FLAG
                    else:
                        old_hutype = row[1]
                        new_hutype = self.gethutype_fromdicts(
                            self.outletscount_dict,
                            self.frontaloutletscount_dict,
                            row[0]
                        )
                        row[2] = new_hutype
                        if not new_hutype:   # both standard and frontal outlets  >>> wrong
                            new_flag = wbd_c.WRONG_FLAG
                        elif not old_hutype or re.match(WHITESPACE, old_hutype):
                            if self.fill_missing:
                                new_flag = wbd_c.FILLED_FLAG
                                row[1] = new_hutype
                                if row[3] in wbd_f.AUTO_FLAGS:
                                    row[3] = new_flag # overwrite nodata flag with filled flag
                            else:
                                new_flag = wbd_c.NODATA_FLAG
                        elif old_hutype != new_hutype:
                            new_flag = wbd_c.WRONG_FLAG
                        else:
                            new_flag = wbd_c.CORRECT_FLAG
                    row[3] = wbd_f.update_flag(row[3], new_flag)
                    cursor.updateRow(row)

    def execute(self, params, messages):
        """
        Executes above functions.  Reminder: this code only works with HUC12s, HUC14s, and HUC16s.
        """
        # parameters
        if self.callfrom_pyt:
            self.polygonfc_updated = params[0].valueAsText
            self.nhdline_fc = params[1].valueAsText
            self.nhdarea_fc = params[2].valueAsText
            self.outlets_fpath = params[3].valueAsText
            self.frontaloutlets_fpath = params[4].valueAsText
            self.fill_missing = params[5]
        else:
            self.polygonfc_updated = params[0]
            self.nhdline_fc = params[1]
            self.nhdarea_fc = params[2]
            self.outlets_fpath = params[3]
            self.frontaloutlets_fpath = params[4]
            self.fill_missing = params[5]


        self.field_names = [x.name for x in arcpy.ListFields(self.polygonfc_updated)]
        _, hudigits = wbd_f.sort_polygons([self.polygonfc_updated])
        hudigit = hudigits[0]
        self.hucfield = wbd_c.HUC + hudigit
        # execute above functions
        self.create_fls()
        self.split_flowlines()
        self.create_points()
        self.find_outlets()
        self.frontal_outlets()
        self.create_dictionaries()
        self.checkfill_data()
        # delete uneeded data
        arcpy.Delete_management(self.start_points)
        arcpy.Delete_management(self.end_points)
        arcpy.Delete_management(self.frontaloutlets_stats)
        arcpy.Delete_management(self.flowlinessplit)

if __name__ == '__main__':
    """
    Execute as standalone script.
    """
    arcpy.env.workspace = r'C:\GIS_Project\WBD\AK\Work\hu19080306\hu19080306.gdb'
    polygonfc_updated = r'C:\GIS_Project\WBD\AK\Work\hu19080306\19080306_prep\19080306_data.gdb\initial_data\hu12_19080306'
    nhd_flowline = r'C:\GIS_Project\WBD\AK\Work\hu19080306\19080306_prep\19080306_data.gdb\initial_data\NHDFlowline_19080306'
    nhd_area = r'C:\GIS_Project\WBD\AK\Work\hu19080306\19080306_prep\19080306_data.gdb\initial_data\NHDArea_19080306'
    outlets_path = r'C:\GIS_Project\WBD\AK\Work\hu19080306\hu19080306.gdb\outlets_old'
    frontaloutlets_path = r'C:\GIS_Project\WBD\AK\Work\hu19080306\hu19080306.gdb\frontaloutlets_old'
    fill_missing = False
    params = (polygonfc_updated, nhd_flowline, nhd_area, outlets_path, frontaloutlets_path, fill_missing)
    tohuc = TohucHutype()
    tohuc.callfrom_pyt = False
    tohuc.execute(params, None)

