from __future__ import annotations

from math import radians
from typing import Iterable, List

import numpy as np
import pinocchio as pin

from .kinematics import validate_se3
from .solver import RM75IkSolver
from .types import IkOptions


DEMO_TRAJECTORIES = ("joint", "line", "arc", "orientation", "combined")


def se3_target_trajectory(
    start_pose: pin.SE3,
    target_pose: pin.SE3,
    points: int,
) -> List[pin.SE3]:
    """Interpolate translation linearly and orientation on SO(3)."""

    if points < 2:
        raise ValueError("trajectory must contain at least two points")
    validate_se3(start_pose, "start_pose")
    validate_se3(target_pose, "target_pose")
    rotation_delta = pin.log3(start_pose.rotation.T @ target_pose.rotation)
    targets = []
    for fraction in np.linspace(0.0, 1.0, points):
        translation = (
            (1.0 - fraction) * start_pose.translation
            + fraction * target_pose.translation
        )
        rotation = start_pose.rotation @ pin.exp3(fraction * rotation_delta)
        targets.append(pin.SE3(rotation, translation))
    return targets


def cartesian_demo_targets(
    kind: str,
    start_pose: pin.SE3,
    points: int = 180,
) -> List[pin.SE3]:
    if kind not in DEMO_TRAJECTORIES[1:]:
        raise ValueError(f"Cartesian trajectory must be one of {DEMO_TRAJECTORIES[1:]}")
    if points < 2:
        raise ValueError("trajectory must contain at least two points")
    validate_se3(start_pose, "start_pose")
    targets: List[pin.SE3] = []
    for fraction in np.linspace(0.0, 1.0, points):
        translation = start_pose.translation.copy()
        rotation = start_pose.rotation.copy()
        if kind == "line":
            translation += np.array([0.04 * fraction, 0.0, 0.0])
        elif kind == "arc":
            angle = 0.5 * np.pi * fraction
            translation += np.array(
                [0.03 * np.sin(angle), 0.03 * (1.0 - np.cos(angle)), 0.0]
            )
        elif kind == "orientation":
            rotation = rotation @ pin.rpy.rpyToMatrix(0.0, 0.0, radians(15) * fraction)
        elif kind == "combined":
            translation += np.array([0.035 * fraction, 0.015 * fraction, 0.0])
            rotation = rotation @ pin.rpy.rpyToMatrix(
                radians(8) * fraction,
                0.0,
                radians(12) * fraction,
            )
        targets.append(pin.SE3(rotation, translation))
    return targets


def joint_demo_trajectory(
    start_q_rad: np.ndarray,
    points: int = 180,
) -> np.ndarray:
    start = np.asarray(start_q_rad, dtype=float)
    if start.shape != (7,) or not np.all(np.isfinite(start)):
        raise ValueError("start_q_rad must be a finite shape-(7,) vector")
    if points < 2:
        raise ValueError("trajectory must contain at least two points")
    offset = np.deg2rad([12.0, -8.0, 10.0, 8.0, -6.0, 7.0, 10.0])
    blend = 0.5 - 0.5 * np.cos(np.linspace(0.0, np.pi, points))
    return start[None, :] + blend[:, None] * offset[None, :]


def solve_pose_trajectory(
    solver: RM75IkSolver,
    targets: Iterable[pin.SE3],
    seed_q_rad: np.ndarray,
    options: IkOptions = IkOptions(
        position_tolerance_m=9e-4,
        orientation_tolerance_rad=radians(0.09),
        max_iterations=200,
    ),
) -> np.ndarray:
    seed = np.asarray(seed_q_rad, dtype=float).copy()
    solutions = []
    for index, target in enumerate(targets):
        result = solver.solve(target, seed, options)
        if not result.success or result.q is None:
            raise RuntimeError(
                f"IK failed at trajectory point {index}: "
                f"{result.status.value} {result.message}"
            )
        seed = result.q.copy()
        solutions.append(seed)
    return np.asarray(solutions)
