#!/usr/bin/env python3
"""Solve gem-slot translations from where you want the gem to actually appear.

Placement cannot be dialled in by hand because the gem model is `item/generated`:
a 2x2 sprite block sitting at the dead centre (8,8) of a 16x16 sheet, so its
geometry is 8px from the model origin on every axis. MIAPI's toMatrix is
T * Rx * Ry * Rz * S with no centring offset, so the rotation pivots on that
origin and flings the gem up to ~12px away. Translation is therefore NOT the
gem's position - it is whatever cancels that throw.

    world = pivot + flip( T + R*S*c ),   flip = negate x,y,   c = (8,8,8)

`flip` is the block-model(+y up) -> ModelPart(+y down) conversion. Inverting:

    T = flip(D) - R*S*c

with D the desired gem centre as an offset from the body part's pivot.

Edit TARGET below (in world coordinates, +x = wearer's LEFT, +y = DOWN,
-z = forward) and run. Rotation and scale are read from the JSON and preserved.
"""
import math, json, collections, os

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
BASE = os.path.join(REPO, "src/main/resources/packs/armory/data/cmmodular"
                          "/miapi/modules")

# vanilla HumanoidModel pivots
PIVOT = {"arm_left": (5, 2, 0), "arm_right": (-5, 2, 0),
         "leg_left": (1.9, 12, 0), "leg_right": (-1.9, 12, 0),
         "life_support_socket": (0, 0, 0), "helmet_socket": (0, 0, 0)}

# Which datapack file holds each slot.
FILE = {n: f"armor/socket/{n}.json" for n in
        ("arm_left", "arm_right", "leg_left", "leg_right")}
FILE["life_support_socket"] = "armor/space/life_support_socket.json"
FILE["helmet_socket"] = "armor/space/helmet_socket.json"

# The helmet gem is the one slot that does NOT go through the z180 flip. It
# chains through the hat slot instead, which scales by 1.25 and lifts by
# -3.925 to sit the helmet on the skull rather than the neck:
#     world = (T + R*S*c) * 1.25 + (0, -3.925, 0)
# Solving it with the flip inverts y and drops the gem past the chin to the
# floor. Verified against two in-game observations: T_y -2.0 rendered just
# under the chin (predicted +3.58) and T_y +11.5 at floor height (+20.45).
HAT_SCALE, HAT_LIFT = 1.25, -3.925
NO_FLIP = {"helmet_socket"}

# Desired gem centre, in WORLD model pixels. Reference geometry:
#   arms  x +-3..7 (surface at +-7), shoulder y 2, hand y 14
#   legs  x 0..+-4 (surface at +-2), hip y 12, knee y 18, foot y 24
#   front face z -2; -z is forward
TARGET = {
    "arm_left":  ( 8.30,  2.00,  0.00),   # outer pauldron face, jutting ~1.3px proud
    "arm_right": (-8.30,  2.00,  0.00),
    "leg_left":  ( 1.90, 13.00, -2.50),   # just under the knee
    "leg_right": (-1.90, 13.00, -2.50),
    # pack occupies x +-2.80, y -1.36..5.91, z 2.00..4.00
    "life_support_socket": ( 0.00,  2.27,  4.50),   # centre of the rear face
    # dome is +-4.25*1.25 about the lift: y -9.24..+1.39, front face z -5.31
    "helmet_socket":       ( 0.00, -9.00, -5.80),   # upper edge of the front face
}

d = math.radians
def rx(a): c, s = math.cos(a), math.sin(a); return [[1,0,0],[0,c,-s],[0,s,c]]
def ry(a): c, s = math.cos(a), math.sin(a); return [[c,0,s],[0,1,0],[-s,0,c]]
def rz(a): c, s = math.cos(a), math.sin(a); return [[c,-s,0],[s,c,0],[0,0,1]]
def mul(A, B): return [[sum(A[i][k]*B[k][j] for k in range(3)) for j in range(3)] for i in range(3)]
def ap(M, v): return [sum(M[i][k]*v[k] for k in range(3)) for i in range(3)]
def rot(r): return mul(rx(d(r["x"])), mul(ry(d(r["y"])), rz(d(r["z"]))))
def flip(v): return [-v[0], -v[1], v[2]]

C = [8.0, 8.0, 8.0]

def main():
    for name, world in TARGET.items():
        path = os.path.join(BASE, FILE[name])
        doc = json.load(open(path), object_pairs_hook=collections.OrderedDict)
        t = doc["data"]["replace"]["slots"]["gem"]["transform"]
        piv = PIVOT[name]
        D = [world[i] - piv[i] for i in range(3)]
        S = [t["scale"][k] for k in "xyz"]
        throw = ap(rot(t["rotation"]), [C[i]*S[i] for i in range(3)])
        if name in NO_FLIP:
            lift = [0.0, HAT_LIFT, 0.0]
            T = [round((D[i] - lift[i]) / HAT_SCALE - throw[i], 4) + 0.0 for i in range(3)]
        else:
            T = [round(flip(D)[i] - throw[i], 4) + 0.0 for i in range(3)]

        old = dict(t["translation"])
        for i, k in enumerate("xyz"):
            t["translation"][k] = T[i]
        open(path, "w").write(json.dumps(doc, indent=2) + "\n")

        if name in NO_FLIP:
            landed = [piv[i] + (T[j] + throw[j]) * HAT_SCALE + [0.0, HAT_LIFT, 0.0][j]
                      for i, j in enumerate(range(3))]
        else:
            landed = [piv[i] + flip([T[j] + throw[j] for j in range(3)])[i] for i in range(3)]
        assert all(abs(landed[i] - world[i]) < 1e-3 for i in range(3)), f"{name} did not solve"
        print(f"{name:9} T {old['x']:+8.3f},{old['y']:+8.3f},{old['z']:+8.3f}"
              f"  ->  {T[0]:+8.3f},{T[1]:+8.3f},{T[2]:+8.3f}"
              f"   world ({landed[0]:+6.2f},{landed[1]:+6.2f},{landed[2]:+6.2f})")

if __name__ == "__main__":
    main()