from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from math import radians
from typing import Optional, Tuple

import numpy as np


class IkStatus(str, Enum):
    SUCCESS = "success"
    INVALID_INPUT = "invalid_input"
    OSQP_FAILURE = "osqp_failure"
    STAGNATED = "stagnated"
    TIME_LIMIT = "time_limit"
    MAX_ITERATIONS = "max_iterations"


@dataclass(frozen=True)
class JointLimits:
    name: str
    lower: np.ndarray
    upper: np.ndarray

    def __post_init__(self) -> None:
        lower = np.asarray(self.lower, dtype=float).copy()
        upper = np.asarray(self.upper, dtype=float).copy()
        if lower.shape != (7,) or upper.shape != (7,):
            raise ValueError("RM75 joint limits must each have shape (7,)")
        if not np.all(np.isfinite(lower)) or not np.all(np.isfinite(upper)):
            raise ValueError("joint limits must be finite")
        if np.any(lower >= upper):
            raise ValueError("every lower joint limit must be below its upper limit")
        lower.setflags(write=False)
        upper.setflags(write=False)
        object.__setattr__(self, "lower", lower)
        object.__setattr__(self, "upper", upper)

    def contains(self, q: np.ndarray, tolerance: float = 1e-10) -> bool:
        values = np.asarray(q, dtype=float)
        return bool(
            values.shape == (7,)
            and np.all(values >= self.lower - tolerance)
            and np.all(values <= self.upper + tolerance)
        )


def physical_joint_limits() -> JointLimits:
    upper = np.deg2rad([178.0, 130.0, 178.0, 135.0, 178.0, 128.0, 360.0])
    return JointLimits("physical", -upper, upper)


def teleop_joint_limits() -> JointLimits:
    lower = np.deg2rad([-150.0, -110.0, -170.0, -130.0, -175.0, -125.0, -179.0])
    upper = np.deg2rad([150.0, 110.0, 170.0, 130.0, 175.0, 125.0, 179.0])
    return JointLimits("teleop", lower, upper)


def joint_limit_profile(name: str) -> JointLimits:
    profiles = {
        "physical": physical_joint_limits,
        "teleop": teleop_joint_limits,
    }
    try:
        return profiles[name]()
    except KeyError as exc:
        raise ValueError(f"unknown joint limit profile: {name!r}") from exc


@dataclass(frozen=True)
class IkOptions:
    position_tolerance_m: float = 1e-3
    orientation_tolerance_rad: float = radians(0.1)
    max_iterations: int = 500
    time_limit_sec: Optional[float] = None
    trust_region_rad: float = 0.05
    task_weights: Tuple[float, float, float, float, float, float] = (
        1.0,
        1.0,
        1.0,
        0.4,
        0.4,
        0.4,
    )
    posture_weight: float = 1e-5
    damping_initial: float = 0.1
    damping_min: float = 0.01
    damping_max: float = 1.0
    damping_reduction: float = 0.95
    stagnation_iterations: int = 40
    stagnation_delta: float = 1e-9

    def __post_init__(self) -> None:
        positive = {
            "position_tolerance_m": self.position_tolerance_m,
            "orientation_tolerance_rad": self.orientation_tolerance_rad,
            "trust_region_rad": self.trust_region_rad,
            "damping_initial": self.damping_initial,
            "damping_min": self.damping_min,
            "damping_max": self.damping_max,
        }
        for name, value in positive.items():
            if not np.isfinite(value) or value <= 0.0:
                raise ValueError(f"{name} must be finite and positive")
        if self.max_iterations <= 0 or self.stagnation_iterations <= 0:
            raise ValueError("iteration limits must be positive")
        if self.time_limit_sec is not None and self.time_limit_sec <= 0.0:
            raise ValueError("time_limit_sec must be positive when set")
        if len(self.task_weights) != 6 or any(weight <= 0.0 for weight in self.task_weights):
            raise ValueError("task_weights must contain six positive values")
        if self.posture_weight < 0.0 or not np.isfinite(self.posture_weight):
            raise ValueError("posture_weight must be finite and non-negative")
        if not self.damping_min <= self.damping_initial <= self.damping_max:
            raise ValueError("damping_initial must be within damping_min and damping_max")
        if not 0.0 < self.damping_reduction <= 1.0:
            raise ValueError("damping_reduction must be in (0, 1]")


@dataclass(frozen=True)
class IkResult:
    status: IkStatus
    q: Optional[np.ndarray]
    position_error_m: float
    orientation_error_rad: float
    iterations: int
    solve_time_sec: float
    osqp_status: str = ""
    message: str = ""

    @property
    def success(self) -> bool:
        return self.status is IkStatus.SUCCESS
