class SimpleCartPole:
    """A simplified, dependency-free cart-pole environment."""

    def __init__(self):
        self.gravity, self.mass_cart, self.mass_pole = 9.8, 1.0, 0.1
        self.length, self.force_mag, self.dt = 0.5, 10.0, 0.02
        self.state_dim, self.n_actions = 4, 2

    def reset(self):
        self.state = np.random.uniform(-0.05, 0.05, size=4)
        self.steps = 0
        return self.state.copy()

    def step(self, action):
        x, x_dot, theta, theta_dot = self.state
        force = self.force_mag if action == 1 else -self.force_mag
        total_mass = self.mass_cart + self.mass_pole
        temp = (
            force + self.mass_pole * self.length * theta_dot ** 2
            * np.sin(theta)
        ) / total_mass
        theta_acc = (
            self.gravity * np.sin(theta) - np.cos(theta) * temp
        ) / (
            self.length * (
                4.0 / 3.0
                - self.mass_pole * np.cos(theta) ** 2 / total_mass
            )
        )
        pole_term = (
            self.mass_pole * self.length * theta_acc
            * np.cos(theta) / total_mass
        )
        x_acc = temp - pole_term

        x += self.dt * x_dot
        x_dot += self.dt * x_acc
        theta += self.dt * theta_dot
        theta_dot += self.dt * theta_acc
        self.state = np.array([x, x_dot, theta, theta_dot])
        self.steps += 1

        done = bool(
            abs(x) > 2.4 or abs(theta) > 0.21 or self.steps >= 200
        )
        reward = 1.0
        return self.state.copy(), reward, done


if __name__ == "__main__":
    env = SimpleCartPole()
    online, returns = train_dqn(
        env, state_dim=4, n_actions=2, episodes=150,
    )
    print("Average return, first 10 episodes:",
          np.mean(returns[:10]))
    print("Average return, last 10 episodes: ",
          np.mean(returns[-10:]))
