// 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.
// 

using System;
using System.Windows.Forms;
using ESRI.ArcGIS.Carto;
using ESRI.ArcGIS.esriSystem;
using ESRI.ArcGIS.Geodatabase;
using ESRI.ArcGIS.Geometry;
using ESRI.ArcGIS.NetworkAnalyst;
using ESRI.ArcGIS.NetworkAnalystUI;

namespace NATraversalResult
{
  public class SelectFeaturesTool : ESRI.ArcGIS.Desktop.AddIns.Tool
  {
    INetworkAnalystExtension m_networkAnalystExtension;

    public SelectFeaturesTool()
    {
      m_networkAnalystExtension = ArcMap.Application.FindExtensionByName("Network Analyst") as INetworkAnalystExtension;
    }

    protected override void OnUpdate()
    {
      Enabled = ArcMap.Application != null;
    }

    private INALayer GetActiveAnalysisLayer()
    {
      if (m_networkAnalystExtension != null)
        return m_networkAnalystExtension.NAWindow.ActiveAnalysis;
      else
        return null;
    }

    protected override bool OnDeactivate()
    {
      return true;
    }

    /// <summary>
    /// Finds the closest feature to the point clicked that corresponds to a NATraversalResultElement in the active analysis layer's traversal result
    /// </summary>
    /// 
    protected override void OnMouseDown(MouseEventArgs e)
    {
      try
      {
        // Get the current network analysis traversal result
        INALayer naLayer = GetActiveAnalysisLayer();
        if (naLayer == null) throw new Exception("Null NALayer");

        INATraversalResult naTraversalResult = (INATraversalResult)naLayer.Context.Result;
        if (naTraversalResult == null || !((INAResult)naTraversalResult).HasValidResult)
          throw new Exception("Active analysis layer has an invalid result.");

        // Get the current map
        IMap map = ArcMap.Document.FocusMap;
        IActiveView activeView = ArcMap.Document.ActiveView;

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

        IProximityOperator proximityOperator = point as IProximityOperator;
        ITopologicalOperator topologicalOperator = point as ITopologicalOperator;
        IPolygon polygon = (IPolygon)topologicalOperator.Buffer(activeView.Extent.Width / 20);
        double closestDistance = activeView.Extent.Width / 20;
        IFeature closestFeature = null;

        // Setup the spatial filter
        ISpatialFilter spatialFilter = 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
        IEnumLayer enumLayer = map.get_Layers(null, true);
        ILayer layer = enumLayer.Next();
        while (layer != null)
        {
          IFeatureLayer featureLayer = layer as IFeatureLayer;
          if (featureLayer != null)
          {
            if (featureLayer.Visible)
            {
              IFeatureClass featureClass = featureLayer.FeatureClass;

              // loop though all the sources in the network, trying to see if it corresponds to a source in the traversal result.
              for (int sourceIndex = 0; sourceIndex < naTraversalResult.SourceCount; sourceIndex++)
              {
                // Check to see if the feature class matches the traversal result source's feature class
                ITable currentTable = featureClass as ITable;
                ITable sourceTable = naTraversalResult.get_Source(sourceIndex).Table;
                if (currentTable == sourceTable)
                {
                  // Since we correspond to a source in the traversal result, find the closest feature to the point clicked.
                  spatialFilter.GeometryField = featureClass.ShapeFieldName;
                  IFeatureCursor featureCursor = featureClass.Search(spatialFilter, false);
                  IFeature feature = featureCursor.NextFeature();
                  while (feature != null)
                  {
                    double distance = proximityOperator.ReturnDistance(feature.Shape);

                    if (distance < closestDistance)
                    {
                      closestFeature = feature;
                      closestDistance = distance;
                    }

                    feature = featureCursor.NextFeature();
                  }
                }
              }
            }
          }

          layer = enumLayer.Next();
        }

        if (closestFeature == null)
          MessageBox.Show("Could not find a feature");
        else
          SelectTraversedFeatures(closestFeature);  // Since we found a feature, select the features traversed to get to it.
      }
      catch (Exception exception)
      {
        MessageBox.Show(exception.Message, "Error");
      }

    }


    /// <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 void SelectTraversedFeatures(IFeature selectedFeature)
    {
      // Get reference to the current map
      IMap map = ArcMap.Document.FocusMap;

      // Get the NATraversalResult
      INALayer naLayer = GetActiveAnalysisLayer();
      INAContext naContext = naLayer.Context;
      INATraversalResult naTraversalResult = (INATraversalResult)naContext.Result;
      INATraversalResultEdit naTraversalResultEdit = (INATraversalResultEdit)naTraversalResult;
      INATraversalResultQuery naTraversalResultQuery = (INATraversalResultQuery)naTraversalResult;

      // 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
      INATraversalResultSource selectedSource = null;
      IFeatureClass[] sourceFeatureClasses = new IFeatureClass[naTraversalResult.SourceCount + 1]; // (will never use index 0)
      ILongArray[] sourceOIDs = new ILongArray[naTraversalResult.SourceCount + 1]; // (will never use index 0)

      for (int sourceIndex = sourceFeatureClasses.GetLowerBound(0); sourceIndex <= sourceFeatureClasses.GetUpperBound(0); sourceIndex++)
      {
        INATraversalResultSource naTraversalResultSource = naTraversalResult.get_SourceByID(sourceIndex);

        if (naTraversalResultSource != null)
        {
          IFeatureClass featureClass = naTraversalResultSource.Table as IFeatureClass;
          if (featureClass != null)
          {
            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 == selectedFeature.Class)
              selectedSource = naTraversalResultSource;

            sourceOIDs[sourceIndex] = new LongArrayClass();
          }
        }
      }

      // Verify the selected feature's class is actually in the network
      if (selectedSource == null)
        throw (new Exception("The selected feature is not part of the network."));

      // 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.
      IQueryFilter queryFilter = 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
      IFeatureClass traversalFClass;
      if (selectedFeature.Shape.GeometryType == esriGeometryType.esriGeometryPoint)
        traversalFClass = naTraversalResultQuery.get_FeatureClass(esriNetworkElementType.esriNETJunction);
      else if (selectedFeature.Shape.GeometryType == esriGeometryType.esriGeometryPolyline)
        traversalFClass = naTraversalResultQuery.get_FeatureClass(esriNetworkElementType.esriNETEdge);
      else
        throw (new Exception("You must start from a point or polyline"));

      if (traversalFClass.FeatureCount(null) == 0)
        throw (new Exception("The traversal result is empty.  Please click solve."));

      // Search through the traversal class to find those corresponding to the selected feature
      INATraversalResultElement naTraversalResultElement;
      IFeatureCursor featureCursor = traversalFClass.Search(queryFilter, true);
      IFeature feature = featureCursor.NextFeature();
      if (feature == null)
        throw (new Exception("You must pick a feature that was traversed."));

      // Loop through all the traversal result elements for that feature
      while (feature != null)
      {
        naTraversalResultElement = (INATraversalResultElement)feature;
        sourceOIDs[naTraversalResultElement.SourceID].Add(naTraversalResultElement.SourceOID);

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

        feature = featureCursor.NextFeature();
      }

      // Select the features in the Map
      SelectFeatures(map, sourceFeatureClasses, sourceOIDs);
    }

    /// <summary> 
    /// This function selects the features corresponding to the arrays of feature classes and OIDs
    /// </summary>
    private void SelectFeatures(IMap map, IFeatureClass[] sourceFeatureClasses, ILongArray[] sourceOIDs)
    {
      IActiveView activeView = (IActiveView)map;

      IEnumLayer enumLayer = map.get_Layers(null, true);
      ILayer layer = enumLayer.Next();
      while (layer != null)
      {
        IFeatureLayer featureLayer = layer as IFeatureLayer;
        IFeatureSelection featureSelection = layer as IFeatureSelection;

        if ((featureLayer != null) && (featureSelection != null) && featureLayer.Visible && featureLayer.Selectable)
        {
          for (int sourceIndex = sourceFeatureClasses.GetLowerBound(0); sourceIndex <= sourceFeatureClasses.GetUpperBound(0); sourceIndex++)
          {
            if (featureLayer.FeatureClass == sourceFeatureClasses[sourceIndex])
            {
              //Clear the current selection on this layer
              featureSelection.Clear();

              activeView.PartialRefresh(esriViewDrawPhase.esriViewGeoSelection, map, null);

              // Get the OIDS we traversed that we want to select
              ILongArray longArray = sourceOIDs[sourceIndex];

              if (longArray.Count > 0)
              {
                // We have some OIDs, ad them to the selection set
                ISelectionSet selectionSet = featureSelection.SelectionSet;

                // Transfer from LongArray to an array of longs
                int numIDs = longArray.Count;
                int[] oids = new int[numIDs];
                for (int oidIndex = oids.GetLowerBound(0); oidIndex < numIDs; oidIndex++)
                  selectionSet.Add(longArray.get_Element(oidIndex));
              }
            }
          }
        }
        layer = enumLayer.Next();
      }

      // Refresh the selection and Map
      ISelectionEvents selectionEvents = (ISelectionEvents)map;
      selectionEvents.SelectionChanged();

      activeView.PartialRefresh(esriViewDrawPhase.esriViewGeoSelection, map, null);
    }


    /// <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 void OutputConnected(INATraversalResultQuery naTraversalResultQuery, INATraversalResultElement inputElement, esriRelDirection searchDirection, ILongArray[] sourceOIDs)
    {
      IFeatureCursor featureCursor;
      if (inputElement.ElementType == esriNetworkElementType.esriNETJunction)
        featureCursor = naTraversalResultQuery.SearchConnected(inputElement, esriNetworkElementType.esriNETEdge, searchDirection, false);
      else
        featureCursor = naTraversalResultQuery.SearchConnected(inputElement, esriNetworkElementType.esriNETJunction, searchDirection, false);

      // Iterate through the connected features.
      // For each feature process it and then recursively call this function again
      // to process the connected features.
      INATraversalResultElement connectedElement = (INATraversalResultElement)featureCursor.NextFeature();
      while (connectedElement != null)
      {
        if (sourceOIDs[connectedElement.SourceID] != null)
          sourceOIDs[connectedElement.SourceID].Add(connectedElement.SourceOID);

        // Recursion
        OutputConnected(naTraversalResultQuery, connectedElement, searchDirection, sourceOIDs);

        connectedElement = (INATraversalResultElement)featureCursor.NextFeature();
      }

    }
  }

}