from __future__ import annotations

from dataclasses import dataclass
from math import pi

import numpy as np


@dataclass
class PDGaitConfig:
    frequency_hz: float = 0.9
    kp: float = 24.0
    kd: float = 2.5
    near_phase_offset: float = 0.0
    far_phase_offset: float = pi


class PDGaitController:
    """Phase-based crank PD controller for the tandem 4-leg linkage."""

    JOINTS = (
        "near_crank_joint",
        "far_crank_joint",
    )

    def __init__(self, model, data, config: PDGaitConfig | None = None):
        self.model = model
        self.data = data
        self.cfg = config or PDGaitConfig()

        self._joint_info = []
        for name in self.JOINTS:
            jid = model.joint(name).id
            actuator_name = "near_crank_motor" if name.startswith("near") else "far_crank_motor"
            aid = model.actuator(actuator_name).id
            qpos_adr = model.jnt_qposadr[jid]
            qvel_adr = model.jnt_dofadr[jid]
            ctrl_min, ctrl_max = model.actuator_ctrlrange[aid]
            self._joint_info.append((aid, qpos_adr, qvel_adr, ctrl_min, ctrl_max))

    def desired_positions(self, t: float) -> np.ndarray:
        w = 2.0 * pi * self.cfg.frequency_hz
        phase_n = w * t + self.cfg.near_phase_offset
        phase_f = phase_n + self.cfg.far_phase_offset

        q_des = np.array(
            [
                phase_n,
                phase_f,
            ],
            dtype=np.float64,
        )
        return q_des

    def step(self, t: float) -> None:
        q_des = self.desired_positions(t)
        for i, (aid, qpos_adr, qvel_adr, ctrl_min, ctrl_max) in enumerate(self._joint_info):
            q = self.data.qpos[qpos_adr]
            qd = self.data.qvel[qvel_adr]
            err = float(q_des[i] - q)
            tau = self.cfg.kp * err - self.cfg.kd * qd
            self.data.ctrl[aid] = np.clip(tau, ctrl_min, ctrl_max)