import numpy as np
from pydrake.common.containers import namedview
from pydrake.common.value import Value
from pydrake.math import RigidTransform, RotationMatrix
from pydrake.systems.analysis import Simulator
from pydrake.systems.framework import BasicVector, LeafSystem
from pydrake.trajectories import PiecewisePolynomialclass MyAdder(LeafSystem):
    def __init__(self):
        super().__init__()  # Don't forget to initialize the base class.
        self._a_port = self.DeclareVectorInputPort(name="a", size=2)
        self._b_port = self.DeclareVectorInputPort(name="b", size=2)
        self.DeclareVectorOutputPort(name="sum", size=2, calc=self.CalcSum)
        self.DeclareVectorOutputPort(name="difference",
                                     size=2,
                                     calc=self.CalcDifference)

    def CalcSum(self, context, output):
        # Evaluate the input ports to obtain the 2x1 vectors.
        a = self._a_port.Eval(context)
        b = self._b_port.Eval(context)

        # Write the sum into the output vector.
        output.SetFromVector(a + b)

    def CalcDifference(self, context, output):
        # Evaluate the input ports to obtain the 2x1 vectors.
        a = self._a_port.Eval(context)
        b = self._b_port.Eval(context)

        # Write the difference into output vector.
        output.SetFromVector(a - b)

# Construct an instance of this system and a context.
system = MyAdder()
context = system.CreateDefaultContext()

# Fix the input ports to some constant values.
system.GetInputPort("a").FixValue(context, [3, 4])
system.GetInputPort("b").FixValue(context, [1, 2])

# Evaluate the output ports.
print(f"sum: {system.GetOutputPort('sum').Eval(context)}")
print(f"difference: {system.GetOutputPort('difference').Eval(context)}")class RotateAboutZ(LeafSystem):
    def __init__(self):
        super().__init__()  # Don't forget to initialize the base class.
        self.DeclareAbstractInputPort(name="in",
                                      model_value=Value(RigidTransform()))
        self.DeclareAbstractOutputPort(
            name="out",
            alloc=lambda: Value(RigidTransform()),
            calc=self.CalcOutput)

    def CalcOutput(self, context, output):
        # Evaluate the input port to obtain the RigidTransform.
        X_1 = system.get_input_port().Eval(context)

        X_2 = RigidTransform(RotationMatrix.MakeZRotation(np.pi / 2)) @ X_1

        # Set the output RigidTransform.
        output.set_value(X_2)

# Construct an instance of this system and a context.
system = RotateAboutZ()
context = system.CreateDefaultContext()

# Fix the input port to a constant value.
system.get_input_port().FixValue(context, RigidTransform())

# Evaluate the output port.
print(f"output: {system.get_output_port(0).Eval(context)}")class SimpleDiscreteTimeSystem(LeafSystem):
    def __init__(self):
        super().__init__()

        state_index = self.DeclareDiscreteState(1)  # One state variable.
        self.DeclareStateOutputPort("y", state_index)  # One output: y=x.
        self.DeclarePeriodicDiscreteUpdateEvent(
            period_sec=1.0,  # One second time step.
            offset_sec=0.0,  # The first event is at time zero.
            update=self.Update) # Call the Update method defined below.

    # x[n+1] = x^3[n].
    def Update(self, context, discrete_state):
        x = context.get_discrete_state_vector().GetAtIndex(0)
        x_next = x**3
        discrete_state.get_mutable_vector().SetAtIndex(0, x_next)

# Instantiate the System.
system = SimpleDiscreteTimeSystem()
simulator = Simulator(system)
context = simulator.get_mutable_context()

# Set the initial conditions: x[0] = [0.9].
context.get_mutable_discrete_state_vector().SetFromVector([0.9])

# Run the simulation.
simulator.AdvanceTo(4.0)
print(context.get_discrete_state_vector())# Define the system.
class SimpleContinuousTimeSystem(LeafSystem):
    def __init__(self):
        super().__init__()

        state_index = self.DeclareContinuousState(1)  # One state variable.
        self.DeclareStateOutputPort("y", state_index)  # One output: y=x.

    # xdot(t) = -x(t) + x^3(t).
    def DoCalcTimeDerivatives(self, context, derivatives):
        x = context.get_continuous_state_vector().GetAtIndex(0)
        xdot = -x + x**3
        derivatives.get_mutable_vector().SetAtIndex(0, xdot)

# Instantiate the System.
system = SimpleContinuousTimeSystem()
simulator = Simulator(system)
context = simulator.get_mutable_context()

# Set the initial conditions: x(0) = [0.9]
context.SetContinuousState([0.9])

# Run the simulation.
simulator.AdvanceTo(4.0)
print(context.get_continuous_state_vector())class AbstractStateSystem(LeafSystem):
    def __init__(self):
        super().__init__()

        self._traj_index = self.DeclareAbstractState(
            Value(PiecewisePolynomial()))
        self.DeclarePeriodicUnrestrictedUpdateEvent(period_sec=1.0,
                                                    offset_sec=0.0,
                                                    update=self.Update)

    def Update(self, context, state):
        t = context.get_time()
        traj = PiecewisePolynomial.FirstOrderHold(
            [t, t + 1],
            np.array([[-np.pi / 2.0 + 1., -np.pi / 2.0 - 1.], [-2., 2.]]))
        # Update the state
        state.get_mutable_abstract_state(int(self._traj_index)).set_value(traj)



system = AbstractStateSystem()
simulator = Simulator(system)
context = simulator.get_mutable_context()

# Set an initial condition for the abstract state.
context.SetAbstractState(0, PiecewisePolynomial())

# Run the simulation.
simulator.AdvanceTo(4.0)
traj = context.get_abstract_state(0).get_value()
print(f"breaks: {traj.get_segment_times()}")
print(f"traj({context.get_time()}) = {traj.value(context.get_time())}")class SystemWithParameters(LeafSystem):
    def __init__(self):
        super().__init__()  # Don't forget to initialize the base class.

        self.DeclareNumericParameter(BasicVector([1.2, 3.4]))
        self.DeclareAbstractParameter(
            Value(RigidTransform(RotationMatrix.MakeXRotation(np.pi / 6))))

        # Declare output ports to demonstrate how to access the parameters in
        # system methods.
        self.DeclareVectorOutputPort(name="numeric",
                                     size=2,
                                     calc=self.OutputNumeric)
        self.DeclareAbstractOutputPort(
            name="abstract",
            alloc=lambda: Value(RigidTransform()),
            calc=self.OutputAbstract)

    def OutputNumeric(self, context, output):
        output.SetFromVector(context.get_numeric_parameter(0).get_value())

    def OutputAbstract(self, context, output):
        output.set_value(context.get_abstract_parameter(0).get_value())

# Construct an instance of this system and a context.
system = SystemWithParameters()
context = system.CreateDefaultContext()

# Evaluate the output ports.
print(f"numeric: {system.get_output_port(0).Eval(context)}")
print(f"abstract: {system.get_output_port(1).Eval(context)}")class MyPublishingSystem(LeafSystem):
    def __init__(self):
        super().__init__()

        # Calling `ForcePublish()` will trigger the callback.
        self.DeclareForcedPublishEvent(self.Publish)

        # Publish once every second.
        self.DeclarePeriodicPublishEvent(period_sec=1,
                                         offset_sec=0,
                                         publish=self.Publish)
        
    def Publish(self, context):
        print(f"Publish() called at time={context.get_time()}")

system = MyPublishingSystem()
simulator = Simulator(system)
simulator.AdvanceTo(5.3)

# We can also "force" a publish at a arbitrary time.
print("\ncalling ForcedPublish:")
system.ForcedPublish(simulator.get_context())
# Define the system.
class NamedViewDemo(LeafSystem):
    MyDiscreteState = namedview("MyDiscreteState", ["a", "b"])
    MyContinuousState = namedview("MyContinuousState", ["x", "z", "theta"])
    MyOutput = namedview("MyOutput", ["x","a"])

    def __init__(self):
        super().__init__()

        self.DeclareDiscreteState(2)
        self.DeclarePeriodicDiscreteUpdateEvent(
            period_sec=1.0,
            offset_sec=0.0,
            update=self.DiscreteUpdate)
        self.DeclareContinuousState(3)
        self.DeclareVectorOutputPort(name="out", size=2, calc=self.CalcOutput)

    def DiscreteUpdate(self, context, discrete_values):
        discrete_state = self.MyDiscreteState(
            context.get_discrete_state_vector().value())
        continuous_state = self.MyContinuousState(
            context.get_continuous_state_vector().CopyToVector())
        next_state = self.MyDiscreteState(discrete_values.get_mutable_value())
        # Now we can compute the next state by referencing each element by name.
        next_state.a = discrete_state.a + 1
        next_state.b = discrete_state.b + continuous_state.x

    def DoCalcTimeDerivatives(self, context, derivatives):
        continuous_state = self.MyContinuousState(
            context.get_continuous_state_vector().CopyToVector())
        dstate_dt = self.MyContinuousState(continuous_state[:])
        dstate_dt.x = -continuous_state.x
        dstate_dt.z = -continuous_state.z
        dstate_dt.theta = -np.arctan2(continuous_state.z, continuous_state.x)
        derivatives.SetFromVector(dstate_dt[:])

    def CalcOutput(self, context, output):
        discrete_state = self.MyDiscreteState(
            context.get_discrete_state_vector().value())
        continuous_state = self.MyContinuousState(
            context.get_continuous_state_vector().CopyToVector())
        out = self.MyOutput(output.get_mutable_value())
        out.x = continuous_state.x
        out.a = discrete_state.a

# Instantiate the System.
system = NamedViewDemo()
simulator = Simulator(system)
context = simulator.get_mutable_context()

# Set the initial conditions.
initial_discrete_state = NamedViewDemo.MyDiscreteState([3, 4])
context.SetDiscreteState(initial_discrete_state[:])
initial_continuous_state = NamedViewDemo.MyContinuousState.Zero()
initial_continuous_state.x = 0.5
initial_continuous_state.z = 0.92
initial_continuous_state.theta = 0.23
context.SetContinuousState(initial_continuous_state[:])

# Run the simulation.
simulator.AdvanceTo(4.0)
print(
    NamedViewDemo.MyDiscreteState(context.get_discrete_state_vector().value()))
print(
    NamedViewDemo.MyContinuousState(
        context.get_continuous_state_vector().CopyToVector()))
print(NamedViewDemo.MyOutput(system.get_output_port().Eval(context)))
from pydrake.systems.framework import LeafSystem_
from pydrake.systems.scalar_conversion import TemplateSystem
from pydrake.autodiffutils import AutoDiffXd
from pydrake.symbolic import Expression

@TemplateSystem.define("RunningCost_")
def RunningCost_(T):

    class Impl(LeafSystem_[T]):

        def _construct(self, converter=None, Q=np.eye(2)):
            super().__init__(converter)
            self._Q = Q
            self._state_port = self.DeclareVectorInputPort("state", 2)
            self._command_port = self.DeclareVectorInputPort("command", 1)
            self.DeclareVectorOutputPort("cost", 1, self.CostOutput)

        def _construct_copy(self, other, converter=None):
            # Any member fields (e.g. configuration values) need to be
            # transferred here from `other` to `self`.
            Impl._construct(self, converter=converter, Q=other._Q)

        def CostOutput(self, context, output):
            x = self._state_port.Eval(context)
            u = self._command_port.Eval(context)[0]
            output[0] = x.dot(self._Q.dot(x)) + u**2

    return Impl

RunningCost = RunningCost_[None]  # Default instantiation.# Having done this, we can still use the system in the original way:
system = RunningCost(Q=np.diag([10, 1]))
context = system.CreateDefaultContext()
system.get_input_port(0).FixValue(context, [1, 2])
system.get_input_port(1).FixValue(context, [3])
print(system.get_output_port().Eval(context))

# But we can also now use an autodiff or symbolic versions of the system,
# either by declaring them directly:
system_ad = RunningCost_[AutoDiffXd]()
system_symbolic = RunningCost_[Expression]()

# or by scalar conversion:
system_ad = system.ToAutoDiffXd()
system_symbolic = system.ToSymbolic()

# We can also convert the time, state, parameters, and (if needed) input ports
# from the original system:
context_symbolic = system_symbolic.CreateDefaultContext()
context_symbolic.SetTimeStateAndParametersFrom(context)
system_symbolic.FixInputPortsFrom(system, context, context_symbolic)
print(system_symbolic.get_output_port().Eval(context_symbolic))