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

from __future__ import annotations

import kcsolve

from .constraints import (
    CoincidentConstraint,
    ConstraintBase,
    DistancePointPointConstraint,
    FixedConstraint,
)
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

# Map BaseJointKind enum values to handler names
_SUPPORTED = {
    kcsolve.BaseJointKind.Coincident,
    kcsolve.BaseJointKind.DistancePointPoint,
    kcsolve.BaseJointKind.Fixed,
}


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
        all_residuals = []
        constraint_objs = []

        for c in 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)
            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
        all_residuals = substitution_pass(all_residuals, params)
        all_residuals = single_equation_pass(all_residuals, params)

        # 5. Newton-Raphson
        converged = newton_solve(
            all_residuals,
            params,
            quat_groups=quat_groups,
            max_iter=100,
            tol=1e-10,
        )

        # 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 _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."""
    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:
        marker_i_quat = tuple(marker_i.quaternion)
        marker_j_quat = tuple(marker_j.quaternion)
        return FixedConstraint(
            body_i,
            marker_i_pos,
            marker_i_quat,
            body_j,
            marker_j_pos,
            marker_j_quat,
        )

    return None