"""CartPole in plain numpy (no gymnasium needed).

Same physics and termination rules as the classic Barto, Sutton & Anderson (1983)
task used by Gymnasium's CartPole-v1:

    state  = (x, x_dot, theta, theta_dot)      cart position/velocity, pole angle/angular velocity
    action = 0 (push left) or 1 (push right)    force of magnitude 10 N
    reward = +1 for every step the pole stays up
    done   = |theta| > 12 deg, |x| > 2.4, or 500 steps

Usage:
    env = CartPole(seed=0)
    s = env.reset()
    s, r, done = env.step(1)
"""

import math

import numpy as np


class CartPole:
    gravity = 9.8
    masscart = 1.0
    masspole = 0.1
    length = 0.5  # half the pole's length
    force_mag = 10.0
    tau = 0.02  # seconds between state updates (explicit Euler)
    theta_threshold = 12 * 2 * math.pi / 360
    x_threshold = 2.4

    n_actions = 2
    obs_dim = 4

    def __init__(self, seed=None, max_steps=500):
        self.rng = np.random.default_rng(seed)
        self.max_steps = max_steps
        self.state = None
        self.t = 0

    def reset(self):
        self.state = self.rng.uniform(-0.05, 0.05, size=4)
        self.t = 0
        return self.state.astype(np.float32).copy()

    def step(self, action):
        assert self.state is not None, "call reset() first"
        x, x_dot, theta, theta_dot = self.state
        force = self.force_mag if action == 1 else -self.force_mag
        total_mass = self.masspole + self.masscart
        polemass_length = self.masspole * self.length
        cos, sin = math.cos(theta), math.sin(theta)

        temp = (force + polemass_length * theta_dot**2 * sin) / total_mass
        theta_acc = (self.gravity * sin - cos * temp) / (
            self.length * (4.0 / 3.0 - self.masspole * cos**2 / total_mass)
        )
        x_acc = temp - polemass_length * theta_acc * cos / total_mass

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

        fell = abs(x) > self.x_threshold or abs(theta) > self.theta_threshold
        done = bool(fell or self.t >= self.max_steps)
        return self.state.astype(np.float32).copy(), 1.0, done