import enum
import nanobind
import pydrake.multibody.plant
import pydrake.systems.analysis
import pydrake.systems.framework
import pydrake.trajectories
from _typeshed import Incomplete
from typing import Any, Callable, ClassVar, overload

DiscreteTimeLinearQuadraticRegulator: nanobind.nb_func
FiniteHorizonLinearQuadraticRegulator: nanobind.nb_func
FittedValueIteration: nanobind.nb_func
LinearProgrammingApproximateDynamicProgramming: nanobind.nb_func
LinearQuadraticRegulator: nanobind.nb_func
MakeFiniteHorizonLinearQuadraticRegulator: nanobind.nb_func

class DynamicProgrammingOptions:
    class PeriodicBoundaryCondition:
        high: Incomplete
        low: Incomplete
        state_index: Incomplete
        def __init__(self, state_index: int, low: float, high: float) -> None: ...
    assume_non_continuous_states_are_fixed: Incomplete
    convergence_tol: float
    discount_factor: float
    input_port_index: pydrake.systems.framework.InputPortSelection | pydrake.systems.framework.InputPortIndex
    periodic_boundary_conditions: Incomplete
    visualization_callback: Incomplete
    def __init__(self) -> None: ...

class FiniteHorizonLinearQuadraticRegulatorOptions:
    N: Incomplete
    Qf: Incomplete
    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): ...
    @property
    def k0(self): ...
    @property
    def s0(self): ...
    @property
    def sx(self): ...
    @property
    def u0(self): ...
    @property
    def x0(self): ...

class InverseDynamics(pydrake.systems.framework.LeafSystem):
    class InverseDynamicsMode(enum.Enum):
        __new__: ClassVar[Callable] = ...
        _generate_next_value_: ClassVar[Callable] = ...
        _hashable_values_: ClassVar[list] = ...
        _member_map_: ClassVar[dict] = ...
        _member_names_: ClassVar[list] = ...
        _member_type_: ClassVar[type[object]] = ...
        _unhashable_values_: ClassVar[list] = ...
        _unhashable_values_map_: ClassVar[dict] = ...
        _use_args_: ClassVar[bool] = ...
        _value2member_map_: ClassVar[dict] = ...
        _value_repr_: ClassVar[None] = ...
        kGravityCompensation: ClassVar[InverseDynamics.InverseDynamicsMode] = ...
        kInverseDynamics: ClassVar[InverseDynamics.InverseDynamicsMode] = ...
        __nb_enum__: ClassVar[PyCapsule] = ...
    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 = ...) -> 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, ki, kd, has_reference_acceleration: bool, plant_context: pydrake.systems.framework.Context | None = ...) -> 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) -> None: ...

class JointStiffnessController(pydrake.systems.framework.LeafSystem):
    def __init__(self, plant: pydrake.multibody.plant.MultibodyPlant, kp, kd) -> 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: Incomplete
    low: Incomplete
    state_index: Incomplete
    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, ki, kd, state_output_port_index: int = ..., plant_input_port_index: int = ...) -> None: ...
    @overload
    def __init__(self, plant: pydrake.systems.framework.System, feedback_selector, 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, feedback_selector, kp, ki, kd, 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, ki, kd) -> None: ...
    @overload
    def __init__(self, state_projection, kp, ki, kd) -> None: ...
    @overload
    def __init__(self, state_projection, output_projection, kp, ki, kd) -> None: ...
    def get_Kd_vector(self, *args, **kwargs): ...
    def get_Ki_vector(self, *args, **kwargs): ...
    def get_Kp_vector(self, *args, **kwargs): ...
    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: ...
