' Copyright 2010 ESRI
' 
' All rights reserved under the copyright laws of the United States
' and applicable international laws, treaties, and conventions.
' 
' You may freely redistribute and use this sample code, with or
' without modification, provided you include the original copyright
' notice and use restrictions.
' 
' See the use restrictions.
' 

Imports System
Imports System.Windows.Forms
Imports System.Runtime.InteropServices
Imports ESRI.ArcGIS.Framework
Imports ESRI.ArcGIS.NetworkAnalyst
Imports ESRI.ArcGIS.Carto
Imports ESRI.ArcGIS.esriSystem
Imports ESRI.ArcGIS.Geometry
Imports ESRI.ArcGIS.Geodatabase
Imports ESRI.ArcGIS.NetworkAnalystUI

Public Class SelectFeaturesTool
  Inherits ESRI.ArcGIS.Desktop.AddIns.Tool

  Dim m_networkAnalystExtension As INetworkAnalystExtension
  Public Sub New()
    m_networkAnalystExtension = My.ArcMap.Application.FindExtensionByName("Network Analyst")
  End Sub

  Protected Overrides Sub OnUpdate()
    Enabled = My.ArcMap.Application IsNot Nothing
  End Sub

  Private Function GetActiveAnalysisLayer() As INALayer
    If Not (m_networkAnalystExtension Is Nothing) Then
      Return m_networkAnalystExtension.NAWindow.ActiveAnalysis
    Else
      Return Nothing
    End If
  End Function

  Protected Overloads Overrides Function OnDeactivate() As Boolean
    Return True
  End Function

  ''' <summary>
  ''' Finds the closest feature to the point clicked that corresponds to a NATraversalResultElement in the active analysis layer's traversal result
  ''' </summary>
  Protected Overrides Sub OnMouseDown(ByVal e As MouseEventArgs)
    Try
      ' Get the current network analysis traversal result
      Dim naLayer As INALayer = GetActiveAnalysisLayer()
      If (naLayer Is Nothing) Then
        Throw New Exception("Null NALayer")
      End If

      Dim naTraversalResult As INATraversalResult = CType(naLayer.Context.Result, INATraversalResult)
      If (naTraversalResult Is Nothing) Then
        Throw New Exception("Active analysis layer does not have a result")
      End If
      Dim naResult = CType(naTraversalResult, INAResult)
      If (Not naResult.HasValidResult) Then
        Throw New Exception("Active analysis layer has an invalid result.")
      End If

      ' Get the current map
      Dim map As IMap = My.ArcMap.Document.FocusMap
      Dim activeView As IActiveView = My.ArcMap.Document.ActiveView

      ' Get the current point and buffer it into a polygon based on the search tolerance
      Dim point As IPoint = activeView.ScreenDisplay.DisplayTransformation.ToMapPoint(e.X, e.Y)

      Dim proximityOperator As IProximityOperator = point
      Dim topologicalOperator As ITopologicalOperator = point
      Dim polygon As IPolygon = topologicalOperator.Buffer(activeView.Extent.Width / 20)
      Dim closestDistance As Double = activeView.Extent.Width / 20
      Dim closestFeature As IFeature = Nothing

      ' Setup the spatial filter
      Dim spatialFilter As ISpatialFilter = New SpatialFilterClass()
      spatialFilter.SpatialRel = esriSpatialRelEnum.esriSpatialRelEnvelopeIntersects
      spatialFilter.Geometry = polygon.Envelope

      'Loop through the layers in the map that participate in the network dataset
      'Find the closest feature
      Dim layers As IEnumLayer = map.Layers(Nothing, True)
      Dim layer As ILayer = layers.Next()

      While Not layer Is Nothing

        If TypeOf layer Is IFeatureLayer Then

          Dim featureLayer As IFeatureLayer = layer

          If featureLayer.Visible Then
            Dim featureClass As IFeatureClass = featureLayer.FeatureClass

            ' loop though all the sources in the network, trying to see if it corresponds to a source in the traversal result.
            Dim sourceIndex As Integer
            For sourceIndex = 0 To naTraversalResult.SourceCount - 1
              ' Check to see if the feature class matches the traversal result source's feature class
              Dim currentTable As ITable = featureClass
              Dim sourceTable As ITable = naTraversalResult.Source(sourceIndex).Table
              If currentTable Is sourceTable Then
                ' Since we correspond to a source in the traversal result, find the closest feature to the point clicked.
                spatialFilter.GeometryField = featureClass.ShapeFieldName
                Dim featureCursor As IFeatureCursor = featureClass.Search(spatialFilter, False)
                Dim feature As IFeature = featureCursor.NextFeature()
                While Not feature Is Nothing
                  Dim distance As Double = proximityOperator.ReturnDistance(feature.Shape)

                  If distance < closestDistance Then
                    closestFeature = feature
                    closestDistance = distance
                  End If

                  feature = featureCursor.NextFeature()
                End While
              End If
            Next
          End If
        End If

        layer = layers.Next()
      End While

      If closestFeature Is Nothing Then
        MessageBox.Show("Could not find a feature")
      Else
        SelectTraversedFeatures(closestFeature)   'Since we found a feature, select the features traversed to get to it.
      End If
    Catch exception As Exception
      MessageBox.Show(exception.Message, "Error")
    End Try

  End Sub

  ''' <summary>
  ''' Selects all the features in the visible feature layers that correspond to NATraversalResultElements
  ''' that were traversed to get to the selected feature.  Do this by using INATraversalResultQuery to walk
  ''' the traversal result.
  ''' </summary> 
  Private Sub SelectTraversedFeatures(ByVal selectedFeature As IFeature)
    ' Get reference to the current map
    Dim map As IMap = My.ArcMap.Document.FocusMap

    ' Get the NATraversalResult
    Dim naLayer As INALayer = GetActiveAnalysisLayer()
    Dim naContext As INAContext = naLayer.Context
    Dim naTraversalResult As INATraversalResult = CType(naContext.Result, INATraversalResult)
    Dim naTraversalResultEdit As INATraversalResultEdit = CType(naTraversalResult, INATraversalResultEdit)
    Dim naTraversalResultQuery As INATraversalResultQuery = CType(naTraversalResult, INATraversalResultQuery)

    ' Cache the FeatureClass of each input source
    ' Also create a LongArray to hold the OIDs of features traversed
    ' This makes the assumption that the sources go from 1 to N
    Dim selectedSource As INATraversalResultSource = Nothing
    Dim sourceFeatureClasses As IFeatureClass() = New IFeatureClass(naTraversalResult.SourceCount + 1) {}
    Dim sourceOIDs As ILongArray() = New ILongArray(naTraversalResult.SourceCount + 1) {}

    Dim sourceIndex As Integer
    For sourceIndex = sourceFeatureClasses.GetLowerBound(0) To sourceFeatureClasses.GetUpperBound(0)
      Dim naTraversalResultSource As INATraversalResultSource = naTraversalResult.SourceByID(sourceIndex)

      If Not naTraversalResultSource Is Nothing Then
        Dim featureClass As IFeatureClass = naTraversalResultSource.Table
        If Not featureClass Is Nothing Then
          sourceFeatureClasses(sourceIndex) = featureClass

          ' If this source is the one of the selected feature,
          ' remember it for later when we find it's traversal elements
          If featureClass Is selectedFeature.Class Then
            selectedSource = naTraversalResultSource
          End If

          sourceOIDs(sourceIndex) = New LongArrayClass()
        End If
      End If
    Next

    ' Verify the selected feature's class is actually in the network
    If selectedSource Is Nothing Then
      Throw (New Exception("The selected feature is not part of the network."))
    End If

    ' Get the traversal edges corresponding to the selected feature
    ' There may be more than one returned if there are multiple facilities
    ' or if the edge was traversed in both directions.
    Dim queryFilter As IQueryFilter = New QueryFilterClass()
    queryFilter.WhereClause = "SourceID = " & selectedSource.ID & " and SourceOID = " & selectedFeature.OID

    ' Figure out if the selected feature is a point or line
    ' and get the correct traversal feature class
    Dim traversalFClass As IFeatureClass
    If selectedFeature.Shape.GeometryType = ESRI.ArcGIS.Geometry.esriGeometryType.esriGeometryPoint Then
      traversalFClass = naTraversalResultQuery.FeatureClass(esriNetworkElementType.esriNETJunction)
    ElseIf selectedFeature.Shape.GeometryType = ESRI.ArcGIS.Geometry.esriGeometryType.esriGeometryPolyline Then
      traversalFClass = naTraversalResultQuery.FeatureClass(esriNetworkElementType.esriNETEdge)
    Else
      Throw (New Exception("You must start from a point or polyline"))
    End If

    If traversalFClass.FeatureCount(Nothing) = 0 Then
      Throw (New Exception("The traversal result is empty.  Please click solve."))
    End If

    ' Search through the traversal class to find those corresponding to the selected feature
    Dim naTraversalResultElement As INATraversalResultElement
    Dim featureCursor As IFeatureCursor = traversalFClass.Search(queryFilter, True)
    Dim feature As IFeature = featureCursor.NextFeature()
    If feature Is Nothing Then
      Throw (New Exception("You must pick a feature that was traversed."))
    End If

    ' Loop through all the traversal result elements for that feature
    While Not feature Is Nothing
      naTraversalResultElement = CType(feature, INATraversalResultElement)
      sourceOIDs(naTraversalResultElement.SourceID).Add(naTraversalResultElement.SourceOID)

      'Get the connected traversal elements going backwards
      OutputConnected(naTraversalResultQuery, naTraversalResultElement, esriRelDirection.esriRelDirectionBackward, sourceOIDs)

      feature = featureCursor.NextFeature()
    End While

    ' Select the features in the Map
    SelectFeatures(map, sourceFeatureClasses, sourceOIDs)
  End Sub

  ''' <summary> 
  ''' This function selects the features corresponding to the arrays of feature classes and OIDs
  ''' </summary>
  Private Sub SelectFeatures(ByVal map As IMap, ByVal sourceFeatureClasses() As IFeatureClass, ByVal sourceOIDs() As ILongArray)
    Dim layers As IEnumLayer = map.Layers(Nothing, True)
    Dim layer As ILayer = layers.Next()
    Dim activeView As IActiveView = CType(map, IActiveView)

    While Not layer Is Nothing
      If (TypeOf layer Is IFeatureLayer) And (TypeOf layer Is IFeatureSelection) Then

        Dim featureLayer As IFeatureLayer = layer
        Dim featureSelection As IFeatureSelection = layer

        If featureLayer.Visible And featureLayer.Selectable Then
          Dim sourceIndex As Integer
          For sourceIndex = sourceFeatureClasses.GetLowerBound(0) To sourceFeatureClasses.GetUpperBound(0)
            If featureLayer.FeatureClass Is sourceFeatureClasses(sourceIndex) Then
              'Clear the current selection on this layer
              featureSelection.Clear()
              activeView.PartialRefresh(esriViewDrawPhase.esriViewGeoSelection, map, Nothing)

              ' Get the OIDS we traversed that we want to select
              Dim longArray As ILongArray = sourceOIDs(sourceIndex)

              If longArray.Count > 0 Then
                ' We have some OIDs, ad them to the selection set
                Dim selectionSet As ISelectionSet = featureSelection.SelectionSet

                ' Transfer from LongArray to an array of longs
                Dim numIDs As Integer = longArray.Count
                Dim oids() As Integer = New Integer(numIDs) {}
                Dim oidIndex As Integer
                For oidIndex = oids.GetLowerBound(0) To numIDs - 1
                  selectionSet.Add(longArray.Element(oidIndex))
                Next

              End If
            End If
          Next
        End If
      End If
      layer = layers.Next()
    End While

    ' Refresh the selection and Map
    Dim selectionEvents As ISelectionEvents = CType(map, ISelectionEvents)
    selectionEvents.SelectionChanged()

    activeView.PartialRefresh(esriViewDrawPhase.esriViewGeoSelection, map, Nothing)
  End Sub

  ''' <summary> 
  ''' This is a recursive function that simply finds the connected traversal elements.
  ''' It adds the OID of the feature to the appropriate array of ObjectIDs.
  ''' </summary>
  Private Sub OutputConnected(ByVal naTraversalResultQuery As INATraversalResultQuery, ByVal inputElement As INATraversalResultElement, ByVal searchDirection As esriRelDirection, ByVal sourceOIDs As ILongArray())
    Dim featureCursor As IFeatureCursor
    If inputElement.ElementType = esriNetworkElementType.esriNETJunction Then
      featureCursor = naTraversalResultQuery.SearchConnected(inputElement, esriNetworkElementType.esriNETEdge, searchDirection, False)
    Else
      featureCursor = naTraversalResultQuery.SearchConnected(inputElement, esriNetworkElementType.esriNETJunction, searchDirection, False)
    End If
    Dim connectedElement As INATraversalResultElement = CType(featureCursor.NextFeature, INATraversalResultElement)
    While Not (connectedElement Is Nothing)
      If Not (sourceOIDs(connectedElement.SourceID) Is Nothing) Then
        sourceOIDs(connectedElement.SourceID).Add(connectedElement.SourceOID)
      End If
      OutputConnected(naTraversalResultQuery, connectedElement, searchDirection, sourceOIDs)
      connectedElement = CType(featureCursor.NextFeature, INATraversalResultElement)
    End While
  End Sub
End Class