import dataclasses
import math

import numpy as np
import torch
from gymnasium import spaces

from src.core.env import BaseEnv, BaseEnvConfig


@dataclasses.dataclass
class RotaryCartPoleState:
    motor_angle: torch.Tensor     # (num_envs,)
    motor_vel: torch.Tensor       # (num_envs,)
    pendulum_angle: torch.Tensor  # (num_envs,)
    pendulum_vel: torch.Tensor    # (num_envs,)


@dataclasses.dataclass
class RotaryCartPoleConfig(BaseEnvConfig):
    """Rotary inverted pendulum (Furuta pendulum) task config.

    The motor rotates the arm horizontally; the pendulum swings freely
    at the arm tip.  Goal: swing the pendulum up and balance it upright.
    """
    # Reward shaping
    reward_upright_scale: float = 1.0      # upright reward ∈ [0, scale]
    alive_bonus: float = 0.25              # per-step survival bonus (must stay alive > die)
    balance_bonus: float = 2.0             # extra reward for upright AND still (beats spinning)
    balance_vel_scale: float = 0.5         # decay rate of the bonus with pendulum speed
    motor_vel_penalty: float = 0.01        # penalise high motor angular velocity
    motor_angle_penalty: float = 0.05      # penalise deviation from centre
    action_penalty: float = 0.05           # penalise large actions (energy cost)
    action_rate_penalty: float = 0.01      # penalise action changes (smoothness —
                                           # critical with ~100 ms real motor lag)

    # ── Initial state randomisation ──
    pendulum_init_range_deg: float = 180.0 # pendulum starts in [-range, +range]

    # ── Software safety limit (env-level, on top of URDF + hardware) ──
    motor_angle_limit_deg: float = 90.0    # terminate episode if exceeded


class RotaryCartPoleEnv(BaseEnv[RotaryCartPoleConfig]):
    """Furuta pendulum / rotary inverted pendulum environment.

    Kinematic chain:  base_link ─(motor_joint, z)─► arm ─(pendulum_joint, y)─► pendulum

    Observations (6):
        [sin(motor), cos(motor), sin(pendulum), cos(pendulum), motor_vel, pendulum_vel]
    Using sin/cos avoids discontinuities at ±π for continuous joints.

    Actions (1):
        Torque applied to the motor_joint (normalised to [-1, 1]).
    """

    def __init__(self, config: RotaryCartPoleConfig):
        super().__init__(config)

    # ── Spaces ───────────────────────────────────────────────────

    @property
    def observation_space(self) -> spaces.Space:
        return spaces.Box(low=-torch.inf, high=torch.inf, shape=(6,))

    @property
    def action_space(self) -> spaces.Space:
        return spaces.Box(low=-1.0, high=1.0, shape=(1,))

    # ── State building ───────────────────────────────────────────

    def build_state(self, qpos: torch.Tensor, qvel: torch.Tensor) -> RotaryCartPoleState:
        return RotaryCartPoleState(
            motor_angle=qpos[:, 0],
            motor_vel=qvel[:, 0],
            pendulum_angle=qpos[:, 1],
            pendulum_vel=qvel[:, 1],
        )

    # ── Observations ─────────────────────────────────────────────

    def compute_observations(self, state: RotaryCartPoleState) -> torch.Tensor:
        obs = [
            torch.sin(state.motor_angle),
            torch.cos(state.motor_angle),
            torch.sin(state.pendulum_angle),
            torch.cos(state.pendulum_angle),
            state.motor_vel,
            state.pendulum_vel,
        ]
        return torch.stack(obs, dim=-1)

    # ── Rewards ──────────────────────────────────────────────────

    def compute_rewards(
        self,
        state: RotaryCartPoleState,
        actions: torch.Tensor,
        prev_actions: torch.Tensor,
    ) -> torch.Tensor:
        # Upright shaping ∈ [0, 1]: 0 hanging down (θ=0), 1 fully upright (θ=π).
        # Non-negative by design so *surviving* always beats ending the episode early
        # (otherwise the optimum is to slam the arm into the ±limit — "suicide policy").
        upright = 0.5 * (1.0 - torch.cos(state.pendulum_angle))

        # Balanced bonus: large ONLY when near the top AND nearly still.  A freely
        # spinning pendulum passes through the top at high speed, so stillness≈0 and
        # it earns ~none of this — making true balancing strictly dominate the
        # "just keep spinning in full loops" local optimum.
        stillness = torch.exp(-self.config.balance_vel_scale * state.pendulum_vel.pow(2))
        balance = self.config.balance_bonus * upright * stillness

        # Per-step alive bonus keeps a not-yet-upright step net-positive despite
        # penalties, so the −10 termination penalty is genuinely a deterrent.
        reward = (upright * self.config.reward_upright_scale
                  + balance
                  + self.config.alive_bonus)

        # Penalise fast motor spinning (discourages violent oscillation)
        reward = reward - self.config.motor_vel_penalty * state.motor_vel.pow(2)

        # Penalise motor deviation from centre (keep arm near zero)
        reward = reward - self.config.motor_angle_penalty * state.motor_angle.pow(2)

        # Penalise large actions (energy efficiency / smoother control)
        reward = reward - self.config.action_penalty * actions.squeeze(-1).pow(2)

        # Penalise rapid action changes — a jittery policy is unrealisable
        # through the real motor's ~100 ms lag and excites unmodeled dynamics.
        action_rate = (actions - prev_actions).squeeze(-1).pow(2)
        reward = reward - self.config.action_rate_penalty * action_rate

        # Penalty for exceeding motor angle limit (episode also terminates)
        limit_rad = math.radians(self.config.motor_angle_limit_deg)
        exceeded = state.motor_angle.abs() >= limit_rad
        reward = torch.where(exceeded, torch.tensor(-10.0, device=reward.device), reward)

        return reward

    # ── Initial state randomization ──────────────────────────────

    def initial_state_ranges(
        self, nq: int, nv: int,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """Small motor perturbation; wide pendulum angle (swing-up task)."""
        qpos_lo = np.full(nq, -0.05)
        qpos_hi = np.full(nq, 0.05)
        qvel_lo = np.full(nv, -0.05)
        qvel_hi = np.full(nv, 0.05)
        pend_range = math.radians(self.config.pendulum_init_range_deg)
        if pend_range > 0 and nq >= 2:
            qpos_lo[1] = -pend_range
            qpos_hi[1] = pend_range
        return qpos_lo, qpos_hi, qvel_lo, qvel_hi

    # ── Terminations ─────────────────────────────────────────────

    def compute_terminations(self, state: RotaryCartPoleState) -> torch.Tensor:
        # Software safety: terminate if motor angle exceeds limit.
        limit_rad = math.radians(self.config.motor_angle_limit_deg)
        exceeded = state.motor_angle.abs() >= limit_rad
        return exceeded

    # ── Reset readiness (for SerialRunner) ───────────────────────

    def is_reset_ready(self, qpos: torch.Tensor, qvel: torch.Tensor) -> bool:
        """Pendulum must be hanging still and motor near center."""
        motor_angle = float(qpos[0, 0])
        pend_angle = float(qpos[0, 1])
        motor_vel = float(qvel[0, 0])
        pend_vel = float(qvel[0, 1])

        # Pendulum near hanging-down (angle ~0) and slow
        angle_ok = abs(pend_angle) < math.radians(2.0)
        vel_ok = abs(pend_vel) < math.radians(5.0)
        # Motor near center and slow
        motor_ok = abs(motor_angle) < math.radians(5.0)
        motor_vel_ok = abs(motor_vel) < math.radians(10.0)
        return angle_ok and vel_ok and motor_ok and motor_vel_ok


