# -*- 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
"""
# TODO if 2 watersheds share more than one line segment (coastal thing).  Intersection >>> multipart line
    # then the exact spatial join fails >>> to both new and old watersheds
# TODO if line segment splits for attribution or not merged... (not polygon based). The intersection >>> one line segment,
    # but original lines have more than one segment.  Exact spatial join fials.
# TODO this is related to dealing with new subdivisions
    # NOTCHECKED at boundary
    # NOTCHECKED for new huc code (>>> lines that make up the new hucs are new)
    # NOTCHECKED polygons touch that did not used to touch (>>> new lines)
    # NOTCHECKED new polygons OR old polygons >>> multipart intersection
        # could explode all multipart for old and new and modify key (to keep unique based on length or position??)
         # then join like normal
    # NOTCHECKED user forgot to merge line segment
    # NOTCHECKED change in attribution >>> line break

# IMPORTS
import arcpy
import os
import sys
import re
from textwrap import wrap
from arcpy.arcobjects.arcobjects import FieldMappings
import wbd_f
import wbd_c
import wbd_params

# Constants
PYTHON_VERSION = sys.version_info.major
MEMORY_FPATH = wbd_f.get_memoryfpath(PYTHON_VERSION)
KEY = 'key'
REGEX_HUC = wbd_c.REGEX_HUC
TNMID = wbd_c.F_TNMID
HUMOD = wbd_c.F_HUMOD
LINESOURCE = wbd_c.F_LINESOURCE
HUDIGIT = wbd_c.F_HUDIGIT
CALC_SUFFIX = wbd_c.CALC_SUFFIX
FLAG_SUFFIX = wbd_c.FLAG_SUFFIX
TNMID_CALC = TNMID + CALC_SUFFIX
TNMID_FLAG = TNMID + FLAG_SUFFIX
HUMOD_CALC = HUMOD + CALC_SUFFIX
HUMOD_FLAG = HUMOD + FLAG_SUFFIX
HUDIGIT_CALC = HUDIGIT + CALC_SUFFIX
HUDIGIT_FLAG = HUDIGIT + FLAG_SUFFIX
CORRECT_FLAG = wbd_c.CORRECT_FLAG
LINESOURCE_CALC = LINESOURCE + CALC_SUFFIX
FLAG_LENGTH = wbd_c.FLAG_LENGTH
HUC = wbd_c.HUC



# WBDline_tnmid
class LinesCheck(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_tocheck (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 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_tocheck,
            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 = [TNMID, HUMOD, 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 self_intersection(self, fc, out_fname):
        return arcpy.Intersect_analysis(
            in_features=[fc,fc],
            # TODO overwrite output?
            out_feature_class=out_fname,
            join_attributes='ALL',
            output_type='LINE'
        )

    def add_linefields(self, oldline_fc, old_intersection, fields, out_fname):
        # field mappings object for old_intersections
        fms = arcpy.FieldMappings()
        fms.addTable(old_intersection)
        # add fields to be joined to field mapping
        for f in fields:
            fm = arcpy.FieldMap()
            # input field
            fm.addInputField(oldline_fc, f)
            # output field
            outfield = fm.outputField
            outfield.name = f
            outfield.aliasName = outfield.name
            outfield.aliasName = outfield.name
            fm.outputField = outfield
            fms.addFieldMap(fm)
        return arcpy.SpatialJoin_analysis(
            old_intersection,
            oldline_fc,
            out_fname,
            join_operation='JOIN_ONE_TO_MANY',
            field_mapping=fms,
            join_type='KEEP_ALL',
            match_option='ARE_IDENTICAL_TO'                                  # TODO this will only include polygons that are seperated by a single line segment
            )

    def add_keyfield(self, fc, huca, hucb, regex):
        arcpy.AddField_management(fc, KEY, 'TEXT')
        fields = [huca, hucb, KEY]
        workspace = os.path.dirname(wbd_f.get_fpath(fc))
        if '.gdb' in workspace and not workspace.endswith('.gdb'):
            workspace = os.path.dirname(workspace)
        with arcpy.da.Editor(workspace) as edit:
            with arcpy.da.UpdateCursor(fc, fields) as cursor:
                for row in cursor:
                    if (not row[0] or
                            not row[1] or
                            not re.match(regex, row[0]) or
                            not re.match(regex, row[1])):
                        row[2] = None
                    else:
                        row[2] = ','.join(sorted([row[0], row[1]]))
                    cursor.updateRow(row)

    def add_fields(self, fc, 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(fc)]
        if TNMID in fields_tocheck:
            if TNMID_CALC not in line_fieldnames:
                arcpy.AddField_management(fc, TNMID_CALC, 'TEXT', field_length=40)
            if TNMID_FLAG not in line_fieldnames:
                arcpy.AddField_management(fc, TNMID_FLAG, 'TEXT', FLAG_LENGTH)
                arcpy.CalculateField_management(fc, TNMID_FLAG, "'{}'".format(CORRECT_FLAG), 'PYTHON')
        if HUMOD in fields_tocheck:
            if HUMOD_CALC not in line_fieldnames:
                arcpy.AddField_management(fc, HUMOD_CALC, 'TEXT', field_length=30)
            if HUMOD_FLAG not in line_fieldnames:
                arcpy.AddField_management(fc, HUMOD_FLAG, 'TEXT', FLAG_LENGTH)
                arcpy.CalculateField_management(fc, HUMOD_FLAG, "'{}'".format(CORRECT_FLAG), 'PYTHON')
        if HUDIGIT in fields_tocheck:
            if HUDIGIT_CALC not in line_fieldnames:
                arcpy.AddField_management(fc, HUDIGIT_CALC, 'LONG')
            if HUDIGIT_FLAG not in line_fieldnames:
                arcpy.AddField_management(fc, HUDIGIT_FLAG, 'TEXT', FLAG_LENGTH)
                arcpy.CalculateField_management(fc, HUDIGIT_FLAG, "'{}'".format(CORRECT_FLAG), 'PYTHON')
        if LINESOURCE in fields_tocheck and LINESOURCE_CALC not in line_fieldnames:
                arcpy.AddField_management(fc, LINESOURCE_CALC, 'TEXT', field_length=75)

    def add_data(self, fields_tocheck, line_dict):
        """
        Checks/fills data.
        """
        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)
        # start edit session to update from dict
        with arcpy.da.Editor(workspace) as edit:
            # TNMID
            if TNMID in fields_tocheck:
                tnmid_dict = {k:v[TNMID] for (k,v) in line_dict.items()}
                fields = ['OID@', 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,
                            tnmid_dict,
                            newlineid,
                            self.fill_missing
                        )
                        cursor.updateRow(row)
            # HUMOD
            if HUMOD in fields_tocheck:
                humod_dict = {k:v[HUMOD] for (k,v) in line_dict.items()}
                fields = ['OID@', 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,
                            humod_dict,
                            newlineid,
                            self.fill_missing
                        )
                        cursor.updateRow(row)
            # HUDIGIT
            if HUDIGIT in fields_tocheck:
                hudigit_dict = {k:v[HUDIGIT] for (k,v) in line_dict.items()}
                fields = ['OID@', HUDIGIT, HUDIGIT_CALC, HUDIGIT_FLAG]
                with arcpy.da.UpdateCursor(self.linefc_updated, fields) as cursor:
                    for row in cursor:
                        newlineid = row[0]
                        row = wbd_f.updaterow_fromdict(                            # TODO fails because this field is a number
                            row,
                            1,
                            2,
                            3,
                            hudigit_dict,
                            newlineid,
                            self.fill_missing
                        )
                        cursor.updateRow(row)
            # LineSource
            if LINESOURCE in fields_tocheck:
                linesource_dict = {k:v[LINESOURCE] for (k,v) in line_dict.items()}
                fields = ['OID@', LINESOURCE, LINESOURCE_CALC]
                with arcpy.da.UpdateCursor(self.linefc_updated, fields) as cursor:
                    for row in cursor:
                        newlineid = row[0]
                        try:
                            new_linesource = linesource_dict[newlineid]
                        except KeyError:
                            new_linesource = None
                        row[2] = new_linesource
                        cursor.updateRow(row)

    def delete_badrows(self, fc, huca, hucb, key):
        with arcpy.EnvManager(overwriteOutput = True):
            huca = arcpy.AddFieldDelimiters(fc, huca)
            hucb = arcpy.AddFieldDelimiters(fc, hucb)
            key = arcpy.AddFieldDelimiters(fc, key)
            where = """{} = {} OR {} IS NULL OR {} IS NULL OR {} IS NULL""".format(
                huca,
                hucb,
                huca,
                hucb,
                key,
            )
            fl = arcpy.MakeFeatureLayer_management(fc, 'fl_temp', where_clause=where)
            arcpy.DeleteFeatures_management(fl)

    def get_hudigit(self, huca, hucb):
        for i, dig in enumerate(str(huca)):
            if str(hucb)[i] != dig:
                # Returned digits must be even rounded up
                if (i + 1) % 2 == 0:
                    return i + 1
                else:
                    return i + 2

    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.fill_missing = parameters[5].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.fill_missing = parameters[5]
        try:
            fields_tocheck = fields_tocheck.split(';')
        except AttributeError:
            pass
        self.add_fields(self.linefc_updated, fields_tocheck)
        if HUDIGIT in fields_tocheck:
            fields_tocheck.remove(HUDIGIT)
            calc_hudigit = True
        else:
            calc_hudigit = False
        huca = wbd_f.get_hucfield(self.polygonfc_old)
        hucb = huca + '_1'
        hudigit = huca.strip(HUC)
        regex_huc = REGEX_HUC.format(str(int(int(hudigit)/2-1)))
        # add old line fields to old intersections
        intersection_old = self.self_intersection(self.polygonfc_old, 'intersection_old')
        wbd_f.delete_extrafields(intersection_old, [huca, hucb])
        self.add_keyfield(intersection_old, huca, hucb, regex_huc)
        self.delete_badrows(intersection_old, huca, hucb, KEY)
        arcpy.DeleteIdentical_management(intersection_old,['Shape'])
        with arcpy.EnvManager(overwriteOutput = True):
            intersection_oldsj = self.add_linefields(self.linefc_old, intersection_old, fields_tocheck, 'intersection_old_sj')

        # add old line fields to new intersections
        intersection_new = self.self_intersection(self.polygonfc_updated, 'intersection_new')
        wbd_f.delete_extrafields(intersection_new, [huca, hucb])
        self.add_keyfield(intersection_new, huca, hucb, regex_huc)
        self.delete_badrows(intersection_new, huca, hucb, KEY)
        arcpy.DeleteIdentical_management(intersection_new,['Shape'])
        arcpy.JoinField_management(intersection_new, KEY, intersection_oldsj, KEY, fields_tocheck)

        # new line ids to new line intersections
        OID_FNAME = arcpy.Describe(self.linefc_updated).OIDFieldName
        LINE_OID = 'line_oid'
        with arcpy.EnvManager(overwriteOutput = True):
            intersection_newsj = wbd_f.spatialjoin_singlefield(
                intersection_new,
                self.linefc_updated,
                OID_FNAME,
                LINE_OID,                                        # TODO field is not joining correctly
                out_fc='intersection_new_sj',
                match_option='ARE_IDENTICAL_TO'                  # TODO this will only include polygons that are seperated by a single line segment
            )
        fields_tocheck.insert(0, 'JOIN_FID')               # TODO implied string name

        # create dictionary with line attributes
        line_dict = {}
        with arcpy.da.SearchCursor(intersection_newsj, fields_tocheck) as cursor:
            for row in cursor:
                line_dict[row[0]] = dict(zip(fields_tocheck[1:], row[1:]))
        if calc_hudigit:
            with arcpy.da.SearchCursor(intersection_newsj, ['JOIN_FID', huca, hucb]) as cursor:
                for row in cursor:
                    line_dict[row[0]][HUDIGIT] = self.get_hudigit(row[1], row[2])
            fields_tocheck.append(HUDIGIT)
        self.add_data(fields_tocheck, line_dict)

        arcpy.Delete_management([intersection_old, intersection_oldsj, intersection_new, intersection_newsj])

if __name__ == '__main__':
    """
    Execute as standalone script.
    """
    arcpy.env.workspace = r'C:\GIS_Project\WBD\AK\Work\hu19010204_redo\hu19010204_redo2.gdb'
    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 = [TNMID, HUMOD, LINESOURCE, HUDIGIT]
    fill_missing = True
    params = (linefc_old, linefc_updated, polygonfc_old, polygonfc_updated, to_transfer, fill_missing)
    linescheck = LinesCheck()
    linescheck.callfrom_pyt = False
    linescheck.execute(params, None)

