# -*- coding: utf-8 -*-
"""
 Tool Name:  CopyLineTNMID
 Description:  CopyLineTNMID transfers the TNMID values from the old WBDLine feature class
               to the updated WBDLine featureclass.
 Author:    Patrick Longley (plongley@usgs.gov)
 Created:   07/31/2020
 Language: Written in python3 (arcpro). Modified to also work in python2 (arcmap).
 History:
    metadata 20201005
    fields as constants 20201005
    check fields in update parameters 20201005
    general review 10/13/2020
"""

# IMPORTS
import arcpy
import os
import sys
import numpy as np
import collections
from scipy.spatial.distance import cdist
from scipy.optimize import linear_sum_assignment
import wbd_f
import wbd_c
import wbd_params

# Constants
PYTHON_VERSION = sys.version_info.major
MEMORY_FPATH = wbd_f.get_memoryfpath(PYTHON_VERSION)
MESSAGE = "Feature {} in {} has an incorrect {}."
TNMID_CALC = wbd_c.F_TNMID + wbd_c.CALC_SUFFIX
TNMID_FLAG = wbd_c.F_TNMID + wbd_c.FLAG_SUFFIX
HUMOD_CALC = wbd_c.F_HUMOD + wbd_c.CALC_SUFFIX
HUMOD_FLAG = wbd_c.F_HUMOD + wbd_c.FLAG_SUFFIX
LINESOURCE_CALC = wbd_c.F_LINESOURCE + wbd_c.CALC_SUFFIX
LINE_ID = 'line_id'
ORIG_FID = 'ORIG_FID'
WHITESPACE = r'^\s*$'

# WBDline_tnmid
class TransferFields(object):
    """
    arcpy tool that transfers the TNMID, HUMod and LineSource values from the old WBDLine feature class 
    to the updated WBDLine featureclass.

    Args:
        linefc_old (line feature class): Line feature class representing the original WBD lines.
            Must contain tnmid and hudigit fields (not case sensitive).
        linefc_new (line feature class): Line feature class representing the updated WBD lines.
        polygonfc_old (polygon feature class): Polygon feature class representing the original WBD polygons.
            Must contain a HUC code field.  Ex: HUC12 or huc8.
        polygonfc_updated (polygon feature class): Polygon feature class representing the updated WBD polygons.
            Must contain the same HUC code field as polygonfc_old.
        linefields_totransfer (list): subset of ['tnmid', 'linesource', 'humod']

    Outputs:
    returns: None
    output parameter: Modifies the polygonfc_updated in place
    """
    def __init__(self):
        """
        Initialize variables
        """
        self.label       = "2) Line: TNMID/HUMod/LineSource Check"
        self.description = "This tool migrates TNMID, HUMod, LineSource from the old WBD line feature class " + \
                           "to the updated WBD line feature class."
        self.callfrom_pyt = True
        self.category = 'Attribution'

    def getParameterInfo(self):
        """
        Define the parameters for use in arcmap/pro.
        """
        parameters = [
            wbd_params.linefc_old,
            wbd_params.linefc_updated,
            wbd_params.polygonfc_old,
            wbd_params.polygonfc_updated,
            wbd_params.linefields_totransfer,
            wbd_params.new_subdivisions,
            wbd_params.fill_missing,
        ]
        return parameters

    def updateMessages(self, params):
        """
        Modify the messages created by internal validation for each tool
        parameter.This method is called after internal validation.
        """
        required_fields = [wbd_c.F_TNMID, wbd_c.F_HUMOD, wbd_c.F_LINESOURCE]
        MESSAGE1 = "Feature class must contain the following fields: {}.".format(', '.join(required_fields))
        MESSAGE2 = "Both WBD feature classes must contain the same huc field."
        if params[0].altered:
            if not wbd_f.check_fieldsexist(params[0].valueAsText, required_fields):
                params[0].setErrorMessage(MESSAGE1)
        if params[1].altered:
            if not wbd_f.check_fieldsexist(params[1].valueAsText, required_fields):
                params[1].setErrorMessage(MESSAGE1)
        if params[2].altered and params[3].altered:
            if not wbd_f.check_hucfieldsmatch([params[2].valueAsText, params[3].valueAsText]):
                params[2].setErrorMessage(MESSAGE2)
                params[3].setErrorMessage(MESSAGE2)

    def subset_lines(self):
        """
        This function does the following:
            1) Converts the updatedpolygon_fc to lines
            2) Adds the object id of the longest line segment in the updated line fc to the newly created featureclass
        """
        UPDATEDLINE_OID = arcpy.Describe(self.linefc_updated).OIDFieldName
        KEY = 'key'
        SHAPE_LENGTH = 'shape_length'
        TARGET_FID = 'TARGET_FID'
        WORKSPACE = arcpy.env.workspace
        with arcpy.EnvManager(overwriteOutput = True):
            # line feature classes without new subdivisions ex: hudigit <= 12
            polygon_tolines = arcpy.FeatureToLine_management(self.polygonfc_updated,
                                                             os.path.join(MEMORY_FPATH, 'polygon_tolines'),
                                                             attributes='NO_ATTRIBUTES')
            unsplit_lines = arcpy.management.Dissolve(
                                                      polygon_tolines,
                                                      'unsplit_lines',
                                                      None,
                                                      None,
                                                      "SINGLE_PART",
                                                      "UNSPLIT_LINES")  # this line is necesary
        # spatial join unsplit_lines OBJECTID onto updated lines feature class as 'key'
        UNSPLITLINES_OID = arcpy.Describe(unsplit_lines).OIDFieldName
        sj = wbd_f.spatialjoin_singlefield(self.linefc_updated,
                                        unsplit_lines,
                                        UNSPLITLINES_OID,
                                        KEY,
                                        os.path.join(MEMORY_FPATH, 'sj'),
                                        join_type='KEEP_COMMON',
                                        match_option='SHARE_A_LINE_SEGMENT_WITH')
        # create data frame with unsplit_lines OID (key), updatedlines_fc OID (TARGET_FID) and shape_length
        df = wbd_f.create_df(sj, [KEY, TARGET_FID, SHAPE_LENGTH])
        # Max length for each dissolved line segment
        df_agg = df.groupby([KEY])[SHAPE_LENGTH].max()
        df = df.merge(df_agg, how='inner')
        # convert to dict where key = unsplit lines OID and value = updatedlines OID (for longest line segment)
        id_dict = dict(zip(df[KEY], df[TARGET_FID]))
        # add data from dictionary to unsplit_lines
        arcpy.AddField_management(unsplit_lines, LINE_ID, field_type='LONG')
        with arcpy.da.Editor(WORKSPACE) as edit:
            with arcpy.da.UpdateCursor(unsplit_lines, ['OID@', LINE_ID]) as cursor:
                for row in cursor:
                    try:
                        row[1] = id_dict[row[0]]
                        cursor.updateRow(row)
                    except KeyError:
                        pass
        return arcpy.MakeFeatureLayer_management(unsplit_lines, 'unsplit_lines_fl')

    def create_fls(self):
        """
        Creates polygon feature layers using the inputed polygon feature classes.
        """
        with arcpy.EnvManager(overwriteOutput=True):
            self.polygonfl_old = arcpy.MakeFeatureLayer_management(self.polygonfc_old, 'polygonfl_old')
            self.polygonfl_updated = arcpy.MakeFeatureLayer_management(self.polygonfc_updated, 'polygonfl_updated')
            self.linefl_old = arcpy.MakeFeatureLayer_management(self.linefc_old, 'oldl')
            if self.new_subdivisions:
                self.linefl_updated = self.subset_lines()
            else:
                self.linefl_updated = arcpy.MakeFeatureLayer_management(self.linefc_updated, 'newl')

    def list_hucs(self):
        """
        Creates a list of HUCs that are present in both the old feature class
        and the updated feature class.
        """
        old_arr = arcpy.da.TableToNumPyArray(self.polygonfc_old, [self.hucfield])
        old_arr = [x[0] for x in old_arr]
        new_arr = arcpy.da.TableToNumPyArray(self.polygonfc_updated, [self.hucfield])
        new_arr = [x[0] for x in new_arr]
        self.huc_arr = np.intersect1d(old_arr, new_arr).astype(str)

    def check_nsegments(self, huc):
        """
        Creates line feature layers using the inputed line feature classes.
        Select by location to only use lines within the polygon specified by the "huc" input parameter.
        Checks if the number of line segments in both featureclasses matches
        """
        # select polygon feature layers by attribute
        where_oldpoly = """ {} = '{}' """.format(
            arcpy.AddFieldDelimiters(self.polygonfl_old, self.hucfield),
            huc
            )
        where_newpoly = """ {} = '{}' """.format(
            arcpy.AddFieldDelimiters(self.polygonfl_updated, self.hucfield),
            huc
            )
        arcpy.SelectLayerByAttribute_management(self.polygonfl_old, where_clause=where_oldpoly)
        arcpy.SelectLayerByAttribute_management(self.polygonfl_updated, where_clause=where_newpoly)
        # spatial selection of line feature layers
        arcpy.SelectLayerByLocation_management(self.linefl_old,
                                               select_features=self.polygonfl_old,
                                               overlap_type='SHARE_A_LINE_SEGMENT_WITH')
        arcpy.SelectLayerByLocation_management(self.linefl_updated,
                                               select_features=self.polygonfl_updated,
                                               overlap_type='SHARE_A_LINE_SEGMENT_WITH')
        n_oldlines = int(arcpy.GetCount_management(self.linefl_old)[0])
        n_newlines = int(arcpy.GetCount_management(self.linefl_updated)[0])
        if n_newlines == n_oldlines and n_newlines > 0:
            return True
        else:
            return False

    def create_points(self):
        """
        Converts line feature classes to point feature classes (saved in memory).
        Adds x/y coordinates to the point feature classes.
        """
        # create points feature layer
        with arcpy.EnvManager(overwriteOutput=True):
            self.pointsfl_old = arcpy.FeatureToPoint_management(self.linefl_old, os.path.join(MEMORY_FPATH,"old_points"))
            self.pointsfl_updated = arcpy.FeatureToPoint_management(self.linefl_updated, os.path.join(MEMORY_FPATH,"new_points"))
        arcpy.AddGeometryAttributes_management(self.pointsfl_old, 'POINT_X_Y_Z_M')
        arcpy.AddGeometryAttributes_management(self.pointsfl_updated, 'POINT_X_Y_Z_M')

    def optimize_pairs(self):
        """
        Converts point feature classes to numpy arrays.
        Pairs the new points and old points together by minimizing the sum of the euclidean distance.
        """
        arr_old = arcpy.da.TableToNumPyArray(self.pointsfl_old, ['POINT_X', 'POINT_Y', wbd_c.F_TNMID, wbd_c.F_HUMOD, wbd_c.F_LINESOURCE])
        if self.new_subdivisions:
            arr_new = arcpy.da.TableToNumPyArray(self.pointsfl_updated, ['POINT_X', 'POINT_Y', LINE_ID])
        else:
            arr_new = arcpy.da.TableToNumPyArray(self.pointsfl_updated, ['POINT_X', 'POINT_Y', ORIG_FID])
        arr_old = np.array([list(x) for x in arr_old])
        arr_new = np.array([list(x) for x in arr_new])
        n_matches = arr_old.shape[0]
        distance_matrix = cdist(arr_old[:,:2].astype(float), arr_new[:,:2])
        _, assignment = linear_sum_assignment(distance_matrix)
        for p in range(n_matches):
            self.data_totransfer.append((int(arr_new[assignment[p], 2]), str(arr_old[p, 2]), str(arr_old[p, 3]), str(arr_old[p, 4])))

    def remove_conflicts(self):
        """
        This makes sure that each OBJECTID corresponds to 1 and only 1 TNMID.
        Creates a dictionary of OBJECTID, TNMID pairs where OBJECTID is the key and TNMID is the value.
        Creates a dictionary of OBJECTID, HUMod pairs where OBJECTID is the key and HUMod is the value.
        Creates a dictionary of OBJECTID, LineSource pairs where OBJECTID is the key and LineSource is the value.
        """
        # convert to list of tuples
        self.data_totransfer = list(set(self.data_totransfer))
        # remove conflicts 
        ids = [x[0] for x in self.data_totransfer]
        tnmids = [x[1] for x in self.data_totransfer]
        duplicate_ids = [x for x, n in collections.Counter(ids).items() if n > 1]
        duplicate_tnmids = [x for x, n in collections.Counter(tnmids).items() if n > 1]
        self.data_totransfer = [x for x in self.data_totransfer if x[0] not in duplicate_ids and x[1] not in duplicate_tnmids]
        # create dictionaries (don't include whitespace or 'None' (Null from arc))
        self.tnmid_dict = dict([(x[0], x[1]) for x in self.data_totransfer if x[1].strip() and x[1] != 'None'])
        self.humod_dict = dict([(x[0], x[2]) for x in self.data_totransfer if x[2].strip() and x[2] != 'None'])
        self.linesource_dict = dict([(x[0], x[3]) for x in self.data_totransfer if x[3].strip() and x[3] != 'None'])

    def add_fields(self, fields_tocheck):
        """
            Add calc and flag fields to the dataset if they do not already exist.
        """
        line_fieldnames = [x.name for x in arcpy.ListFields(self.linefc_updated)]
        if wbd_c.F_TNMID in fields_tocheck:
            if TNMID_CALC not in line_fieldnames:
                arcpy.AddField_management(self.linefc_updated, TNMID_CALC, 'TEXT', field_length=40)
            if TNMID_FLAG not in line_fieldnames:
                arcpy.AddField_management(self.linefc_updated, TNMID_FLAG, 'TEXT', wbd_c.FLAG_LENGTH)
                arcpy.CalculateField_management(self.linefc_updated, TNMID_FLAG, "'{}'".format(wbd_c.CORRECT_FLAG), 'PYTHON')
        if wbd_c.F_HUMOD in fields_tocheck:
            if HUMOD_CALC not in line_fieldnames:
                arcpy.AddField_management(self.linefc_updated, HUMOD_CALC, 'TEXT', field_length=30)
            if HUMOD_FLAG not in line_fieldnames:
                arcpy.AddField_management(self.linefc_updated, HUMOD_FLAG, 'TEXT', wbd_c.FLAG_LENGTH)
                arcpy.CalculateField_management(self.linefc_updated, HUMOD_FLAG, "'{}'".format(wbd_c.CORRECT_FLAG), 'PYTHON')
        if wbd_c.F_LINESOURCE in fields_tocheck and LINESOURCE_CALC not in line_fieldnames:
                arcpy.AddField_management(self.linefc_updated, LINESOURCE_CALC, 'TEXT', field_length=75)

    def add_data(self, fields_tocheck):
        """
        Checks/fills data.
        """
        self.add_fields(fields_tocheck)
        workspace = os.path.dirname(wbd_f.get_fpath(self.linefc_updated))
        if '.gdb' in workspace and not workspace.endswith('.gdb'):
            workspace = os.path.dirname(workspace)
        OID = arcpy.Describe(self.linefc_updated).OIDFieldName
        # start edit session to update from dict
        with arcpy.da.Editor(workspace) as edit:
            # TNMID
            if wbd_c.F_TNMID in fields_tocheck:
                fields = [OID, wbd_c.F_TNMID, TNMID_CALC, TNMID_FLAG]
                with arcpy.da.UpdateCursor(self.linefc_updated, fields) as cursor:
                    for row in cursor:
                        newlineid = row[0]
                        row = wbd_f.updaterow_fromdict(
                            row,
                            1,
                            2,
                            3,
                            self.tnmid_dict,
                            newlineid,
                            self.fill_missing
                        )
                        cursor.updateRow(row)
            # HUMOD
            if wbd_c.F_HUMOD in fields_tocheck:
                fields = [OID, wbd_c.F_HUMOD, HUMOD_CALC, HUMOD_FLAG]
                with arcpy.da.UpdateCursor(self.linefc_updated, fields) as cursor:
                    for row in cursor:
                        newlineid = row[0]
                        row = wbd_f.updaterow_fromdict(
                            row,
                            1,
                            2,
                            3,
                            self.humod_dict,
                            newlineid,
                            self.fill_missing
                        )
                        cursor.updateRow(row)
            # LineSource
            if wbd_c.F_LINESOURCE in fields_tocheck:
                fields = [OID, wbd_c.F_LINESOURCE, LINESOURCE_CALC]
                with arcpy.da.UpdateCursor(self.linefc_updated, fields) as cursor:
                    for row in cursor:
                        newlineid = row[0]
                        try:
                            new_linesource = self.linesource_dict[newlineid]
                        except KeyError:
                            new_linesource = None
                        row[2] = new_linesource
                        cursor.updateRow(row)

    def execute(self, parameters, messages):
        """
        Loops through HUCS and executes the above functions to add data to the updated wbd line feature class.
        """
        # parameters
        if self.callfrom_pyt:
            self.linefc_old = parameters[0].valueAsText
            self.linefc_updated = parameters[1].valueAsText
            self.polygonfc_old = parameters[2].valueAsText
            self.polygonfc_updated = parameters[3].valueAsText
            fields_tocheck = parameters[4].valueAsText
            self.new_subdivisions = parameters[5].value
            self.fill_missing = parameters[6].value
        else:
            self.linefc_old = params[0]
            self.linefc_updated = params[1]
            self.polygonfc_old = params[2]
            self.polygonfc_updated = params[3]
            fields_tocheck = parameters[4]
            self.new_subdivisions = parameters[5]
            self.fill_missing = parameters[6]
        try:
            fields_tocheck = fields_tocheck.split(';')
        except AttributeError:
            pass
        self.hucfield = wbd_f.get_hucfield(self.polygonfc_old)
        self.tokeep = [wbd_c.F_TNMID, wbd_c.F_HUMOD, wbd_c.F_LINESOURCE]
        # prep functions
        self.create_fls()
        self.list_hucs()
        # loop through HUCs
        self.data_totransfer = []
        for huc in self.huc_arr:
            nlines_match = self.check_nsegments(huc)
            if nlines_match:
                self.create_points()
                self.optimize_pairs()
        self.remove_conflicts()
        self.add_data(fields_tocheck)

if __name__ == '__main__':
    """
    Execute as standalone script.
    """
    linefc_old = r'C:\GIS_Project\WBD\AK\Work\hu19010204_redo\19010204_prep\19010204_data.gdb\initial_data\WBDLine'
    linefc_updated = r'C:\GIS_Project\WBD\AK\Work\hu19010204_redo\hu19010204_redo2.gdb\Layers\NewLines'
    polygonfc_old = r'C:\GIS_Project\WBD\AK\Work\hu19010204_redo\19010204_prep\19010204_data.gdb\initial_data\HU14_19010204'
    polygonfc_updated = r'C:\GIS_Project\WBD\AK\Work\hu19010204_redo\hu19010204_redo2.gdb\Layers\HU14'
    to_transfer = [wbd_c.F_TNMID, wbd_c.F_HUMOD, wbd_c.F_LINESOURCE]
    new_subdivisions = False
    fill_missing = True
    params = (linefc_old, linefc_updated, polygonfc_old, polygonfc_updated, to_transfer, new_subdivisions, fill_missing)
    transferfields = TransferFields()
    transferfields.callfrom_pyt = False
    transferfields.execute(params, None)

