' Copyright 2008 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 use restrictions at <your ArcGIS install location>/developerkit/userestrictions.txt.
' 




'The New Upstream trace task performs an upstream trace, the same as the Find Upstream
'trace task that is provided by ArcMap. To use the New Upstream trace task, you must
'place at least one flag on the current network. This trace task also uses the
'barriers, disabled layers, selected features, and weight filters that you specify
'using the Utility Network Analyst toolbar to perform the trace.

Option Explicit

Implements ITraceTask
Implements ITraceTaskResults

'declare a module-level variable to hold a reference the Utility Network Analysis extension
Private m_pUtilityNetworkAnalysisExt As esriEditorExt.IUtilityNetworkAnalysisExt

'declare module-level variables to hold the result edges and junctions of the last trace
Private m_pResultJunctions As IEnumNetEID
Private m_pResultEdges As IEnumNetEID
Private m_lJuncEID As Long

Private Property Get ITraceTask_EnableSolve() As Boolean
'determines whether to enable or disable the Solve button on the toolbar.
'this trace task can be enabled if there is at least one flag on the network.

  Dim pNetworkAnalysisExtFlags As esriEditorExt.INetworkAnalysisExtFlags
  Dim pNetworkAnalysisExt As esriEditorExt.INetworkAnalysisExt
  Dim lngNetworkCount As Long
  
  'by default, the Solve button will be disabled
  ITraceTask_EnableSolve = False
  
  'QI for the INetworkAnalysisExt interface using IUtilityNetworkAnalysisExt
  Set pNetworkAnalysisExt = m_pUtilityNetworkAnalysisExt
  'determine the number of networks currently loaded
  lngNetworkCount = pNetworkAnalysisExt.NetworkCount
  'if there are no networks loaded, then the Solve button is disabled
  If lngNetworkCount = 0 Then Exit Property

  'QI for the INetworkAnalysisExtFlags interface using IUtilityNetworkAnalysisExt
  Set pNetworkAnalysisExtFlags = m_pUtilityNetworkAnalysisExt
  'if there is one flag on the network then enable the task
  With pNetworkAnalysisExtFlags
    If .JunctionFlagCount = 1 Or .EdgeFlagCount = 1 Then
      'there is at least one flag on the current network
      ITraceTask_EnableSolve = True
    Else
      'there are no flags on the current network
      Exit Property
    End If
  End With
  
End Property

Private Property Get ITraceTask_Name() As String
  ITraceTask_Name = "Next Upstream Device"
End Property

Private Sub ITraceTask_OnCreate(ByVal utilityNetworkAnalysis As esriEditorExt.IUtilityNetworkAnalysisExt)
'assign the input argument to the module-level variable.
'this maintains a reference to the Utility Network Analysis extension.
  Set m_pUtilityNetworkAnalysisExt = utilityNetworkAnalysis
End Sub

Private Sub ITraceTask_OnTraceExecution()
'perform the upstream trace task, using the flags, barriers, disabled layers,
'and weight filters that were specified using the toolbar.
  On Error GoTo ErrorHandler
  
  Dim pNetworkAnalysisExtFlags As esriEditorExt.INetworkAnalysisExtFlags
  Dim pTraceFlowSolver As esriNetworkAnalysis.ITraceFlowSolver
  Dim pResultEdges As esriGeoDatabase.IEnumNetEID
  Dim pResultJunctions As esriGeoDatabase.IEnumNetEID
  Dim pNetworkAnalysisExtResults As esriEditorExt.INetworkAnalysisExtResults
    
  'prepare the network solver
  m_lJuncEID = -1
  Set pTraceFlowSolver = UTIL_CoreTraceSetup
  If pTraceFlowSolver Is Nothing Then Exit Sub
      
  'Run the trace the first time to get all upstream devices
  pTraceFlowSolver.FindFlowElements esriFMUpstream, esriFEJunctions, _
   pResultJunctions, pResultEdges
   
  'Reset the trace for the second pass
  Set pTraceFlowSolver = UTIL_CoreTraceSetup
  
  'Create a barrier set based on the disabled layers that are set and the junction returned
  'from the first trace.
  UTIL_SetBarriers pResultJunctions, pTraceFlowSolver
  
  pTraceFlowSolver.FindFlowEndElements esriFMUpstream, esriFEJunctions, _
   pResultJunctions, pResultEdges
  
  'copy the results to the module level
  Set m_pResultJunctions = pResultJunctions
  Set m_pResultEdges = pResultEdges
  
  'update the extension with the results
  'QI for the INetworkAnalysisExtResults interface using IUtilityNetworkAnalysisExt
  Set pNetworkAnalysisExtResults = m_pUtilityNetworkAnalysisExt
  'first, clear the previous results
  pNetworkAnalysisExtResults.ClearResults
  
  'Create the selection
  pNetworkAnalysisExtResults.CreateSelection pResultJunctions, pResultEdges
  
  ResultNotification
  
  Exit Sub
  
ErrorHandler:
  'an unexpected error occured somewhere in the function
  'notify the user that an error occurred and exit the sub
  MsgBox "ITraceTask_OnTraceExecution - " & Err.Description
End Sub

Private Sub ResultNotification()
On Error GoTo ErrHand:
  If m_pResultJunctions.Count = 1 Then
    Dim pApp As IApplication, pDoc As IMxDocument, pMap As IMap, lLoop As Long
    Dim pFeatSel As IFeatureSelection
    Set pApp = New AppRef
    Set pDoc = pApp.Document
    Set pMap = pDoc.FocusMap
    For lLoop = 0 To pMap.LayerCount - 1
      If TypeOf pMap.Layer(lLoop) Is IFeatureSelection Then
        Set pFeatSel = pMap.Layer(lLoop)
        If pFeatSel.SelectionSet.Count = 1 Then
          MsgBox "1 " & pMap.Layer(lLoop).Name & " feature selected."
          Exit For
        End If
      End If
    Next lLoop
  Else
    MsgBox m_pResultJunctions.Count & " feature(s) selected."
  End If

  Exit Sub
ErrHand:
  MsgBox "ResultNotification - " & Err.Description
End Sub

'the properties of the ITraceTaskResults interface return the
'results of the last trace performed by this trace task.
Private Property Get ITraceTaskResults_ResultEdges() As IEnumNetEID
  Set ITraceTaskResults_ResultEdges = m_pResultEdges
End Property

Private Property Get ITraceTaskResults_ResultJunctions() As IEnumNetEID
  Set ITraceTaskResults_ResultJunctions = m_pResultJunctions
End Property

Private Sub UTIL_SetBarriers(pResultJunctions As IEnumNetEID, pNetSolver As INetSolver)
On Error GoTo ErrHand:
  Dim lLoop As Integer, pFeatureLayer As IFeatureLayer
  Dim pNetworkAnalysisExt As esriEditorExt.INetworkAnalysisExt
'  Dim pNetElementBarrier As INetElementBarriers2
  Dim pNetElementBarrier As INetElementBarriers
  Dim pNetworkAnalysisExtBarriers As INetworkAnalysisExtBarriers
  Dim lNetEID As Long, lLoop2 As Long, lFeatClass As Long
  Dim lNetEIDs() As Long, lCount As Long
  Dim pNetElems As INetElements
  Dim lUserIDClass As Long, lUserID As Long, lSub As Long
  
  'Initialize the barrier set
  Set pNetworkAnalysisExt = m_pUtilityNetworkAnalysisExt
  Set pNetworkAnalysisExtBarriers = pNetworkAnalysisExt
  Set pNetElementBarrier = New NetElementBarriers
  pNetElementBarrier.ElementType = esriETJunction
  Set pNetElementBarrier.Network = pNetworkAnalysisExt.CurrentNetwork.Network
  Set pNetElems = pNetworkAnalysisExt.CurrentNetwork.Network
  
  'Make a string of the disabled layer feature class IDs
  For lLoop = 0 To pNetworkAnalysisExt.FeatureLayerCount - 1
    Set pFeatureLayer = pNetworkAnalysisExt.FeatureLayer(lLoop)
    If pNetworkAnalysisExtBarriers.GetDisabledLayer(pFeatureLayer) Then
      lFeatClass = pFeatureLayer.FeatureClass.FeatureClassID
    
      'Add barriers to the isolation features that are not currently part of the FindPath
      lCount = 0
      pResultJunctions.Reset
      For lLoop2 = 1 To pResultJunctions.Count
        lNetEID = pResultJunctions.Next
        If m_lJuncEID <> lNetEID Then
          pNetElems.QueryIDs lNetEID, esriETJunction, lUserIDClass, lUserID, lSub
          If lUserIDClass = lFeatClass Then
            ReDim Preserve lNetEIDs(lCount)
            lNetEIDs(lCount) = lUserID
            lCount = lCount + 1
          End If
        End If
      Next lLoop2
        
      If lCount > 0 Then
        pNetElementBarrier.SetBarriers lFeatClass, lCount, lNetEIDs(0)
        Set pNetSolver.ElementBarriers(esriETJunction) = pNetElementBarrier
      End If
    End If
  Next lLoop
    
  Exit Sub

ErrHand:
  MsgBox "UTIL_SetBarriers - " & Err.Description
End Sub

Public Function UTIL_CoreTraceSetup() As esriNetworkAnalysis.ITraceFlowSolver
'this function prepares the network for tracing.
  
  On Error GoTo ErrorHandler
  
  Dim pNetworkAnalysisExt As esriEditorExt.INetworkAnalysisExt
  Dim pNetSolver As esriNetworkAnalysis.INetSolver
  Dim pNetworkAnalysisExtFlags As esriEditorExt.INetworkAnalysisExtFlags
  Dim pFlagDisplay As esriNetworkAnalysis.IFlagDisplay
  Dim pEdgeFlagDisplay As esriNetworkAnalysis.IEdgeFlagDisplay
  Dim pEdgeFlags() As esriNetworkAnalysis.IEdgeFlag
  Dim pJunctionFlags() As esriNetworkAnalysis.IJunctionFlag
  Dim pNetFlag As esriNetworkAnalysis.INetFlag
  Dim pEdgeFlag As esriNetworkAnalysis.IEdgeFlag
  Dim pTraceFlowSolver As esriNetworkAnalysis.ITraceFlowSolver
  Dim pTraceTasks As esriEditorExt.ITraceTasks
  Dim lngEdgeFlagCount As Long
  Dim lngJunctionFlagCount As Long
  Dim lLoop As Long
  Dim pNetElems As INetElements
  
  Set pNetworkAnalysisExt = m_pUtilityNetworkAnalysisExt
  Set pNetElems = pNetworkAnalysisExt.CurrentNetwork.Network
  
  'initialize the trace flow solver
  'co-create a new TraceFlowSolver object
  Set pNetSolver = New esriNetworkAnalysis.TraceFlowSolver
  'set the source network for the solver
  Set pNetSolver.SourceNetwork = pNetElems
    
  'assign the flags to the network solver
  'first, get the edge flags
  'QI for the ITraceFlowSolver interface using INetSolver interface
  Set pTraceFlowSolver = pNetSolver
  'QI for the INetworkAnalysisExtFlags interface using IUtilityNetworkAnalysisExt
  Set pNetworkAnalysisExtFlags = m_pUtilityNetworkAnalysisExt
  'determine the number of edge flags on the current network
  lngEdgeFlagCount = pNetworkAnalysisExtFlags.EdgeFlagCount
  'only execute this next bit if there are any edge flags
  If Not lngEdgeFlagCount = 0 Then
    'redimension the array to hold the correct number of edge flags
    ReDim pEdgeFlags(0 To lngEdgeFlagCount - 1)
    For lLoop = 0 To lngEdgeFlagCount - 1
      'assign to a local IFlagDisplay and IEdgeFlagDisplay variables
      Set pFlagDisplay = pNetworkAnalysisExtFlags.EdgeFlag(lLoop)
      Set pEdgeFlagDisplay = pFlagDisplay
      'co-create a new EdgeFlag object
      Set pNetFlag = New esriNetworkAnalysis.EdgeFlag
      Set pEdgeFlag = pNetFlag
      'assign the properties of the EdgeFlagDisplay object to the EdgeFlag object
      pEdgeFlag.Position = pEdgeFlagDisplay.Percentage
      pNetFlag.UserClassID = pFlagDisplay.FeatureClassID
      pNetFlag.UserID = pFlagDisplay.FID
      pNetFlag.UserSubID = pFlagDisplay.SubID
      'add the new EdgeFlag object to the array
      Set pEdgeFlags(lLoop) = pNetFlag
    Next lLoop
    'add the edge flags to the network solver
    pTraceFlowSolver.PutEdgeOrigins lngEdgeFlagCount, pEdgeFlags(0)
  End If
  
  'next, get the junction flags
  'determine the number of junction flags on the network
  lngJunctionFlagCount = pNetworkAnalysisExtFlags.JunctionFlagCount
  'only execute this next bit if there are junction flags
  If Not lngJunctionFlagCount = 0 Then
    'redimension the array to hold the correct number of junction flags
    ReDim pJunctionFlags(0 To lngJunctionFlagCount - 1)
    For lLoop = 0 To lngJunctionFlagCount - 1
      'assign to a local IFlagDisplay variable
      Set pFlagDisplay = pNetworkAnalysisExtFlags.JunctionFlag(lLoop)
      'co-create a new JunctionFlag object
      Set pNetFlag = New esriNetworkAnalysis.JunctionFlag
      'assign the properties of the JunctionFlagDisplay object to the JunctionFlag object
      pNetFlag.UserClassID = pFlagDisplay.FeatureClassID
      pNetFlag.UserID = pFlagDisplay.FID
      pNetFlag.UserSubID = pFlagDisplay.SubID
      'add the new junction flag to the array of junction flags
      Set pJunctionFlags(lLoop) = pNetFlag
      m_lJuncEID = pNetElems.GetEID(pFlagDisplay.FeatureClassID, pFlagDisplay.FID, _
       pFlagDisplay.SubID, esriETJunction)
    Next lLoop
    'add the junction flags to the network solver
    pTraceFlowSolver.PutJunctionOrigins lngJunctionFlagCount, pJunctionFlags(0)
  End If
  
  'set the option for tracing on indeterminate flow
  'QI for the ITraceTasks interface using IUtilityNetworkAnalysisExt
  Set pTraceTasks = m_pUtilityNetworkAnalysisExt
  pTraceFlowSolver.TraceIndeterminateFlow = pTraceTasks.TraceIndeterminateFlow
  
  'pass the TraceFlowSolver object back the network solver
  Set UTIL_CoreTraceSetup = pTraceFlowSolver
  
  Exit Function
  
ErrorHandler:
  'an unexpected error occured somewhere in the function
  'notify the user that an error occurred and exit the sub
  MsgBox "Unexpected error:" & vbCrLf & Err.Description, vbOKOnly, "My Upstream Trace"
  
End Function