"""Rolling gait: contract cables (lower ctrl rest length), measure CoM motion."""
from __future__ import annotations

import itertools
from dataclasses import dataclass, field

import mujoco
import numpy as np

from .model import build_model, cable_pairs, node_positions


@dataclass
class Robot:
    stiffness: float = 300.0
    slope_deg: float = 0.0
    prestretch: float = 0.9
    friction: float = 1.0
    model: mujoco.MjModel = field(init=False)
    data: mujoco.MjData = field(init=False)
    nominal_rest: np.ndarray = field(init=False)
    frame_every: int = 25  # record one frame every N physics steps
    frames: list | None = field(default=None, init=False)  # set to [] to record node positions


    def __post_init__(self) -> None:
        self.model = build_model(stiffness=self.stiffness, slope_deg=self.slope_deg, friction=self.friction)
        self.data = mujoco.MjData(self.model)
        self.reset()

    def reset(self) -> None:
        mujoco.mj_resetData(self.model, self.data)
        mujoco.mj_forward(self.model, self.data)
        self.nominal_rest = self.data.ten_length.copy() * self.prestretch
        self.data.ctrl[:] = self.nominal_rest
        self.run(1.5)

    def run(self, seconds: float) -> None:
        for index in range(int(seconds / self.model.opt.timestep)):
            mujoco.mj_step(self.model, self.data)
            if self.frames is not None and index % self.frame_every == 0:
                self.frames.append(self.node_positions().round(4).tolist())

    def node_positions(self) -> np.ndarray:
        """World xyz of the 12 bar-end nodes, in node index order."""
        ids = [mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, f"n{i}") for i in range(12)]
        return self.data.site_xpos[ids].copy()

    def cable_tension(self) -> np.ndarray:
        return -self.data.actuator_force.copy()

    def com(self) -> np.ndarray:
        return self.data.subtree_com[0].copy()

    def save(self) -> np.ndarray:
        size = mujoco.mj_stateSize(self.model, mujoco.mjtState.mjSTATE_INTEGRATION)
        state = np.empty(size)
        mujoco.mj_getState(self.model, self.data, state, mujoco.mjtState.mjSTATE_INTEGRATION)
        return state

    def load(self, state: np.ndarray) -> None:
        self.data.warning[mujoco.mjtWarning.mjWARN_BADQACC].number = 0
        mujoco.mj_setState(self.model, self.data, state, mujoco.mjtState.mjSTATE_INTEGRATION)
        mujoco.mj_forward(self.model, self.data)

    def step_action(self, cables: tuple[int, ...], ratio: float = 0.4,
                    hold: float = 1.0, release: float = 0.0) -> float:
        """Contract `cables` by `ratio` of rest length for `hold` s, then relax all cables
        for `release` s. With release=0 the contraction stays until the next action
        (the "hold" strategy, which keeps the robot from rolling back on slopes).
        Return CoM displacement along +x (uphill)."""
        start = self.com()
        self.data.ctrl[:] = self.nominal_rest
        self.data.ctrl[list(cables)] = self.nominal_rest[list(cables)] * (1 - ratio)
        self.run(hold)
        if release > 0:
            self.data.ctrl[:] = self.nominal_rest
            self.run(release)
        if self.unstable():
            return float("nan")
        return float(self.com()[0] - start[0])

    def unstable(self) -> bool:
        """MuJoCo silently resets data after a NaN; its warning counter tells us."""
        return bool(self.data.warning[mujoco.mjtWarning.mjWARN_BADQACC].number > 0)


def single_actions() -> list[tuple[int, ...]]:
    return [(k,) for k in range(24)]


def pair_actions() -> list[tuple[int, ...]]:
    """Pairs of cables that share a node (12 nodes x C(4,2) = 72)."""
    cables = cable_pairs(node_positions(1.0))
    return [(a, b) for a, b in itertools.combinations(range(24), 2) if set(cables[a]) & set(cables[b])]


def greedy_gait(robot: Robot, steps: int = 8, ratio: float = 0.4,
                candidates: list[tuple[int, ...]] | None = None, release: float = 0.0) -> list[dict]:
    """At every step try each candidate action from the same saved state, keep the best.

    This is a one-step lookahead search: it finds a good gait sequence without
    knowing rolling theory in advance, which is what the study wants to compare
    against hand-designed sequences."""
    if candidates is None:
        candidates = single_actions()
    log = []
    recording = robot.frames
    for step in range(steps):
        state = robot.save()
        scores = []
        robot.frames = None  # trial rollouts are not recorded
        for action in candidates:
            robot.load(state)
            scores.append(robot.step_action(action, ratio, release=release))
        best = int(np.nanargmax(scores))
        robot.load(state)
        robot.frames = recording
        dx = robot.step_action(candidates[best], ratio, release=release)
        log.append({"step": step, "action": list(candidates[best]), "dx": dx, "x": float(robot.com()[0])})
    return log