"""Nonlinear least-squares geometry constraint solver."""

from __future__ import annotations

import math
from collections import defaultdict

import numpy as np
from scipy.optimize import least_squares

from statement_to_tikz.ir import (
    AngleMeasure,
    CircleDef,
    Collinear,
    EqualAngle,
    EqualLength,
    GeometryIR,
    Intersection,
    Length,
    Midpoint,
    OnCircle,
    OnLine,
    Parallel,
    Perpendicular,
    SolveMode,
    SolvedScene,
)

DEFAULT_TOL = 1e-4


def _get(coords: dict[str, tuple[float, float]], name: str) -> np.ndarray:
    return np.asarray(coords[name], dtype=float)


def _vec(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return b - a


def _cross2(u: np.ndarray, v: np.ndarray) -> float:
    return float(u[0] * v[1] - u[1] * v[0])


def _dot(u: np.ndarray, v: np.ndarray) -> float:
    return float(np.dot(u, v))


def _norm(u: np.ndarray) -> float:
    return float(np.linalg.norm(u))


def _angle_cos(a: np.ndarray, v: np.ndarray, b: np.ndarray) -> float:
    u = a - v
    w = b - v
    nu, nw = _norm(u), _norm(w)
    if nu < 1e-12 or nw < 1e-12:
        return 1.0
    return _dot(u, w) / (nu * nw)


def _circle_geometry(
    ir: GeometryIR, coords: dict[str, tuple[float, float]], circle_id: str
) -> tuple[np.ndarray, float] | None:
    circ = next((c for c in ir.circles if c.id == circle_id), None)
    if circ is None:
        return None
    return _circle_center_radius_from_def(circ, coords)


def _circle_center_radius_from_def(
    circ: CircleDef, coords: dict[str, tuple[float, float]]
) -> tuple[np.ndarray, float] | None:
    if circ.center is not None and circ.center in coords:
        c = _get(coords, circ.center)
        if circ.radius is not None:
            return c, float(circ.radius)
        if circ.through:
            for p in circ.through:
                if p in coords:
                    return c, _norm(_get(coords, p) - c)
        return c, 1.0
    if circ.through and len(circ.through) >= 3:
        pts = [p for p in circ.through[:3] if p in coords]
        if len(pts) < 3:
            return None
        a, b, c = (_get(coords, p) for p in pts)
        ax, ay = a
        bx, by = b
        cx, cy = c
        d = 2 * (ax * (by - cy) + bx * (cy - ay) + cx * (ay - by))
        if abs(d) < 1e-12:
            return None
        ux = (
            (ax**2 + ay**2) * (by - cy)
            + (bx**2 + by**2) * (cy - ay)
            + (cx**2 + cy**2) * (ay - by)
        ) / d
        uy = (
            (ax**2 + ay**2) * (cx - bx)
            + (bx**2 + by**2) * (ax - cx)
            + (cx**2 + cy**2) * (bx - ax)
        ) / d
        center = np.array([ux, uy])
        return center, _norm(a - center)
    return None


def constraint_residuals(
    ir: GeometryIR, coords: dict[str, tuple[float, float]]
) -> list[float]:
    """Return list of scalar residuals for each constraint (may be multi-valued flattened)."""
    r: list[float] = []
    for c in ir.constraints:
        if isinstance(c, EqualLength):
            la = _norm(_get(coords, c.a2) - _get(coords, c.a1))
            lb = _norm(_get(coords, c.b2) - _get(coords, c.b1))
            r.append(la - lb)
        elif isinstance(c, Length):
            la = _norm(_get(coords, c.b) - _get(coords, c.a))
            r.append(la - c.value)
        elif isinstance(c, EqualAngle):
            cos1 = _angle_cos(
                _get(coords, c.a1), _get(coords, c.v1), _get(coords, c.b1)
            )
            cos2 = _angle_cos(
                _get(coords, c.a2), _get(coords, c.v2), _get(coords, c.b2)
            )
            r.append(cos1 - cos2)
        elif isinstance(c, AngleMeasure):
            cos_t = math.cos(math.radians(c.degrees))
            cos_m = _angle_cos(
                _get(coords, c.a), _get(coords, c.vertex), _get(coords, c.b)
            )
            r.append(cos_m - cos_t)
        elif isinstance(c, Perpendicular):
            u = _vec(_get(coords, c.a1), _get(coords, c.a2))
            v = _vec(_get(coords, c.b1), _get(coords, c.b2))
            r.append(_dot(u, v))
        elif isinstance(c, Parallel):
            u = _vec(_get(coords, c.a1), _get(coords, c.a2))
            v = _vec(_get(coords, c.b1), _get(coords, c.b2))
            r.append(_cross2(u, v))
        elif isinstance(c, OnLine):
            a, b, p = _get(coords, c.a), _get(coords, c.b), _get(coords, c.point)
            r.append(_cross2(b - a, p - a))
        elif isinstance(c, OnCircle):
            geom = _circle_geometry(ir, coords, c.circle)
            if geom is None:
                r.append(0.0)
            else:
                center, rad = geom
                r.append(_norm(_get(coords, c.point) - center) - rad)
        elif isinstance(c, Midpoint):
            m = _get(coords, c.point)
            mid = 0.5 * (_get(coords, c.a) + _get(coords, c.b))
            r.extend([float(m[0] - mid[0]), float(m[1] - mid[1])])
        elif isinstance(c, Collinear):
            pts = [_get(coords, p) for p in c.points]
            base = pts[1] - pts[0]
            for p in pts[2:]:
                r.append(_cross2(base, p - pts[0]))
        elif isinstance(c, Intersection):
            a, b = _get(coords, c.a), _get(coords, c.b)
            cc, d = _get(coords, c.c), _get(coords, c.d)
            p = _get(coords, c.point)
            r.append(_cross2(b - a, p - a))
            r.append(_cross2(d - cc, p - cc))
        else:
            raise TypeError(f"unknown constraint {type(c)}")

    for circ in ir.circles:
        if circ.center and circ.through and circ.center in coords:
            center = _get(coords, circ.center)
            if circ.radius is not None:
                rad = circ.radius
            else:
                refs = [p for p in circ.through if p in coords]
                if not refs:
                    continue
                rad = _norm(_get(coords, refs[0]) - center)
            for p in circ.through:
                if p in coords:
                    r.append(_norm(_get(coords, p) - center) - rad)
        elif circ.through and len(circ.through) >= 3:
            pts = [p for p in circ.through if p in coords]
            if len(pts) >= 3:
                geom = _circle_center_radius_from_def(circ, coords)
                if geom is not None:
                    center, rad = geom
                    for p in pts:
                        r.append(_norm(_get(coords, p) - center) - rad)

    return r


def evaluate_residuals(
    ir: GeometryIR, coords: dict[str, tuple[float, float]]
) -> list[float]:
    return [abs(x) for x in constraint_residuals(ir, coords)]


def _initial_coords(ir: GeometryIR) -> dict[str, tuple[float, float]]:
    coords: dict[str, tuple[float, float]] = {}
    n = len(ir.points)
    for i, p in enumerate(ir.points):
        if p.hint is not None:
            coords[p.name] = (float(p.hint[0]), float(p.hint[1]))
        else:
            ang = 2 * math.pi * i / max(n, 1) - math.pi / 2
            coords[p.name] = (math.cos(ang), math.sin(ang))
    return coords


def _gauge_point_names(ir: GeometryIR) -> tuple[str | None, str | None]:
    """Choose two points to fix: prefer a length-constrained segment as base."""
    names = ir.point_names()
    if not names:
        return None, None

    for c in ir.constraints:
        if isinstance(c, Length):
            return c.a, c.b

    for c in ir.constraints:
        if isinstance(c, Midpoint):
            return c.a, c.b

    if ir.segments:
        return ir.segments[0].a, ir.segments[0].b

    p0 = names[0]
    p1 = names[1] if len(names) > 1 else None
    return p0, p1


def _triangle_area_residuals(
    ir: GeometryIR, coords: dict[str, tuple[float, float]]
) -> list[float]:
    """Soft non-degeneracy: penalize near-zero area for 3-cycles of segments."""
    adj: dict[str, set[str]] = defaultdict(set)
    for seg in ir.segments:
        adj[seg.a].add(seg.b)
        adj[seg.b].add(seg.a)

    residuals: list[float] = []
    seen: set[tuple[str, str, str]] = set()
    for a, nbrs in adj.items():
        for b in nbrs:
            for c in nbrs:
                if b >= c:
                    continue
                if c not in adj[b]:
                    continue
                key = tuple(sorted((a, b, c)))
                if key in seen:
                    continue
                seen.add(key)
                if not all(p in coords for p in key):
                    continue
                pa, pb, pc = (_get(coords, p) for p in key)
                area2 = abs(_cross2(pb - pa, pc - pa))
                residuals.append(max(0.0, 0.5 - area2))
    return residuals


def _remap_init_to_gauge(
    init: dict[str, tuple[float, float]],
    p0: str,
    p1: str,
    gauged: dict[str, tuple[float, float]],
) -> dict[str, tuple[float, float]]:
    """Rigidly map hinted positions onto the gauged base segment."""
    if p0 not in init or p1 not in init:
        return gauged
    o0 = np.asarray(init[p0], dtype=float)
    o1 = np.asarray(init[p1], dtype=float)
    n0 = np.asarray(gauged[p0], dtype=float)
    n1 = np.asarray(gauged[p1], dtype=float)
    o_len = _norm(o1 - o0) or 1.0
    n_len = _norm(n1 - n0) or 1.0
    scale = n_len / o_len
    od = (o1 - o0) / o_len
    nd = (n1 - n0) / n_len
    ang = math.atan2(nd[1], nd[0]) - math.atan2(od[1], od[0])
    ca, sa = math.cos(ang), math.sin(ang)
    rot = np.array([[ca, -sa], [sa, ca]])
    remapped = dict(gauged)
    for n, xy in init.items():
        if n in (p0, p1):
            continue
        local = (np.asarray(xy, dtype=float) - o0) * scale
        remapped[n] = tuple(n0 + rot @ local)
    return remapped


def _degrees_of_freedom(ir: GeometryIR) -> tuple[int, int]:
    """Rough DOF count: 2*|points| - 3 (gauge) vs number of scalar residuals."""
    n_coords = 2 * len(ir.points)
    gauge = min(3, n_coords)
    probe = _initial_coords(ir)
    n_res = len(constraint_residuals(ir, probe))
    return n_coords - gauge, n_res


def solve_geometry(
    ir: GeometryIR,
    *,
    tol: float = DEFAULT_TOL,
    max_nfev: int = 2000,
) -> SolvedScene:
    """Solve for point coordinates. Returns exact mode if residuals within tol."""
    names = ir.point_names()
    if not names:
        return SolvedScene(
            ir=ir,
            coordinates={},
            mode=SolveMode.failed,
            max_residual=0.0,
            message="no points",
        )

    init = _initial_coords(ir)
    dof, n_res = _degrees_of_freedom(ir)
    if n_res == 0 or n_res < max(1, dof - 1):
        res = evaluate_residuals(ir, init)
        return SolvedScene(
            ir=ir,
            coordinates=init,
            mode=SolveMode.failed,
            max_residual=max(res) if res else 0.0,
            residuals=res,
            message=f"underconstrained (dof≈{dof}, residuals={n_res})",
        )

    p0, p1 = _gauge_point_names(ir)

    free_names: list[str] = [n for n in names if n != p0]

    def apply_gauge(partial: dict[str, tuple[float, float]]) -> dict[str, tuple[float, float]]:
        full = dict(partial)
        if p0:
            full[p0] = (0.0, 0.0)
        if p1 and p1 in full:
            x, _y = full[p1]
            if abs(x) < 1e-9:
                x = 1.0
            full[p1] = (abs(x), 0.0)
        elif p1:
            full[p1] = (1.0, 0.0)
        return full

    def pack_free(coords: dict[str, tuple[float, float]]) -> np.ndarray:
        vals: list[float] = []
        for n in free_names:
            x, y = coords[n]
            if n == p1:
                vals.append(x if abs(x) > 1e-9 else 1.0)
            else:
                vals.extend([x, y])
        return np.asarray(vals, dtype=float)

    def unpack_free(v: np.ndarray) -> dict[str, tuple[float, float]]:
        partial: dict[str, tuple[float, float]] = {}
        idx = 0
        for n in free_names:
            if n == p1:
                partial[n] = (float(v[idx]), 0.0)
                idx += 1
            else:
                partial[n] = (float(v[idx]), float(v[idx + 1]))
                idx += 2
        return apply_gauge(partial)

    def fun(v: np.ndarray) -> np.ndarray:
        coords = unpack_free(v)
        res = list(constraint_residuals(ir, coords))
        res.extend(_triangle_area_residuals(ir, coords))
        if not res:
            extras = []
            if p1 and p1 in coords:
                extras.append(coords[p1][0] - 1.0)
            return np.asarray(extras or [0.0], dtype=float)
        return np.asarray(res, dtype=float)

    gauged_init = apply_gauge(init)
    if p0 and p1:
        gauged_init = apply_gauge(_remap_init_to_gauge(init, p0, p1, gauged_init))

    x0 = pack_free(gauged_init)
    if x0.size == 0:
        coords = gauged_init
        res = evaluate_residuals(ir, coords)
        max_r = max(res) if res else 0.0
        mode = SolveMode.exact if max_r <= tol else SolveMode.failed
        return SolvedScene(
            ir=ir,
            coordinates=coords,
            mode=mode,
            max_residual=max_r,
            residuals=res,
            message="only gauge points",
        )

    try:
        result = least_squares(
            fun, x0, ftol=1e-12, xtol=1e-12, gtol=1e-12, max_nfev=max_nfev
        )
        coords = unpack_free(result.x)
    except Exception as exc:  # noqa: BLE001
        coords = gauged_init
        res = evaluate_residuals(ir, coords)
        return SolvedScene(
            ir=ir,
            coordinates=coords,
            mode=SolveMode.failed,
            max_residual=max(res) if res else float("inf"),
            residuals=res,
            message=f"solver error: {exc}",
        )

    res = evaluate_residuals(ir, coords)
    max_r = max(res) if res else 0.0
    if max_r <= tol:
        mode = SolveMode.exact
        msg = "constraints satisfied"
    else:
        mode = SolveMode.failed
        msg = f"residuals too large ({max_r:.4g})"

    return SolvedScene(
        ir=ir,
        coordinates={k: (float(v[0]), float(v[1])) for k, v in coords.items()},
        mode=mode,
        max_residual=float(max_r),
        residuals=res,
        message=msg,
    )


def scene_from_coordinates(
    ir: GeometryIR,
    coords: dict[str, tuple[float, float]],
    *,
    mode: SolveMode,
    message: str = "",
) -> SolvedScene:
    res = evaluate_residuals(ir, coords)
    max_r = max(res) if res else 0.0
    return SolvedScene(
        ir=ir,
        coordinates=coords,
        mode=mode,
        max_residual=float(max_r),
        residuals=res,
        message=message,
    )