import flags
import numpy
import pydrake.math
import pydrake.multibody.plant
import pydrake.symbolic
import pydrake.systems.analysis
import pydrake.systems.framework
import pydrake.systems.primitives
import pydrake.trajectories
from _typeshed import Incomplete
from typing import Any, Callable, ClassVar, overload

class DynamicProgrammingOptions:
    class PeriodicBoundaryCondition:
        high: float
        low: float
        state_index: int
        def __init__(self, state_index: int, low: float, high: float) -> None: ...
    assume_non_continuous_states_are_fixed: bool
    convergence_tol: float
    discount_factor: float
    input_port_index: pydrake.systems.framework.InputPortSelection | pydrake.systems.framework.InputPortIndex
    periodic_boundary_conditions: list[DynamicProgrammingOptions.PeriodicBoundaryCondition]
    visualization_callback: Callable[[int, pydrake.math.BarycentricMesh, numpy.ndarray[numpy.float64[1, n]], numpy.ndarray[numpy.float64[m, n]]], None]
    def __init__(self) -> None: ...

class FiniteHorizonLinearQuadraticRegulatorOptions:
    N: numpy.ndarray[numpy.float64[m, n]] | None
    Qf: numpy.ndarray[numpy.float64[m, n]] | None
    input_port_index: pydrake.systems.framework.InputPortSelection | pydrake.systems.framework.InputPortIndex
    simulator_config: pydrake.systems.analysis.SimulatorConfig
    u0: pydrake.trajectories.Trajectory
    ud: Incomplete
    use_square_root_method: bool
    x0: pydrake.trajectories.Trajectory
    xd: Any
    def __init__(self) -> None: ...

class FiniteHorizonLinearQuadraticRegulatorResult:
    def __init__(self, *args, **kwargs) -> None: ...
    @property
    def K(self) -> pydrake.trajectories.Trajectory: ...
    @property
    def S(self) -> pydrake.trajectories.Trajectory: ...
    @property
    def k0(self) -> pydrake.trajectories.Trajectory: ...
    @property
    def s0(self) -> pydrake.trajectories.Trajectory: ...
    @property
    def sx(self) -> pydrake.trajectories.Trajectory: ...
    @property
    def u0(self) -> pydrake.trajectories.Trajectory: ...
    @property
    def x0(self) -> pydrake.trajectories.Trajectory: ...

class InverseDynamics(pydrake.systems.framework.LeafSystem):
    class InverseDynamicsMode:
        __members__: ClassVar[dict] = ...  # read-only
        __entries: ClassVar[dict] = ...
        kGravityCompensation: ClassVar[InverseDynamics.InverseDynamicsMode] = ...
        kInverseDynamics: ClassVar[InverseDynamics.InverseDynamicsMode] = ...
        def __init__(self, value: int) -> None: ...
        def __eq__(self, other: object) -> bool: ...
        def __hash__(self) -> int: ...
        def __index__(self) -> int: ...
        def __int__(self) -> int: ...
        def __ne__(self, other: object) -> bool: ...
        @property
        def name(self) -> str: ...
        @property
        def value(self) -> int: ...
    kGravityCompensation: ClassVar[InverseDynamics.InverseDynamicsMode] = ...
    kInverseDynamics: ClassVar[InverseDynamics.InverseDynamicsMode] = ...
    def __init__(self, plant: pydrake.multibody.plant.MultibodyPlant, mode: InverseDynamics.InverseDynamicsMode = ..., plant_context: pydrake.systems.framework.Context = ...) -> None: ...
    def get_input_port_desired_acceleration(self) -> pydrake.systems.framework.InputPort: ...
    def get_input_port_estimated_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_output_port_generalized_force(self) -> pydrake.systems.framework.OutputPort: ...
    def is_pure_gravity_compensation(self) -> bool: ...

class InverseDynamicsController(pydrake.systems.framework.Diagram):
    def __init__(self, robot: pydrake.multibody.plant.MultibodyPlant, kp: numpy.ndarray[numpy.float64[m, 1]], ki: numpy.ndarray[numpy.float64[m, 1]], kd: numpy.ndarray[numpy.float64[m, 1]], has_reference_acceleration: bool, plant_context: pydrake.systems.framework.Context = ...) -> None: ...
    def get_input_port_desired_acceleration(self) -> pydrake.systems.framework.InputPort: ...
    def get_input_port_desired_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_input_port_estimated_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_multibody_plant_for_control(self) -> pydrake.multibody.plant.MultibodyPlant: ...
    def get_output_port_control(self) -> pydrake.systems.framework.OutputPort: ...
    def get_output_port_generalized_force(self) -> pydrake.systems.framework.OutputPort: ...
    def set_integral_value(self, arg0: pydrake.systems.framework.Context, arg1: numpy.ndarray[numpy.float64[m, 1]]) -> None: ...

class JointStiffnessController(pydrake.systems.framework.LeafSystem):
    def __init__(self, plant: pydrake.multibody.plant.MultibodyPlant, kp: numpy.ndarray[numpy.float64[m, 1]], kd: numpy.ndarray[numpy.float64[m, 1]]) -> None: ...
    def get_input_port_desired_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_input_port_estimated_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_multibody_plant(self) -> pydrake.multibody.plant.MultibodyPlant: ...
    def get_output_port_actuation(self) -> pydrake.systems.framework.OutputPort: ...

class PeriodicBoundaryCondition:
    high: float
    low: float
    state_index: int
    def __init__(self, state_index: int, low: float, high: float) -> None: ...

class PidControlledSystem(pydrake.systems.framework.Diagram):
    @overload
    def __init__(self, plant: pydrake.systems.framework.System, kp: float, ki: float, kd: float, state_output_port_index: int = ..., plant_input_port_index: int = ...) -> None: ...
    @overload
    def __init__(self, plant: pydrake.systems.framework.System, kp: numpy.ndarray[numpy.float64[m, 1]], ki: numpy.ndarray[numpy.float64[m, 1]], kd: numpy.ndarray[numpy.float64[m, 1]], state_output_port_index: int = ..., plant_input_port_index: int = ...) -> None: ...
    @overload
    def __init__(self, plant: pydrake.systems.framework.System, feedback_selector: numpy.ndarray[numpy.float64[m, n]], kp: float, ki: float, kd: float, state_output_port_index: int = ..., plant_input_port_index: int = ...) -> None: ...
    def get_control_input_port(self) -> pydrake.systems.framework.InputPort: ...
    def get_state_input_port(self) -> pydrake.systems.framework.InputPort: ...
    def get_state_output_port(self) -> pydrake.systems.framework.OutputPort: ...

class PidController(pydrake.systems.framework.LeafSystem):
    @overload
    def __init__(self, kp: numpy.ndarray[numpy.float64[m, 1]], ki: numpy.ndarray[numpy.float64[m, 1]], kd: numpy.ndarray[numpy.float64[m, 1]]) -> None: ...
    @overload
    def __init__(self, state_projection: numpy.ndarray[numpy.float64[m, n]], kp: numpy.ndarray[numpy.float64[m, 1]], ki: numpy.ndarray[numpy.float64[m, 1]], kd: numpy.ndarray[numpy.float64[m, 1]]) -> None: ...
    @overload
    def __init__(self, state_projection: numpy.ndarray[numpy.float64[m, n]], output_projection: numpy.ndarray[numpy.float64[m, n]], kp: numpy.ndarray[numpy.float64[m, 1]], ki: numpy.ndarray[numpy.float64[m, 1]], kd: numpy.ndarray[numpy.float64[m, 1]]) -> None: ...
    def get_Kd_vector(self) -> numpy.ndarray[numpy.float64[m, 1]]: ...
    def get_Ki_vector(self) -> numpy.ndarray[numpy.float64[m, 1]]: ...
    def get_Kp_vector(self) -> numpy.ndarray[numpy.float64[m, 1]]: ...
    def get_input_port_desired_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_input_port_estimated_state(self) -> pydrake.systems.framework.InputPort: ...
    def get_output_port_control(self) -> pydrake.systems.framework.OutputPort: ...

def DiscreteTimeLinearQuadraticRegulator(A: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], B: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], Q: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], R: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], N: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous] = ...) -> tuple[numpy.ndarray[numpy.float64[m, n]], numpy.ndarray[numpy.float64[m, n]]]: ...
def FiniteHorizonLinearQuadraticRegulator(system: pydrake.systems.framework.System, context: pydrake.systems.framework.Context, t0: float, tf: float, Q: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], R: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], options: FiniteHorizonLinearQuadraticRegulatorOptions = ...) -> FiniteHorizonLinearQuadraticRegulatorResult: ...
def FittedValueIteration(arg0: pydrake.systems.analysis.Simulator, arg1: Callable[[pydrake.systems.framework.Context], float], arg2: list[set[float]], arg3: list[set[float]], arg4: float, arg5: DynamicProgrammingOptions) -> tuple[pydrake.systems.primitives.BarycentricMeshSystem, numpy.ndarray[numpy.float64[1, n]]]: ...
def LinearProgrammingApproximateDynamicProgramming(arg0: pydrake.systems.analysis.Simulator, arg1: Callable[[pydrake.systems.framework.Context], float], arg2: Callable[[numpy.ndarray[numpy.float64[m, 1]], numpy.ndarray[object[m, 1]]], pydrake.symbolic.Expression], arg3: int, arg4: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], arg5: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], arg6: float, arg7: DynamicProgrammingOptions) -> numpy.ndarray[numpy.float64[m, 1]]: ...
def LinearQuadraticRegulator(A: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], B: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], Q: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], R: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], N: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous] = ..., F: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous] = ...) -> tuple[numpy.ndarray[numpy.float64[m, n]], numpy.ndarray[numpy.float64[m, n]]]: ...
def MakeFiniteHorizonLinearQuadraticRegulator(system: pydrake.systems.framework.System, context: pydrake.systems.framework.Context, t0: float, tf: float, Q: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], R: numpy.ndarray[numpy.float64[m, n], flags.f_contiguous], options: FiniteHorizonLinearQuadraticRegulatorOptions = ...) -> pydrake.systems.framework.System: ...
