"""EFC geometry, contact row assembly and force projection (standard library)."""
import json
from pathlib import Path


def model_data():
    source = (Path(__file__).resolve().parents[1]/"assets/contact-model.js").read_text()
    return json.loads(source.split("=", 1)[1].strip().removesuffix(";"))


def dot(a, b):
    return sum(x*y for x, y in zip(a, b))


def matvec(matrix, vector):
    return [dot(row, vector) for row in matrix]


def transpose_product(matrix, force):
    return [sum(row[j]*f for row, f in zip(matrix, force))
            for j in range(len(matrix[0]))]


def contact_state(height, radius, margin, gap):
    distance = height-radius
    return {"distance": distance, "error": distance-margin,
            "candidate": distance <= margin+gap, "active": distance < margin}


def contact_rows(condim=3, pyramidal=False, friction=None):
    """Rows for the translation-only sphere: generalized velocity is (vx,vy,vz)."""
    if condim not in (1, 3, 4, 6):
        raise ValueError("condim must be 1, 3, 4 or 6")
    physical = [[0, 0, 1], [1, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 0], [0, 0, 0]][:condim]
    if condim == 1 or not pyramidal:
        return physical
    friction = friction or [0.6, 0.6, 0.01, 0.02, 0.02]
    return [[n+s*friction[k-1]*t for n, t in zip(physical[0], physical[k])]
            for k in range(1, condim) for s in (1, -1)]


def inverse_inertia(rows, inverse_mass):
    """A=J M^-1 J^T, inverse_mass is a diagonal for the teaching model."""
    return [[sum(a*b*m for a, b, m in zip(ri, rj, inverse_mass))
             for rj in rows] for ri in rows]


def pyramid_force(edge_force, friction):
    physical = [sum(edge_force)]
    for k, mu in enumerate(friction):
        physical.append(mu*(edge_force[2*k]-edge_force[2*k+1]))
    return physical


if __name__ == "__main__":
    c = model_data()
    rows = contact_rows()
    force = [20, -4, 1]
    print(contact_state(c["position"][2], c["radius"], c["margin"], c["gap"]))
    print("Jv =", matvec(rows, c["velocity"]))
    print("A =", inverse_inertia(rows, [1/c["mass"]]*3))
    generalized = transpose_product(rows, force)
    print("J^T f =", generalized)
    print("a =", [g+f/c["mass"] for g, f in zip(c["gravity"], generalized)])