"""The player model and our pose maths, ported from the library so poses can be rendered in Python.

Mirrors core's LimbBend and Bezier and the fabric mod's PoseApplier: same part boxes and pivots, same
rotation order (ModelPart's Z, then Y, then X), the same bends, and the same carrying of the head and
arms by a bent torso. check_port.py compares it against the Java classes.

Model space is Minecraft's: pixels, Y down, the player faces -Z, and its right side is at -X.
"""
from dataclasses import dataclass

import numpy as np

BONES = ["root", "head", "body", "right_arm", "left_arm", "right_leg", "left_leg"]

# Default pivots and boxes (from, size) of the vanilla player model with wide arms.
PIVOTS = {
    "head": (0, 0, 0),
    "body": (0, 0, 0),
    "right_arm": (-5, 2, 0),
    "left_arm": (5, 2, 0),
    "right_leg": (-1.9, 12, 0),
    "left_leg": (1.9, 12, 0),
}
BOXES = {
    "head": ((-4, -8, -4), (8, 8, 8)),
    "body": ((-4, 0, -2), (8, 12, 4)),
    "right_arm": ((-3, -2, -2), (4, 12, 4)),
    "left_arm": ((-1, -2, -2), (4, 12, 4)),
    "right_leg": ((-2, 0, -2), (4, 12, 4)),
    "left_leg": ((-2, 0, -2), (4, 12, 4)),
}
# The outer skin layer (hat, jacket, sleeves, pants) is this much bigger on every side.
OVERLAY = {"head": 0.5, "body": 0.25, "right_arm": 0.25, "left_arm": 0.25, "right_leg": 0.25, "left_leg": 0.25}
# Where each part's base and outer layer start in a 64x64 skin.
UV = {
    "head": ((0, 0), (32, 0)),
    "body": ((16, 16), (16, 32)),
    "right_arm": ((40, 16), (40, 32)),
    "left_arm": ((32, 48), (48, 48)),
    "right_leg": ((0, 16), (0, 32)),
    "left_leg": ((16, 48), (0, 48)),
}

SHARPNESS = 0.5
SEGMENTS = 12
ARC_SAMPLES = 64
ROOT_PIVOT_Y = 12.8


@dataclass
class Transform:
    """BoneTransform: offset in pixels, rotations and bend in radians."""
    x: float = 0
    y: float = 0
    z: float = 0
    pitch: float = 0
    yaw: float = 0
    roll: float = 0
    bend: float = 0
    axis: float = 0


def rotation_zyx(pitch, yaw, roll):
    """The matrix ModelPart.translateAndRotate applies: Rz(roll) Ry(yaw) Rx(pitch)."""
    cx, sx = np.cos(pitch), np.sin(pitch)
    cy, sy = np.cos(yaw), np.sin(yaw)
    cz, sz = np.cos(roll), np.sin(roll)
    rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
    ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
    rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])
    return rz @ ry @ rx


def axis_angle(axis, angle):
    x, y, z = axis
    c, s = np.cos(angle), np.sin(angle)
    k = np.array([[0, -z, y], [z, 0, -x], [-y, x, 0]])
    return np.eye(3) + s * k + (1 - c) * (k @ k)


def _bezier(points, t):
    """Points on a cubic Bézier at each t (an array), like core's Bezier.point."""
    t = np.asarray(t, float)[:, None]
    p0, p1, p2, p3 = points
    u = 1 - t
    return u ** 3 * p0 + 3 * u * u * t * p1 + 3 * u * t * t * p2 + t ** 3 * p3


def _bezier_derivative(points, t):
    t = np.asarray(t, float)[:, None]
    p0, p1, p2, p3 = points
    u = 1 - t
    return 3 * u * u * (p1 - p0) + 6 * u * t * (p2 - p1) + 3 * t * t * (p3 - p2)


class LimbBend:
    """Port of core's LimbBend. Vectorised over points."""

    def __init__(self, start, end, angle, axis, sharpness=SHARPNESS, segments=SEGMENTS, anchored_at_end=False):
        self.start = start
        self.length = end - start
        self.anchored_at_end = anchored_at_end
        self.segments = segments
        self.angle = angle
        self.dir = np.array([np.sin(axis), 0, -np.cos(axis)])
        self.n = np.array([self.dir[2], 0, -self.dir[0]])
        self.ring = np.zeros((segments + 1, 3))
        self.ring_angle = np.zeros(segments + 1)
        if self.straight:
            self.ring[:, 1] = start + self.length * np.arange(segments + 1) / segments
            return
        self._build(self._spine(min(1, max(0, sharpness))))

    @property
    def straight(self):
        return abs(self.angle) < 1e-4

    def _spine(self, sharpness):
        half = self.length / 2
        top = np.array([0, self.start, 0.0])
        joint = np.array([0, self.start + half, 0.0])
        lower = np.array([self.dir[0] * np.sin(self.angle), np.cos(self.angle), self.dir[2] * np.sin(self.angle)])
        bottom = joint + lower * half
        pull = 2 / 3 + sharpness / 3
        return [top, top + (joint - top) * pull, bottom + (joint - bottom) * pull, bottom]

    def _build(self, curve):
        ts = np.arange(ARC_SAMPLES + 1) / ARC_SAMPLES
        pts = _bezier(curve, ts)
        lengths = np.concatenate([[0], np.cumsum(np.linalg.norm(np.diff(pts, axis=0), axis=1))])
        total = lengths[-1]
        scale = self.length / total
        top = pts[0]
        # For each ring, the arc-length sample it falls in (LimbBend's while loop), then t within it.
        target = total * np.arange(self.segments + 1) / self.segments
        sample = np.clip(np.searchsorted(lengths, target, side="left") - 1, 0, ARC_SAMPLES - 1)
        span = lengths[sample + 1] - lengths[sample]
        f = np.where(span > 0, (target - lengths[sample]) / np.where(span > 0, span, 1), 0)
        t = ts[sample] + (ts[sample + 1] - ts[sample]) * f
        self.ring = top + (_bezier(curve, t) - top) * scale
        tangent = _bezier_derivative(curve, t)
        self.ring_angle = np.arctan2(tangent @ self.dir, tangent[:, 1])

    def _mirror(self, y):
        return 2 * self.start + self.length - y

    def _rotate(self, v, a):
        """Rodrigues' rotation of rows of v around n by per-row angles a."""
        n = self.n
        cos, sin = np.cos(a)[:, None], np.sin(a)[:, None]
        cross = np.cross(np.broadcast_to(n, v.shape), v)
        dot = (v @ n)[:, None]
        return v * cos + cross * sin + n * dot * (1 - cos)

    def _from_start(self, p):
        along = (p[:, 1] - self.start) / self.length * self.segments
        ring = np.clip(np.floor(along).astype(int), 0, self.segments - 1)
        f = np.clip(along - ring, 0, 1)
        extra = np.where(along <= 0, p[:, 1] - self.start,
                         np.where(along >= self.segments, p[:, 1] - (self.start + self.length), 0))
        f = np.where(along <= 0, 0, np.where(along >= self.segments, 1, f))
        center = self.ring[ring] + (self.ring[ring + 1] - self.ring[ring]) * f[:, None]
        a = self.ring_angle[ring] + (self.ring_angle[ring + 1] - self.ring_angle[ring]) * f
        local = np.stack([p[:, 0], extra, p[:, 2]], axis=1)
        return self._rotate(local, a) + center

    def bend_points(self, p):
        if self.anchored_at_end:
            q = p.copy()
            q[:, 1] = self._mirror(q[:, 1])
            out = self._from_start(q)
            out[:, 1] = self._mirror(out[:, 1])
            return out
        return self._from_start(p)

    def _angle_from_start(self, y):
        along = np.clip((y - self.start) / self.length * self.segments, 0, self.segments)
        ring = min(int(along), self.segments - 1)
        return self.ring_angle[ring] + (self.ring_angle[ring + 1] - self.ring_angle[ring]) * (along - ring)

    def angle_at(self, y):
        return -self._angle_from_start(self._mirror(y)) if self.anchored_at_end else self._angle_from_start(y)


def limb_bend(bone, t):
    if t is None or t.bend == 0:
        return None
    if bone in ("right_arm", "left_arm"):
        return LimbBend(-2, 10, t.bend, t.axis)
    if bone in ("right_leg", "left_leg"):
        return LimbBend(0, 12, t.bend, t.axis)
    if bone == "body":
        return LimbBend(0, 12, t.bend, t.axis, anchored_at_end=True)
    return None


def textured_points(bone, skin, spacing=0.5):
    """Points on a part's base and outer layer with their skin colours, in the part's own space.

    Follows ModelPart.Cube: each face's corners and UV rectangle, and Polygon's corner-to-UV order.
    Outer layer points are only kept where the skin has pixels there.
    """
    (fx, fy, fz), (dx, dy, dz) = BOXES[bone]
    pts, cols = [], []
    for layer, (u, v) in enumerate(UV[bone]):
        g = OVERLAY[bone] if layer else 0
        f, gg, h = fx - g, fy - g, fz - g
        i, j, k = fx + dx + g, fy + dy + g, fz + dz + g
        c = [(f, gg, h), (i, gg, h), (i, j, h), (f, j, h), (f, gg, k), (i, gg, k), (i, j, k), (f, j, k)]
        l, m, n, o = u, u + dz, u + dz + dx, u + dz + dx + dx
        pp, q = u + dz + dx + dz, u + dz + dx + dz + dx
        r, ss, t = v, v + dz, v + dz + dy
        faces = [  # corners, (u1, v1, u2, v2)
            ((5, 4, 0, 1), (m, r, n, ss)),   # DOWN: minY, the top in model space
            ((2, 3, 7, 6), (n, ss, o, r)),   # UP: maxY, the bottom
            ((0, 4, 7, 3), (l, ss, m, t)),   # WEST: -X, the player's right
            ((1, 0, 3, 2), (m, ss, n, t)),   # NORTH: -Z, the front
            ((5, 1, 2, 6), (n, ss, pp, t)),  # EAST: +X, the player's left
            ((4, 5, 6, 7), (pp, ss, q, t)),  # SOUTH: +Z, the back
        ]
        for corners, (u1, v1, u2, v2) in faces:
            a, b, _, d = (np.array(c[x], float) for x in corners)
            ns = max(2, int(np.ceil(np.linalg.norm(b - a) / spacing)) + 1)
            nt = max(2, int(np.ceil(np.linalg.norm(d - a) / spacing)) + 1)
            # Sample cell centres so every point falls inside one texture pixel.
            sv, tv = np.meshgrid((np.arange(ns) + 0.5) / ns, (np.arange(nt) + 0.5) / nt, indexing="ij")
            sv, tv = sv.ravel(), tv.ravel()
            face = a + sv[:, None] * (b - a) + tv[:, None] * (d - a)
            tu = np.clip(np.floor(u2 + sv * (u1 - u2)).astype(int), min(u1, u2), max(u1, u2) - 1)
            tvv = np.clip(np.floor(v1 + tv * (v2 - v1)).astype(int), min(v1, v2), max(v1, v2) - 1)
            rgba = skin[tvv, tu]
            keep = rgba[:, 3] > 0 if layer else np.ones(len(face), bool)
            pts.append(face[keep])
            cols.append(rgba[keep, :3])
    return np.concatenate(pts), np.concatenate(cols).astype(float) / 255


def surface_points(bone, spacing=0.5, overlay=True):
    """Points on the surface of a part's box (outer layer included), in the part's own space."""
    (fx, fy, fz), (sx, sy, sz) = BOXES[bone]
    g = OVERLAY[bone] if overlay else 0
    lo = np.array([fx - g, fy - g, fz - g])
    hi = np.array([fx + sx + g, fy + sy + g, fz + sz + g])
    axes = [np.linspace(lo[i], hi[i], max(2, int(np.ceil((hi[i] - lo[i]) / spacing)) + 1)) for i in range(3)]
    pts = []
    for fixed in range(3):
        a, b = [i for i in range(3) if i != fixed]
        u, v = np.meshgrid(axes[a], axes[b], indexing="ij")
        for side in (lo[fixed], hi[fixed]):
            face = np.empty((u.size, 3))
            face[:, a], face[:, b], face[:, fixed] = u.ravel(), v.ravel(), side
            pts.append(face)
    return np.unique(np.round(np.concatenate(pts), 4), axis=0)


class Model:
    """Poses the player's parts. pose() returns model-space points for each part.

    With a skin (an RGBA array), the points carry the skin's colours in self.colors, and the outer layer
    only exists where the skin has pixels. Without one, parts are plain boxes including the outer layer.
    """

    def __init__(self, spacing=0.5, skin=None):
        if skin is None:
            self.local = {bone: surface_points(bone, spacing) for bone in PIVOTS}
            self.colors = None
        else:
            textured = {bone: textured_points(bone, skin, spacing) for bone in PIVOTS}
            self.local = {bone: pts for bone, (pts, _) in textured.items()}
            self.colors = np.concatenate([cols for _, cols in textured.values()])

    def part_frames(self, pose):
        """Each part's pivot, rotation and bend after PoseApplier.apply, in root space."""
        frames = {}
        for bone in PIVOTS:
            t = pose.get(bone)
            pivot = np.array(PIVOTS[bone], float)
            rot = np.eye(3)
            if t is not None:
                pivot = pivot + [t.x, t.y, t.z]
                rot = rotation_zyx(t.pitch, t.yaw, t.roll)
            frames[bone] = [pivot, rot, limb_bend(bone, t)]

        body_pivot, body_rot, body_bend = frames["body"]
        if body_bend is not None and not body_bend.straight:
            for bone in ("head", "right_arm", "left_arm"):
                pivot, rot, bend = frames[bone]
                local = body_rot.T @ (pivot - body_pivot)
                moved = body_bend.bend_points(local[None, :])[0]
                turn = axis_angle(body_bend.n, body_bend.angle_at(local[1]))
                frames[bone] = [body_pivot + body_rot @ moved, body_rot @ turn @ body_rot.T @ rot, bend]
        return frames

    def pose(self, pose):
        frames = self.part_frames(pose)
        root = pose.get("root")
        if root is not None:
            r_rot = rotation_zyx(root.pitch, root.yaw, root.roll)
            pivot = np.array([0, ROOT_PIVOT_Y, 0])
            r_pos = np.array([root.x, root.y, root.z]) + pivot - r_rot @ pivot
        else:
            r_rot, r_pos = np.eye(3), np.zeros(3)
        out = {}
        for bone, (pivot, rot, bend) in frames.items():
            p = self.local[bone]
            if bend is not None and not bend.straight:
                p = bend.bend_points(p)
            out[bone] = (p @ rot.T + pivot) @ r_rot.T + r_pos
        return out


# --- The rig the fitter works in -------------------------------------------------------------------
#
# Moving every pivot freely lets parts drift apart, so the fitter uses a jointed rig instead and
# converts it to BoneTransforms: the legs hang from the hips, the torso's hip end sits on the hips,
# and the head and arms hang from the torso (PoseApplier then carries them along its bend).

RIG = (
    ["root_x", "root_y", "root_z", "root_pitch", "root_yaw", "root_roll", "hip_x", "hip_y", "hip_z"]
    + [f"body_{c}" for c in ("pitch", "yaw", "roll", "bend", "axis")]
    + [f"head_{c}" for c in ("pitch", "yaw", "roll")]
    + [f"{limb}_{c}" for limb in ("right_arm", "left_arm", "right_leg", "left_leg")
       for c in ("pitch", "yaw", "roll", "bend", "axis")]
)
INDEX = {name: i for i, name in enumerate(RIG)}


def rig_to_pose(v):
    g = lambda name: float(v[INDEX[name]])
    hip = np.array([g("hip_x"), g("hip_y"), g("hip_z")])
    body_rot = rotation_zyx(g("body_pitch"), g("body_yaw"), g("body_roll"))
    # The torso's bottom middle (0, 12, 0) in its own space stays on the hips; a bend keeps it fixed.
    hips_rest = np.array([0, 12, 0.0])
    body_pivot = hips_rest + hip - body_rot @ hips_rest
    pose = {
        "root": Transform(g("root_x"), g("root_y"), g("root_z"), g("root_pitch"), g("root_yaw"), g("root_roll")),
        "body": Transform(*body_pivot, g("body_pitch"), g("body_yaw"), g("body_roll"), g("body_bend"), g("body_axis")),
        "head": Transform(*body_pivot, g("head_pitch"), g("head_yaw"), g("head_roll")),
    }
    for arm in ("right_arm", "left_arm"):
        rest = np.array(PIVOTS[arm], float)
        # Where the shoulder sits on the (unbent) tipped torso; PoseApplier adds the bend's carry.
        offset = body_pivot + body_rot @ rest - rest
        pose[arm] = Transform(*offset, *(g(f"{arm}_{c}") for c in ("pitch", "yaw", "roll", "bend", "axis")))
    for leg in ("right_leg", "left_leg"):
        pose[leg] = Transform(*hip, *(g(f"{leg}_{c}") for c in ("pitch", "yaw", "roll", "bend", "axis")))
    return pose


def rest_rig():
    return np.zeros(len(RIG))