"""KindredSolver — IKCSolver implementation bridging KCSolve to the
expression-based Newton-Raphson solver."""

from __future__ import annotations

import kcsolve

from .bfgs import bfgs_solve
from .constraints import (
    AngleConstraint,
    BallConstraint,
    CamConstraint,
    CoincidentConstraint,
    ConcentricConstraint,
    ConstraintBase,
    CylindricalConstraint,
    DistanceCylSphConstraint,
    DistancePointPointConstraint,
    FixedConstraint,
    GearConstraint,
    LineInPlaneConstraint,
    ParallelConstraint,
    PerpendicularConstraint,
    PlanarConstraint,
    PointInPlaneConstraint,
    PointOnLineConstraint,
    RackPinionConstraint,
    RevoluteConstraint,
    ScrewConstraint,
    SliderConstraint,
    SlotConstraint,
    TangentConstraint,
    UniversalConstraint,
)
from .decompose import decompose, solve_decomposed
from .dof import count_dof
from .entities import RigidBody
from .newton import newton_solve
from .params import ParamTable
from .prepass import single_equation_pass, substitution_pass

# Assemblies with fewer free bodies than this use the monolithic path.
_DECOMPOSE_THRESHOLD = 8

# All BaseJointKind values this solver can handle
_SUPPORTED = {
    # Phase 1
    kcsolve.BaseJointKind.Coincident,
    kcsolve.BaseJointKind.DistancePointPoint,
    kcsolve.BaseJointKind.Fixed,
    # Phase 2: point constraints
    kcsolve.BaseJointKind.PointOnLine,
    kcsolve.BaseJointKind.PointInPlane,
    # Phase 2: orientation
    kcsolve.BaseJointKind.Parallel,
    kcsolve.BaseJointKind.Perpendicular,
    kcsolve.BaseJointKind.Angle,
    # Phase 2: axis/surface
    kcsolve.BaseJointKind.Concentric,
    kcsolve.BaseJointKind.Tangent,
    kcsolve.BaseJointKind.Planar,
    kcsolve.BaseJointKind.LineInPlane,
    # Phase 2: kinematic joints
    kcsolve.BaseJointKind.Ball,
    kcsolve.BaseJointKind.Revolute,
    kcsolve.BaseJointKind.Cylindrical,
    kcsolve.BaseJointKind.Slider,
    kcsolve.BaseJointKind.Screw,
    kcsolve.BaseJointKind.Universal,
    # Phase 2: mechanical
    kcsolve.BaseJointKind.Gear,
    kcsolve.BaseJointKind.RackPinion,
}


class KindredSolver(kcsolve.IKCSolver):
    """Expression-based Newton-Raphson constraint solver."""

    def name(self):
        return "Kindred (Newton-Raphson)"

    def supported_joints(self):
        return list(_SUPPORTED)

    def solve(self, ctx):
        params = ParamTable()
        bodies = {}  # part_id -> RigidBody

        # 1. Build entities from parts
        for part in ctx.parts:
            pos = tuple(part.placement.position)
            quat = tuple(part.placement.quaternion)  # (w, x, y, z)
            body = RigidBody(
                part.id,
                params,
                position=pos,
                quaternion=quat,
                grounded=part.grounded,
            )
            bodies[part.id] = body

        # 2. Build constraint residuals (track index mapping for decomposition)
        all_residuals = []
        constraint_objs = []
        constraint_indices = []  # parallel to constraint_objs: index in ctx.constraints

        for idx, c in enumerate(ctx.constraints):
            if not c.activated:
                continue
            body_i = bodies.get(c.part_i)
            body_j = bodies.get(c.part_j)
            if body_i is None or body_j is None:
                continue

            marker_i_pos = tuple(c.marker_i.position)
            marker_j_pos = tuple(c.marker_j.position)

            obj = _build_constraint(
                c.type,
                body_i,
                marker_i_pos,
                body_j,
                marker_j_pos,
                c.marker_i,
                c.marker_j,
                c.params,
            )
            if obj is None:
                continue
            constraint_objs.append(obj)
            constraint_indices.append(idx)
            all_residuals.extend(obj.residuals())

        # 3. Add quaternion normalization residuals for non-grounded bodies
        quat_groups = []
        for body in bodies.values():
            if not body.grounded:
                all_residuals.append(body.quat_norm_residual())
                quat_groups.append(body.quat_param_names())

        # 4. Pre-passes on full system
        all_residuals = substitution_pass(all_residuals, params)
        all_residuals = single_equation_pass(all_residuals, params)

        # 5. Solve (decomposed for large assemblies, monolithic for small)
        n_free_bodies = sum(1 for b in bodies.values() if not b.grounded)
        if n_free_bodies >= _DECOMPOSE_THRESHOLD:
            grounded_ids = {pid for pid, b in bodies.items() if b.grounded}
            clusters = decompose(ctx.constraints, grounded_ids)
            if len(clusters) > 1:
                converged = solve_decomposed(
                    clusters,
                    bodies,
                    constraint_objs,
                    constraint_indices,
                    params,
                )
            else:
                converged = _monolithic_solve(
                    all_residuals,
                    params,
                    quat_groups,
                )
        else:
            converged = _monolithic_solve(all_residuals, params, quat_groups)

        # 6. DOF
        dof = count_dof(all_residuals, params)

        # 7. Build result
        result = kcsolve.SolveResult()
        result.status = (
            kcsolve.SolveStatus.Success if converged else kcsolve.SolveStatus.Failed
        )
        result.dof = dof

        env = params.get_env()
        placements = []
        for body in bodies.values():
            if body.grounded:
                continue
            pr = kcsolve.SolveResult.PartResult()
            pr.id = body.part_id
            pr.placement = kcsolve.Transform()
            pr.placement.position = list(body.extract_position(env))
            pr.placement.quaternion = list(body.extract_quaternion(env))
            placements.append(pr)

        result.placements = placements
        return result

    def is_deterministic(self):
        return True


def _monolithic_solve(all_residuals, params, quat_groups):
    """Newton-Raphson solve with BFGS fallback on the full system."""
    converged = newton_solve(
        all_residuals,
        params,
        quat_groups=quat_groups,
        max_iter=100,
        tol=1e-10,
    )
    if not converged:
        converged = bfgs_solve(
            all_residuals,
            params,
            quat_groups=quat_groups,
            max_iter=200,
            tol=1e-10,
        )
    return converged


def _build_constraint(
    kind,
    body_i,
    marker_i_pos,
    body_j,
    marker_j_pos,
    marker_i,
    marker_j,
    c_params,
) -> ConstraintBase | None:
    """Create the appropriate constraint object from a BaseJointKind."""
    marker_i_quat = tuple(marker_i.quaternion)
    marker_j_quat = tuple(marker_j.quaternion)

    # -- Phase 1 constraints --------------------------------------------------

    if kind == kcsolve.BaseJointKind.Coincident:
        return CoincidentConstraint(body_i, marker_i_pos, body_j, marker_j_pos)

    if kind == kcsolve.BaseJointKind.DistancePointPoint:
        distance = c_params[0] if c_params else 0.0
        return DistancePointPointConstraint(
            body_i,
            marker_i_pos,
            body_j,
            marker_j_pos,
            distance,
        )

    if kind == kcsolve.BaseJointKind.Fixed:
        return FixedConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    # -- Phase 2: point constraints -------------------------------------------

    if kind == kcsolve.BaseJointKind.PointOnLine:
        return PointOnLineConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    if kind == kcsolve.BaseJointKind.PointInPlane:
        offset = c_params[0] if c_params else 0.0
        return PointInPlaneConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
            offset=offset,
        )

    # -- Phase 2: orientation constraints -------------------------------------

    if kind == kcsolve.BaseJointKind.Parallel:
        return ParallelConstraint(body_i, marker_i_quat, body_j, marker_j_quat)

    if kind == kcsolve.BaseJointKind.Perpendicular:
        return PerpendicularConstraint(body_i, marker_i_quat, body_j, marker_j_quat)

    if kind == kcsolve.BaseJointKind.Angle:
        angle = c_params[0] if c_params else 0.0
        return AngleConstraint(body_i, marker_i_quat, body_j, marker_j_quat, angle)

    # -- Phase 2: axis/surface constraints ------------------------------------

    if kind == kcsolve.BaseJointKind.Concentric:
        distance = c_params[0] if c_params else 0.0
        return ConcentricConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
            distance=distance,
        )

    if kind == kcsolve.BaseJointKind.Tangent:
        return TangentConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    if kind == kcsolve.BaseJointKind.Planar:
        offset = c_params[0] if c_params else 0.0
        return PlanarConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
            offset=offset,
        )

    if kind == kcsolve.BaseJointKind.LineInPlane:
        offset = c_params[0] if c_params else 0.0
        return LineInPlaneConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
            offset=offset,
        )

    # -- Phase 2: kinematic joints --------------------------------------------

    if kind == kcsolve.BaseJointKind.Ball:
        return BallConstraint(body_i, marker_i_pos, body_j, marker_j_pos)

    if kind == kcsolve.BaseJointKind.Revolute:
        return RevoluteConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    if kind == kcsolve.BaseJointKind.Cylindrical:
        return CylindricalConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    if kind == kcsolve.BaseJointKind.Slider:
        return SliderConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    if kind == kcsolve.BaseJointKind.Screw:
        pitch = c_params[0] if c_params else 1.0
        return ScrewConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
            pitch=pitch,
        )

    if kind == kcsolve.BaseJointKind.Universal:
        return UniversalConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    # -- Phase 2: mechanical constraints --------------------------------------

    if kind == kcsolve.BaseJointKind.Gear:
        radius_i = c_params[0] if len(c_params) > 0 else 1.0
        radius_j = c_params[1] if len(c_params) > 1 else 1.0
        return GearConstraint(
            body_i,
            marker_i_quat,
            body_j,
            marker_j_quat,
            radius_i,
            radius_j,
        )

    if kind == kcsolve.BaseJointKind.RackPinion:
        pitch_radius = c_params[0] if c_params else 1.0
        return RackPinionConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
            pitch_radius=pitch_radius,
        )

    # -- Stubs (accepted but produce no residuals) ----------------------------

    if kind == kcsolve.BaseJointKind.Cam:
        return CamConstraint()

    if kind == kcsolve.BaseJointKind.Slot:
        return SlotConstraint()

    if kind == kcsolve.BaseJointKind.DistanceCylSph:
        return DistanceCylSphConstraint()

    return None