import numpy as np
import pydot
from IPython.display import SVG, display
from pydrake.examples import PendulumPlant
from pydrake.systems.controllers import PidController
from pydrake.systems.framework import Diagram, DiagramBuilder, LeafSystem
from pydrake.systems.primitives import AffineSystem, LogVectorOutputdef MakePidControlledPendulum():
    # Use a DiagramBuilder to add and connect the subsystems.
    builder = DiagramBuilder()

    # First add the pendulum.
    pendulum = builder.AddNamedSystem("pendulum", PendulumPlant())

    # Add a PID controller.
    controller = builder.AddNamedSystem("controller",
                                        PidController(kp=[10.], ki=[1.], kd=[1.]))

    # Now "wire up" the controller to the plant.
    builder.Connect(pendulum.get_state_output_port(),
                    controller.get_input_port_estimated_state())
    builder.Connect(controller.get_output_port_control(), pendulum.get_input_port())

    # Make the desired_state input of the controller an input to the diagram.
    builder.ExportInput(controller.get_input_port_desired_state())
    # Make the pendulum state an output from the diagram.
    builder.ExportOutput(pendulum.get_state_output_port())

    # Log the state of the pendulum.
    logger = LogVectorOutput(pendulum.get_state_output_port(), builder)
    logger.set_name("logger")

    pid_controlled_pendulum = builder.Build()
    pid_controlled_pendulum.set_name("PID-controlled Pendulum")

    return pid_controlled_pendulum, pendulum

pid_controlled_pendulum, pendulum = MakePidControlledPendulum()

# Visualize the diagram.
display(
    SVG(
        pydot.graph_from_dot_data(
            pid_controlled_pendulum.GetGraphvizString(
                max_depth=2))[0].create_svg()))
second_builder = DiagramBuilder()

try:
    second_builder.AddSystem(pendulum)
except RuntimeError as err:
    print(err)builder = DiagramBuilder()

# We make the PID-controlled pendulum again here, to avoid C++ ownership issues.
pid_controlled_pendulum, pendulum = MakePidControlledPendulum()
builder.AddSystem(pid_controlled_pendulum)

# Shift desired theta by PI with the system y= u + [pi;0]
shift_desired = builder.AddNamedSystem(
    "shift_desired", AffineSystem(D=np.eye(2), y0=[np.pi, 0]))
builder.ExportInput(shift_desired.get_input_port(),
                    "controller_desired_state")

# Connect the shift to the *exported input port* from the subdiagram.
builder.Connect(shift_desired.get_output_port(),
                pid_controlled_pendulum.get_input_port())

# Shift actual theta by -PI.
shift_actual = builder.AddNamedSystem(
    "shift_actual", AffineSystem(D=np.eye(2), y0=[-np.pi, 0]))
builder.Connect(pid_controlled_pendulum.get_output_port(),
                shift_actual.get_input_port())
builder.ExportOutput(shift_actual.get_output_port(), "pendulum_state")

diagram = builder.Build()
diagram.set_name("PID-controlled Pendulum (with θ shifted by π)")

# Visualize the diagram (max depth=1)
display(
    SVG(
        pydot.graph_from_dot_data(
            diagram.GetGraphvizString(max_depth=1))[0].create_svg()))
# Visualize the diagram (max depth=2)
display(
    SVG(
        pydot.graph_from_dot_data(
            diagram.GetGraphvizString(max_depth=2))[0].create_svg()))
assert(pid_controlled_pendulum.GetSubsystemByName("pendulum") == pendulum)diagram_context = diagram.CreateDefaultContext()
print(diagram_context)pendulum_context = pendulum.GetMyContextFromRoot(diagram_context)
pendulum_context.SetContinuousState([1.2, 0.5])
# The PendulumPlant class provides some helper methods for working with its
# Context.
pendulum_state = pendulum.get_state(pendulum_context)
print(f"θ = {pendulum_state.theta()}, θ̇ = {pendulum_state.thetadot()}")
# Observe that the root system's Context was also updated.
print(diagram_context)builder = DiagramBuilder()

pid_controlled_pendulum, pendulum = MakePidControlledPendulum()
builder.AddSystem(pid_controlled_pendulum)

# Shift desired theta by PI with the system y= u + [pi;0]
shift_desired = builder.AddNamedSystem(
    "shift_desired", AffineSystem(D=np.eye(2), y0=[np.pi, 0]))

# ** WRONG ** This doesn't work, but gives a helpful message.
controller = pid_controlled_pendulum.GetSubsystemByName("controller")
try:
    builder.Connect(shift_desired.get_output_port(),
                    controller.get_input_port_desired_state())
except RuntimeError as err:
    print(err)class MyLeafSystem(LeafSystem):
    def __init__(self):
        super().__init__()

        self.DeclareContinuousState(1)
        self.DeclareForcedPublishEvent(self.Publish)
        self.DeclarePeriodicPublishEvent(period_sec=1,
                                         offset_sec=0,
                                         publish=self.Publish)

    def DoCalcTimeDerivatives(self, context, derivatives):
        x = context.get_continuous_state_vector().GetAtIndex(0)
        print(f"{self.get_name()}: DoCalcTimeDerivatives()")
        derivatives.get_mutable_vector().SetAtIndex(0, -x)

    def Publish(self, context):
        print(f"{self.get_name()}: Publish()")

builder = DiagramBuilder()
builder.AddNamedSystem("system1", MyLeafSystem())
builder.AddNamedSystem("system2", MyLeafSystem())
diagram = builder.Build()

context = diagram.CreateDefaultContext()

# To evaluate the time derivatives of the diagram, diagram evaluates the time
# derivatives of the subsystems.
print("diagram.EvalTimeDerivatives()")
diagram.EvalTimeDerivatives(context)

# A ForcedPublish on the diagram calls ForcedPublish on the subsystems.
print("diagram.ForcedPublish()")
diagram.ForcedPublish(context)builder = DiagramBuilder()
# AffineSystem is a primitive which supports all scalar types, and
# scalar-conversion.
builder.AddSystem(AffineSystem(y0=[2, 3]))
builder.AddSystem(AffineSystem(D=np.eye(2)))
diagram = builder.Build()

diagram_autodiff = diagram.ToAutoDiffXd()
diagram_symbolic = diagram.ToSymbolic()builder = DiagramBuilder()
# MyLeafSystem (defined above) did not implement scalar support, so the
# resulting Diagram will not, either.
builder.AddSystem(MyLeafSystem())
diagram = builder.Build()

print("calling ToAutoDiffXd()")
try:
    diagram_autodiff = diagram.ToAutoDiffXd()
except RuntimeError as err:
    print(err)

print("\ncalling ToSymbolic()")
try:
    diagram_symbolic = diagram.ToSymbolic()
except RuntimeError as err:
    print(err)
class CustomDiagram(Diagram):

    def __init__(self):
        Diagram.__init__(self)
        builder = DiagramBuilder()
        builder = DiagramBuilder()
        self.system1 = builder.AddNamedSystem("system1", MyLeafSystem())
        self.system2 = builder.AddNamedSystem("system2", MyLeafSystem())
        # Instead of builder.Build(), we call
        builder.BuildInto(self)

    def get_system1(self):
        return self.system1
    
    def get_system2(self):
        return self.system2
    
diagram = CustomDiagram()
context = diagram.CreateDefaultContext()
print(context)
