# Path: deluca/agents/_lqr.py
# class LQR(Agent):
#     """
#     LQR
#     """
# 
#     # TODO: need to address problem of LQR with jax.lax.scan
# 
#     def __init__(
#         self, A: jnp.ndarray, B: jnp.ndarray, Q: jnp.ndarray = None, R: jnp.ndarray = None
#     ) -> None:
#         """
#         Description: Initialize the infinite-time horizon LQR.
#         Args:
#             A (jnp.ndarray): system dynamics
#             B (jnp.ndarray): system dynamics
#             Q (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
#             R (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
# 
#         Returns:
#             None
#         """
# 
#         state_size, action_size = B.shape
# 
#         if Q is None:
#             Q = jnp.identity(state_size, dtype=jnp.float32)
# 
#         if R is None:
#             R = jnp.identity(action_size, dtype=jnp.float32)
# 
#         # solve the ricatti equation
#         X = dare(A, B, Q, R)
# 
#         # compute LQR gain
#         self.K = jnp.linalg.inv(B.T @ X @ B + R) @ (B.T @ X @ A)
# 
#     def __call__(self, state) -> jnp.ndarray:
#         """
#         Description: Return the action based on current state and internal parameters.
# 
#         Args:
#             state (float/numpy.ndarray): current state
# 
#         Returns:
#            jnp.ndarray: action to take
#         """
#         return -self.K @ state


from numbers import Real
from deluca.agents._lqr import LQR
from deluca.agents.core import Agent
import jax.numpy as jnp
import numpy as np
import numpy.random as random
        """
        Description: Initialize the dynamics of the model.

        Args:
            A (jnp.ndarray): system dynamics
            B (jnp.ndarray): system dynamics
            Q (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
            R (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
            K (jnp.ndarray): Starting policy (optional). Defaults to LQR gain.
            start_time (int):
            H (postive int): history of the controller
            lr_scale (Real):
            decay (boolean):
        """

        self.d_state, self.d_action = B.shape  # State & Action Dimensions

        self.A, self.B = A, B  # System Dynamics

        self.t = 0  # Time Counter (for decaying learning rate)

        self.H = H

        self.lr_scale, self.decay = lr_scale, decay

        self.delta = delta

        # Model Parameters
        # initial linear policy / perturbation contributions / bias
        # TODO: need to address problem of LQR with jax.lax.scan

# Path: deluca/agents/_lqr.py
# class LQR(Agent):
#     """
#     LQR
#     """
# 
#     # TODO: need to address problem of LQR with jax.lax.scan
# 
#     def __init__(
#         self, A: jnp.ndarray, B: jnp.ndarray, Q: jnp.ndarray = None, R: jnp.ndarray = None
#     ) -> None:
#         """
#         Description: Initialize the infinite-time horizon LQR.
#         Args:
#             A (jnp.ndarray): system dynamics
#             B (jnp.ndarray): system dynamics
#             Q (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
#             R (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
# 
#         Returns:
#             None
#         """
# 
#         state_size, action_size = B.shape
# 
#         if Q is None:
#             Q = jnp.identity(state_size, dtype=jnp.float32)
# 
#         if R is None:
#             R = jnp.identity(action_size, dtype=jnp.float32)
# 
#         # solve the ricatti equation
#         X = dare(A, B, Q, R)
# 
#         # compute LQR gain
#         self.K = jnp.linalg.inv(B.T @ X @ B + R) @ (B.T @ X @ A)
# 
#     def __call__(self, state) -> jnp.ndarray:
#         """
#         Description: Return the action based on current state and internal parameters.
# 
#         Args:
#             state (float/numpy.ndarray): current state
# 
#         Returns:
#            jnp.ndarray: action to take
#         """
#         return -self.K @ state


# Path: deluca/agents/_bpc.py
from numbers import Real
from deluca.agents._lqr import LQR
from deluca.agents.core import Agent
import jax.numpy as jnp
import numpy as np
import numpy.random as random

        """
        Description: Initialize the dynamics of the model.

        Args:
            A (jnp.ndarray): system dynamics
            B (jnp.ndarray): system dynamics
            Q (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
            R (jnp.ndarray): cost matrices (i.e. cost = x^TQx + u^TRu)
            K (jnp.ndarray): Starting policy (optional). Defaults to LQR gain.
            start_time (int):
            H (postive int): history of the controller
            lr_scale (Real):
            decay (boolean):
        """

        self.d_state, self.d_action = B.shape  # State & Action Dimensions

        self.A, self.B = A, B  # System Dynamics

        self.t = 0  # Time Counter (for decaying learning rate)

        self.H = H

        self.lr_scale, self.decay = lr_scale, decay

        self.delta = delta

        # Model Parameters
        # initial linear policy / perturbation contributions / bias
        # TODO: need to address problem of LQR with jax.lax.scan

        self.K = K if K is not None else LQR(self.A, self.B, Q, R).K