#  Copyright 2025 Zeppelin Bend Pty Ltd
#
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
from zepben.ewb.services.common.resolver import per_length_impedance
from zepben.ewb import (
    NetworkService, AcLineSegment, PerLengthSequenceImpedance, Switch, Breaker,
    ConductingEquipment, NameType, Meter, EnergySource, Terminal
)

# A `NetworkService` is a mutable node breaker network model that implements a subset of
# IEC61968 and IEC61970 CIM classes. It is essentially a collection of `IdentifiedObject`s,
# and they may be added and removed as desired.
network = NetworkService()

print("\n##################\n# ADDING OBJECTS #\n##################\n")
# We start by adding a line segment and a breaker to the network model.
line = AcLineSegment(mrid="acls_123")
breaker = Breaker(mrid="b_456")

print(f"{line} added? {network.add(line)}")
print(f"{breaker} added? {network.add(breaker)}")

# Objects with duplicate mRIDs are not added.
invalid = EnergySource(mrid="acls_123")
print(f"{invalid} added? {network.add(invalid)}")


print("\n####################\n# QUERYING OBJECTS #\n\n####################\n")
# Use the `get` method to query the network model for an object with the specified mRID.
print(f"Identified object with mrid acls_123: {network.get('acls_123')}")
print(f"Identified object with mrid b_456: {network.get('b_456')}")

# A `KeyError` is raised if no object with the specified mRID is in the network model.
try:
    network.get("not_in_network")
except KeyError as error:
    print(error)

# Narrow the desired type with the second parameter. In makes the intent clearer, and
# lets IDEs lint and autocomplete according to the requested type.
print(f"Switch with mrid b_456 is open? {network.get('b_456', Switch).is_open()}")

# A `TypeError` is raised if the object exists in the network model, but is not the correct type.
try:
    network.get("acls_123", Switch)
except TypeError as error:
    print(error)

print("\n##################\n# QUERYING TYPES #\n##################\n")
# You may use the `objects` method to iterate over all objects that inherit a specified type.
# Because the breaker is the only object in the network model that inherits from the `Switch`
# class, this will print "Switch: Breaker{b_456}".
for switch in network.objects(Switch):
    print(f"Switch: {switch}")

# However, both the line and the breaker inherit from `ConductingEquipment`.
# The following line prints "Conducting equipment: AcLineSegment{acls_123}"
# and "Conducting equipment: Breaker{b_456}".
for conducting_equipment in network.objects(ConductingEquipment):
    print(f"Conducting equipment: {conducting_equipment}")

# Remark: Objects generated by network.objects(BaseType) are ordered by the name of their leaf
# class, so all `AcLineSegment`s will appear before all `Breaker`s.
# The `len_of` method returns the number of objects that inherit a specified type.
print(f"Number of switches: {network.len_of(Switch)}")
print(f"Number of conducting equipment: {network.len_of(ConductingEquipment)}")

print("\n#############\n# RESOLVERS #\n#############\n")
# There may be times when you need to reconstruct a network model from an unordered collection
# of identified objects. `NetworkService` allows you to add reference resolvers, which complete
# associations when the remaining object in an association is added.
network.resolve_or_defer_reference(per_length_impedance(line), "plsi_789")

print(f"Network has unresolved references? {network.has_unresolved_references()}")
print(f"plsi_789 has unresolved references? {network.has_unresolved_references('plsi_789')}")
for unresolved_reference in network.get_unresolved_references_from("acls_123"):
    print(f"Unresolved reference from acls_123: {unresolved_reference}")
for unresolved_reference in network.get_unresolved_references_to("plsi_789"):
    print(f"Unresolved reference to plsi_789: {unresolved_reference}")
print(f"Number of unresolved references to plsi_789: {network.num_unresolved_references('plsi_789')}")
print(f"Total unresolved references: {network.num_unresolved_references()}")
print("Adding plsi_789 to the network...")
network.add(PerLengthSequenceImpedance(mrid="plsi_789"))
print(f"Total unresolved references: {network.num_unresolved_references()}")
print(f"PerLengthSequenceImpedance of acls_123: {line.per_length_sequence_impedance}")

print("\n########################\n# CONNECTING TERMINALS #\n########################\n")
# Terminals in a `NetworkService` may be connected using the `connect_terminals` method.
# This automatically creates a connectivity node between the terminals, unless one of the
# terminals is already assigned to one.
t1, t2, t3 = (Terminal(mrid=f"t{i+1}") for i in range(3))
network.add(t1)
network.add(t2)
network.add(t3)
network.connect_terminals(t1, t2)
cn = t1.connectivity_node
print(f"Connected to node {cn}:")
for terminal in cn.terminals:
    print(f"\t{terminal}")

# The mrid of the connectivity node may also be used to connect a terminal
network.connect_by_mrid(t3, cn.mrid)
print(f"Connected to node {cn}:")
for terminal in cn.terminals:
    print(f"\t{terminal}")


print("\n#########\n# NAMES #\n#########\n")
# Apart from identified objects, a `NetworkService` also supports names. Each identified object has
# exactly one mRID, but can have any number of names. Each name has a name type. In this example,
# we add two names of type "NMI" to the network model.
meter1 = Meter()
meter2 = Meter()

name_type = NameType(name="NMI", description="National Meter Identifier")
name_type.get_or_add_name("987654321", line)
name_type.get_or_add_name("546372819", breaker)

network.add_name_type(name_type)
for name in network.get_name_type("NMI").names:
    print(f"NMI name {name.name} is assigned to {name.identified_object}")
for name_type in network.name_types:
    print(f"Network has name type {name_type}")

# Remark: In practice, NMI names are not assigned to lines and breakers.

print("\n####################\n# REMOVING OBJECTS #\n####################\n")

# You may use the `remove` method to remove objects from the network model.
network.remove(line)
print(f"{line} removed successfully.")
network.remove(breaker)
print(f"{breaker} removed successfully.")

# The object does not need to be the one that was added. It just needs to match the type and mRID.
plsi = PerLengthSequenceImpedance(mrid="plsi_789")
network.remove(plsi)
print(f"{plsi} removed successfully.")

# KeyError is raised if no matching object is found.
try:
    network.remove(line)
except KeyError as error:
    print(error)
