from __future__ import annotations

from dataclasses import replace
from time import perf_counter
from typing import Iterable, List

import numpy as np
import osqp
import pinocchio as pin
from scipy import sparse

from .kinematics import RM75Kinematics, pose_errors, validate_se3
from .types import IkOptions, IkResult, IkStatus, JointLimits


class RM75IkSolver:
    """Single-seed differential IK solved with a reused OSQP workspace."""

    def __init__(self, kinematics: RM75Kinematics) -> None:
        self.kinematics = kinematics
        self.model = kinematics.model
        self.data = kinematics.data
        self.frame_id = kinematics.flange_frame_id
        self._n = 7

        pattern = sparse.triu(np.ones((self._n, self._n)), format="csc")
        self._p_rows = pattern.indices.copy()
        self._p_cols = np.repeat(np.arange(self._n), np.diff(pattern.indptr))
        constraints = sparse.eye(self._n, format="csc")
        self._osqp = osqp.OSQP()
        self._osqp.setup(
            P=pattern,
            q=np.zeros(self._n),
            A=constraints,
            l=-np.ones(self._n),
            u=np.ones(self._n),
            verbose=False,
            warm_start=True,
            polish=False,
            eps_abs=1e-6,
            eps_rel=1e-6,
            max_iter=1000,
        )

    def solve(
        self,
        target_se3: pin.SE3,
        seed_rad: np.ndarray,
        options: IkOptions = IkOptions(),
    ) -> IkResult:
        started = perf_counter()
        try:
            validate_se3(target_se3, "target_se3")
            q = self.kinematics.validate_q(seed_rad)
        except (TypeError, ValueError) as exc:
            return IkResult(
                IkStatus.INVALID_INPUT,
                None,
                float("inf"),
                float("inf"),
                0,
                perf_counter() - started,
                message=str(exc),
            )

        q_reference = q.copy()
        weights = np.diag(np.asarray(options.task_weights, dtype=float))
        damping = options.damping_initial
        previous_error = float("inf")
        best_error = float("inf")
        stagnant_iterations = 0
        last_osqp_status = ""
        position_error = float("inf")
        orientation_error = float("inf")

        for iteration in range(options.max_iterations + 1):
            elapsed = perf_counter() - started
            if options.time_limit_sec is not None and elapsed >= options.time_limit_sec:
                return IkResult(
                    IkStatus.TIME_LIMIT,
                    None,
                    position_error,
                    orientation_error,
                    iteration,
                    elapsed,
                    last_osqp_status,
                    "IK time budget exhausted",
                )

            pin.computeJointJacobians(self.model, self.data, q)
            pin.framesForwardKinematics(self.model, self.data, q)
            current = self.data.oMf[self.frame_id]
            position_error, orientation_error = pose_errors(current, target_se3)
            if (
                position_error <= options.position_tolerance_m
                and orientation_error <= options.orientation_tolerance_rad
            ):
                solution = q.copy()
                solution.setflags(write=False)
                return IkResult(
                    IkStatus.SUCCESS,
                    solution,
                    position_error,
                    orientation_error,
                    iteration,
                    perf_counter() - started,
                    last_osqp_status,
                )

            if iteration == options.max_iterations:
                break

            error_transform = current.actInv(target_se3)
            error_vector = pin.log6(error_transform).vector
            error_norm = float(np.linalg.norm(error_vector))
            if error_norm < best_error - options.stagnation_delta:
                best_error = error_norm
                stagnant_iterations = 0
            else:
                stagnant_iterations += 1
                if stagnant_iterations >= options.stagnation_iterations:
                    return IkResult(
                        IkStatus.STAGNATED,
                        None,
                        position_error,
                        orientation_error,
                        iteration,
                        perf_counter() - started,
                        last_osqp_status,
                        "SE(3) error stopped improving",
                    )

            if error_norm > previous_error * 1.1 and iteration > 10:
                damping = min(options.damping_max, damping * 1.5)
            else:
                damping = max(options.damping_min, damping * options.damping_reduction)

            jacobian = pin.getFrameJacobian(
                self.model,
                self.data,
                self.frame_id,
                pin.ReferenceFrame.LOCAL,
            )
            effective_jacobian = pin.Jlog6(error_transform) @ jacobian
            hessian = effective_jacobian.T @ weights @ effective_jacobian
            hessian += (
                damping * damping + options.posture_weight
            ) * np.eye(self._n)
            gradient = -effective_jacobian.T @ weights @ error_vector
            gradient += options.posture_weight * (q - q_reference)

            lower = np.maximum(
                -options.trust_region_rad,
                self.kinematics.limits.lower - q,
            )
            upper = np.minimum(
                options.trust_region_rad,
                self.kinematics.limits.upper - q,
            )
            p_values = hessian[self._p_rows, self._p_cols]
            self._osqp.update(Px=p_values, q=gradient, l=lower, u=upper)
            osqp_result = self._osqp.solve()
            last_osqp_status = str(osqp_result.info.status)
            if last_osqp_status.lower() != "solved" or osqp_result.x is None:
                return IkResult(
                    IkStatus.OSQP_FAILURE,
                    None,
                    position_error,
                    orientation_error,
                    iteration,
                    perf_counter() - started,
                    last_osqp_status,
                    "OSQP did not return a solved step",
                )
            step = np.asarray(osqp_result.x, dtype=float)
            if step.shape != (7,) or not np.all(np.isfinite(step)):
                return IkResult(
                    IkStatus.OSQP_FAILURE,
                    None,
                    position_error,
                    orientation_error,
                    iteration,
                    perf_counter() - started,
                    last_osqp_status,
                    "OSQP returned a non-finite step",
                )
            q = pin.integrate(self.model, q, step)
            q = np.clip(
                q,
                self.kinematics.limits.lower,
                self.kinematics.limits.upper,
            )
            previous_error = error_norm

        return IkResult(
            IkStatus.MAX_ITERATIONS,
            None,
            position_error,
            orientation_error,
            options.max_iterations,
            perf_counter() - started,
            last_osqp_status,
            "maximum IK iterations reached",
        )

    def solve_multistart(
        self,
        target_se3: pin.SE3,
        seeds_rad: Iterable[np.ndarray],
        options: IkOptions = IkOptions(),
    ) -> IkResult:
        started = perf_counter()
        last_result: IkResult | None = None
        for index, seed in enumerate(seeds_rad, start=1):
            result = self.solve(target_se3, seed, options)
            if result.success:
                return replace(
                    result,
                    solve_time_sec=perf_counter() - started,
                    message=f"converged from recovery seed {index}",
                )
            last_result = result
        if last_result is None:
            return IkResult(
                IkStatus.INVALID_INPUT,
                None,
                float("inf"),
                float("inf"),
                0,
                perf_counter() - started,
                message="no recovery seeds were provided",
            )
        return replace(
            last_result,
            solve_time_sec=perf_counter() - started,
            message=f"all recovery seeds failed; last error: {last_result.message}",
        )


def deterministic_recovery_seeds(
    limits: JointLimits,
    count: int = 8,
    random_seed: int = 75,
) -> List[np.ndarray]:
    if count <= 0:
        raise ValueError("recovery seed count must be positive")
    seeds = [np.clip(np.zeros(7), limits.lower, limits.upper)]
    rng = np.random.default_rng(random_seed)
    while len(seeds) < count:
        seeds.append(rng.uniform(limits.lower, limits.upper))
    return seeds

