# coding=utf-8

"""
@name: 按位置选择工具（支持反选空间关系、过滤对象（多个）、过滤缓冲值设置），返回选择要素Id集合或结果导出位置
@author: 陈蔡东
@datetime: 2024.04.16
"""

import os.path
import time

import arcpy
import sys
import json
import traceback
import codecs


# region 相关参数

# 1.输入数据
inputDatabasePath = ""
# 2.输入图层
inputLayer = ""
# 3.过滤条件
inputLayerWhere = ""

# 4.相关数据
relatedDatabasePath = ""
# 5.相关图层
relatedLayer = ""
# 6.相关图层过滤条件
relatedLayerWhere = ""

# 7.空间关系
spatialRelation = ""
# 8.搜索距离
searchDistance = -1.0
# 9.是否反选空间关系
invertSelect = True

# 10.过滤对象
filterObjects = []
# 11.过滤缓冲值
filterBuffer = -1

# 12.导出位置(.gdb,已存在)
exportLocation = ""
# 13.溯源字段
traceField = ""
# 14.导出字段[导出所有字段、导出指定字段]
exportFields = []

# endregion


# region 常量值

# 过滤数据路径 键名
FILTER_DATABASE_PATH_KEY = "databasePath"
# 过滤图层 键名
FILTER_LAYER_NAME_KEY = "layerName"
# 过滤条件 键名
FILTER_WHERE_CLAUSE_KEY = "whereClause"

# 外相接空间关系（相邻）
# EXTERNAL_TOUCH = "externalTouch"

# endregion


# 初始化
def init():
    # 参数初始化
    if len(sys.argv) < 2:
        raise Exception("未找到输入参数文件")

    argumentsFile = sys.argv[1]
    if not os.path.exists(argumentsFile):
        raise Exception("输入参数文件[" + argumentsFile + "]不存在")

    with codecs.open(argumentsFile, 'r', 'utf-8') as f:
        arguments = json.load(f)

    if len(arguments) < 1:
        raise Exception("输入参数文件[" + argumentsFile + ", 内容为空")

    if len(arguments) < 12:
        raise Exception("输入参数至少为13位")

    global inputDatabasePath, inputLayer, inputLayerWhere, relatedDatabasePath, relatedLayer, relatedLayerWhere
    global spatialRelation, searchDistance, invertSelect
    global filterObjects, filterBuffer, exportLocation, traceField
    inputDatabasePath = arguments[0]
    inputLayer = arguments[1]
    inputLayerWhere = "" if arguments[2] is None else arguments[2]
    outputMsg("输入数据: " + inputDatabasePath + "\\" + inputLayer)
    if bool(inputLayerWhere.strip()):
        outputMsg("过滤条件: " + inputLayerWhere)

    relatedDatabasePath = arguments[3]
    relatedLayer = arguments[4]
    relatedLayerWhere = "" if arguments[5] is None else arguments[5]
    outputMsg("相关数据: " + relatedDatabasePath + "\\" + relatedLayer)
    if bool(relatedLayerWhere.strip()):
        outputMsg("过滤条件: " + relatedLayerWhere)

    spatialRelation = arguments[6]
    outputMsg("空间关系: " + spatialRelation)
    searchDistance = -1 if arguments[7] is None else float(arguments[7])
    if searchDistance >= 0:
        outputMsg("搜索距离: " + str(searchDistance))
    invertSelect = False if arguments[8] is None else bool(arguments[8])
    outputMsg("反选空间关系:" + str(invertSelect))

    filterObjects = [] if arguments[9] is None else arguments[9]
    if len(filterObjects) > 0:
        outputMsg("过滤对象: ")
        outputMsg(filterObjects)
    filterBuffer = -1 if arguments[10] is None else float(arguments[10])
    if filterBuffer >= 0:
        outputMsg("过滤缓冲值:" + str(filterBuffer))

    exportLocation = "" if arguments[11] is None else arguments[11]
    if bool(exportLocation.strip()):
        outputMsg("是否导出结果：" + exportLocation)

    traceField = "" if arguments[12] is None else arguments[12]
    if bool(traceField.strip()):
        outputMsg("溯源字段：" + traceField)
        
    exportFields = [] if len(arguments) < 14 or (arguments[13] is None) else arguments[13]
    if len(exportFields) > 0:
        outputMsg("导出字段：")
        outputMsg(exportFields)


# 图层创建
def createLayer(databasePath, layerName, whereClause, newLayerName):
    arcpy.env.workspace = databasePath
    if bool(whereClause.strip()):
        arcpy.MakeFeatureLayer_management(layerName, newLayerName, whereClause)
    else:
        arcpy.MakeFeatureLayer_management(layerName, newLayerName)

    # 清除历史选择，防止对选择结果造成影响
    # arcpy.SelectLayerByAttribute_management(newLayerName, "CLEAR_SELECTION")
    return newLayerName


# 过滤对象
def filterData(inputLayerName):
    if len(filterObjects) < 1:
        return

    # 循环过滤 过滤对象
    for filterObject in filterObjects:
        databasePath = filterObject[FILTER_DATABASE_PATH_KEY]
        layerName = filterObject[FILTER_LAYER_NAME_KEY]

        whereClause = filterObject[FILTER_WHERE_CLAUSE_KEY]

        # 数据不存在
        if not os.path.exists(databasePath):
            raise Exception("过滤对象: " + databasePath + "不存在")

        # 图层为空
        if not (str(layerName).strip()):
            continue

        # 创建过滤图层
        filterLayerName = "related" + str(int(time.time() * 100)) + "_" + layerName
        createLayer(databasePath, layerName, whereClause, filterLayerName)

        # 执行相交选择，移除与过滤对象相交要素
        if filterBuffer >= 0:
            arcpy.SelectLayerByLocation_management(inputLayerName, "intersect", filterLayerName, filterBuffer, "REMOVE_FROM_SELECTION")
        else:
            arcpy.SelectLayerByLocation_management(inputLayerName, "intersect", filterLayerName, selection_type="REMOVE_FROM_SELECTION")
        # 移除图层
        arcpy.Delete_management(filterLayerName)


# 结果输出
def outputResult(layerName):

    desc = arcpy.Describe(inputDatabasePath + "\\" + inputLayer)
    idFieldName = desc.OIDFieldName

    # 导出结果
    if bool(exportLocation.strip()):

        # 保存溯源字段
        fieldMap = arcpy.FieldMap()
        fieldMap.addInputField(layerName, idFieldName)

        outputMappingField = fieldMap.outputField
        outputMappingField.name = traceField
        outputMappingField.aliasName = traceField
        outputMappingField.type = "long"
        fieldMap.outputField = outputMappingField

        fieldMap.mergeRule = "First"
        fieldMaps = arcpy.FieldMappings()
        fieldMaps.addFieldMap(fieldMap)
        
        for fieldTemp in arcpy.ListFields(layerName):
        	if fieldTemp.name == idFieldName or fieldTemp.type=="OID" or fieldTemp.type == "Geometry" or (len(exportFields) > 0 and not (fieldTemp.name .lower() in map(str.lower, exportFields))):
        		continue
        	fieldMapTemp = arcpy.FieldMap()
        	fieldMapTemp.addInputField(layerName, fieldTemp.name)
        	
        	outputMappingFieldTemp = fieldMapTemp.outputField
        	outputMappingFieldTemp.name = fieldTemp.name
        	outputMappingFieldTemp.aliasName = fieldTemp.name
        	outputMappingFieldTemp.type = fieldTemp.type
        	fieldMapTemp.outputField = outputMappingFieldTemp

        	fieldMapTemp.mergeRule = "First"
        	fieldMaps.addFieldMap(fieldMapTemp)
        
        arcpy.FeatureClassToFeatureClass_conversion(layerName, exportLocation, inputLayer,  field_mapping=fieldMaps)

    cursor = arcpy.SearchCursor(layerName)
    row = cursor.next()
    resultContent = ""
    count = 0
    while row:
        resultContent = resultContent + str(row.getValue(idFieldName)) + ";"
        count = count + 1
        row = cursor.next()

    # outputMsg("共查询到: " + str(count) + "条结果")
    outputMsg("Data" + resultContent)


# 消息写出
def outputMsg(msg):
    try:
        #if not isinstance(msg, str) and not isinstance(msg, unicode):
        #    arcpy.AddMessage(msg)
       #     return
        msg = msg.decode('utf-8')
        msg = repr(msg)
        arcpy.AddMessage(msg)
    except Exception, e:
        temp = e
        arcpy.AddMessage(msg)


# 输出错误信息
def outputErrorMsg(msg):
    try:
        #if not isinstance(msg, str) and not isinstance(msg, unicode):
         #   arcpy.AddError(msg)
        #    return
        msg = msg.decode('utf-8')
        msg = repr(msg)
        arcpy.AddError(msg)
    except Exception, e:
        temp = e
        arcpy.AddError(msg)


# 执行空间关系查询
def executeSpatialQuery(inputLayerName, relatedLayerName):
    if searchDistance >= 0:
        arcpy.SelectLayerByLocation_management(inputLayerName, spatialRelation, relatedLayerName, searchDistance)
    else:
        arcpy.SelectLayerByLocation_management(inputLayerName, spatialRelation, relatedLayerName)
    return


# 程序入口
def main():
    try:
        # 0.系统编码设置
        reload(sys)
        sys.setdefaultencoding('utf-8')

        # 1.参数初始化
        outputMsg("开始初始化参数...")
        init()

        outputMsg("正在打开数据...")
        # 2.创建图层
        inputLayerName = "inputLayer"
        createLayer(inputDatabasePath, inputLayer, inputLayerWhere, inputLayerName)
        relatedLayerName = "relatedLayer"
        createLayer(relatedDatabasePath, relatedLayer, relatedLayerWhere, relatedLayerName)

        outputMsg("开始执行空间选择...")
        arcpy.env.workspace = inputDatabasePath
        # 3.执行选择
        executeSpatialQuery(inputLayerName, relatedLayerName)

        # 4.反向选择
        if invertSelect:
            arcpy.SelectLayerByAttribute_management(inputLayerName, "SWITCH_SELECTION")

        # 5.过滤对象
        filterData(inputLayerName)

        outputMsg("开始输出结果...")
        # 6.结果保存
        outputResult(inputLayerName)

        # 7.移除图层，防止对下次查询造成影响
        arcpy.Delete_management(inputLayerName)
        arcpy.Delete_management(relatedLayerName)

    except Exception, e:
        outputErrorMsg("按位置选择执行异常")
        outputMsg(sys.argv)
        outputErrorMsg(e.message)
        outputErrorMsg(traceback.format_exc())
        raise e


# 开始执行
main()
