"""Geometric entities that own solver parameters.

Phase 1 provides RigidBody — a 7-DOF entity (position + unit quaternion).
"""

from __future__ import annotations

from .expr import Const, Expr
from .params import ParamTable
from .quat import world_point as _world_point


class RigidBody:
    """A rigid body with 7 parameters: tx, ty, tz, qw, qx, qy, qz.

    Grounded bodies have all parameters marked fixed.
    """

    def __init__(
        self,
        part_id: str,
        params: ParamTable,
        position: tuple[float, float, float],
        quaternion: tuple[float, float, float, float],
        grounded: bool = False,
    ):
        self.part_id = part_id
        self.grounded = grounded

        pfx = part_id + "/"
        self.tx = params.add(pfx + "tx", position[0], fixed=grounded)
        self.ty = params.add(pfx + "ty", position[1], fixed=grounded)
        self.tz = params.add(pfx + "tz", position[2], fixed=grounded)
        # quaternion: (w, x, y, z) — matches KCSolve convention
        self.qw = params.add(pfx + "qw", quaternion[0], fixed=grounded)
        self.qx = params.add(pfx + "qx", quaternion[1], fixed=grounded)
        self.qy = params.add(pfx + "qy", quaternion[2], fixed=grounded)
        self.qz = params.add(pfx + "qz", quaternion[3], fixed=grounded)

        self._param_names = [
            pfx + "tx",
            pfx + "ty",
            pfx + "tz",
            pfx + "qw",
            pfx + "qx",
            pfx + "qy",
            pfx + "qz",
        ]

    def world_point(self, lx: float, ly: float, lz: float) -> tuple[Expr, Expr, Expr]:
        """Transform a local-frame point to world-frame Expr triple."""
        return _world_point(
            self.tx,
            self.ty,
            self.tz,
            self.qw,
            self.qx,
            self.qy,
            self.qz,
            lx,
            ly,
            lz,
        )

    def quat_norm_residual(self) -> Expr:
        """qw^2 + qx^2 + qy^2 + qz^2 - 1.  Added as a constraint."""
        return (
            self.qw * self.qw
            + self.qx * self.qx
            + self.qy * self.qy
            + self.qz * self.qz
            - Const(1.0)
        )

    def quat_param_names(self) -> tuple[str, str, str, str]:
        """Return the 4 quaternion parameter names (for renormalization)."""
        pfx = self.part_id + "/"
        return pfx + "qw", pfx + "qx", pfx + "qy", pfx + "qz"

    def extract_position(self, env: dict[str, float]) -> tuple[float, float, float]:
        pfx = self.part_id + "/"
        return env[pfx + "tx"], env[pfx + "ty"], env[pfx + "tz"]

    def extract_quaternion(
        self, env: dict[str, float]
    ) -> tuple[float, float, float, float]:
        pfx = self.part_id + "/"
        return env[pfx + "qw"], env[pfx + "qx"], env[pfx + "qy"], env[pfx + "qz"]
