' Copyright 2012 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 ESRI.ArcGIS.Schematic

Public Class StationsInContainers
  Inherits ESRI.ArcGIS.Desktop.AddIns.Extension

  ' The OnAfterLoadDiagram is used to define the relations between elements
  ' contained in a diagram. It is called by the AfterLoadDiagram 
  ' event defined in the schematic project. The diagram contains Stations and 
  ' Containers elements. For the Station element type, a particular attribute 
  ' named RelatedFeeder has been created. This attribute is used to identify 
  ' the container the station is related to. Near the top of the procedure, 
  ' a new schematic relation is created (SchematicContainerManager). The 
  ' ISchematicRelationControllerEdit CreateRelation method is used to specify 
  ' that the station is related to its container. The value set for the 
  ' RelatedFeeder attribute specifies whether or not the station is related 
  ' to a container.

  Private m_DatasetMgr As SchematicDatasetManager

  Protected Overrides Sub OnStartup()
    m_DatasetMgr = New SchematicDatasetManager
    AddHandler m_DatasetMgr.AfterLoadDiagram, AddressOf OnAfterLoadDiagram
    AddHandler m_DatasetMgr.AfterRefreshDiagram, AddressOf OnAfterRefreshDiagram
  End Sub

  Protected Overrides Sub OnShutdown()
    RemoveHandler m_DatasetMgr.AfterLoadDiagram, AddressOf OnAfterLoadDiagram
    RemoveHandler m_DatasetMgr.AfterRefreshDiagram, AddressOf OnAfterRefreshDiagram
    m_DatasetMgr = Nothing
  End Sub

  Public Sub OnAfterLoadDiagram(ByVal inMemoryDiagram As ESRI.ArcGIS.Schematic.ISchematicInMemoryDiagram)
    If (State <> ESRI.ArcGIS.Desktop.AddIns.ExtensionState.Enabled) Then Return

    Dim schemContainerClass As ISchematicElementClass = Nothing
    Dim schemElementClass As ISchematicElementClass
    Dim schemStationClass As ISchematicElementClass = Nothing
    Dim enumElementsInContainer As IEnumSchematicInMemoryFeature
    Dim enumContainerElements As IEnumSchematicInMemoryFeature
    Dim schemFeature As ISchematicInMemoryFeature = Nothing
    Dim schemContainerFeature As ISchematicInMemoryFeature = Nothing
    Dim feederOID As Object
    Dim containerNameID As String
    Dim enumElementClass As IEnumSchematicElementClass
    Dim schemAttributeContainer As ISchematicAttributeContainer
    Dim schemAttributeRelatedFeeder As ISchematicAttribute
    Dim schemElement As ISchematicElement = Nothing
    Dim schemRelationController As ISchematicRelationController
    Dim schemRelationControllerEdit As ISchematicRelationControllerEdit
    Dim colContElem As New Collection

    ' Getting SchematicFeature Class Stations and Containers
    enumElementClass = inMemoryDiagram.SchematicDiagramClass.AssociatedSchematicElementClasses
    enumElementClass.Reset()
    schemElementClass = enumElementClass.Next
    While (schemElementClass IsNot Nothing)
      If schemElementClass.Name = "Stations" Then schemStationClass = schemElementClass
      If schemElementClass.Name = "Containers" Then schemContainerClass = schemElementClass
      If (schemStationClass IsNot Nothing AndAlso schemContainerClass IsNot Nothing) Then Exit While

      schemElementClass = enumElementClass.Next
    End While
    ' go out if schemStationClass or schemContainerClass are null
    If (schemStationClass Is Nothing OrElse schemContainerClass Is Nothing) Then Return

    ' Getting the Stations elements that will be displayed in the containers
    enumElementsInContainer = inMemoryDiagram.GetSchematicInMemoryFeaturesByClass(schemStationClass)
    If (enumElementsInContainer Is Nothing) Then Return

    ' Creating the Schematic Container Manager
    schemRelationController = New SchematicRelationController

    ' Creating the Schematic Container Editor that will be used to define the relation between the stations and their container
    schemRelationControllerEdit = schemRelationController

    ' Defining each Container element as a schematic container
    enumContainerElements = inMemoryDiagram.GetSchematicInMemoryFeaturesByClass(schemContainerClass)

    ' Add Container Element to a collection
    enumContainerElements.Reset()
    schemContainerFeature = enumContainerElements.Next
    While (schemContainerFeature IsNot Nothing)
      colContElem.Add(schemContainerFeature, schemContainerFeature.Name)
      schemContainerFeature = enumContainerElements.Next
    End While

    ' Setting the relation between each station and its related container
    enumElementsInContainer.Reset()
    schemFeature = enumElementsInContainer.Next

    While (schemFeature IsNot Nothing)
      ' The relation is specified by the RelatedFeeder attribute value defined for each station
      schemAttributeContainer = CType(schemFeature.SchematicElementClass, ISchematicAttributeContainer)
      schemAttributeRelatedFeeder = schemAttributeContainer.GetSchematicAttribute("RelatedFeeder")

      If schemAttributeRelatedFeeder IsNot Nothing Then
        feederOID = schemAttributeRelatedFeeder.GetValue(CType(schemFeature, ISchematicObject))
        If feederOID IsNot Nothing Then
          containerNameID = "Container-" & feederOID.ToString()

          Try
            ' Retrieve Container Element in the collection
            schemContainerFeature = CType(colContElem(containerNameID), ISchematicInMemoryFeature)
            ' Create relation
            schemRelationControllerEdit.CreateRelation(schemFeature, schemContainerFeature)
          Catch
          End Try
        End If
      End If


      schemContainerFeature = Nothing
      schemFeature = enumElementsInContainer.Next
    End While
  End Sub

  Public Sub OnAfterRefreshDiagram(ByVal inMemoryDiagram As ESRI.ArcGIS.Schematic.ISchematicInMemoryDiagram)
    OnAfterLoadDiagram(inMemoryDiagram)
  End Sub

End Class