#!/usr/bin/env python3
"""
frames2d.py : 2D rigid transforms, the way Lecture 2 writes them.

    python3 frames2d.py          # the worked chain from the slides, the three bugs, and the inverse

A transform is a 3 x 3 homogeneous matrix [[R, t], [0, 0, 1]] stored as nested lists, so no packages are needed.
T_a_b maps coordinates in frame b to coordinates in frame a; chains compose left to right, inside subscripts cancel:
    p_map = T_map_base @ T_base_laser @ p_laser
Use it to check your Problem Set 1 answers.
"""
import math


def make(x, y, theta_deg):
    """The transform of a frame at (x, y), rotated by theta degrees."""
    c, s = math.cos(math.radians(theta_deg)), math.sin(math.radians(theta_deg))
    return [[c, -s, x], [s, c, y], [0.0, 0.0, 1.0]]


def compose(A, B):
    """A @ B: first B, then A."""
    return [[sum(A[i][k] * B[k][j] for k in range(3)) for j in range(3)] for i in range(3)]


def invert(T):
    """[[R^T, -R^T t], [0, 1]], not [[R^T, -t], [0, 1]]."""
    r00, r01, r10, r11, tx, ty = T[0][0], T[0][1], T[1][0], T[1][1], T[0][2], T[1][2]
    return [[r00, r10, -(r00 * tx + r10 * ty)], [r01, r11, -(r01 * tx + r11 * ty)], [0.0, 0.0, 1.0]]


def apply(T, p):
    """T @ (x, y, 1): a point, so the translation counts."""
    x, y = p
    return (T[0][0] * x + T[0][1] * y + T[0][2], T[1][0] * x + T[1][1] * y + T[1][2])


def pose(T):
    """(x, y, theta in degrees) of a transform."""
    return (T[0][2], T[1][2], math.degrees(math.atan2(T[1][0], T[0][0])))


def main():
    T_map_base, T_base_laser, p_laser = make(2.00, 1.50, 30.0), make(0.06, 0.00, 0.0), (1.200, 0.000)
    p_map = apply(compose(T_map_base, T_base_laser), p_laser)
    print('p_map = T_map_base @ T_base_laser @ p_laser = (%.3f, %.3f)' % p_map)
    wrong = {'forgot the rotation': apply(make(2.00, 1.50, 0.0), apply(T_base_laser, p_laser)),
             'flipped the yaw sign': apply(compose(make(2.00, 1.50, -30.0), T_base_laser), p_laser),
             'skipped the laser link': apply(T_map_base, p_laser)}
    for name, q in wrong.items():
        print('  %-24s (%.3f, %.3f), %3.0f cm off' % (name, q[0], q[1], 100 * math.dist(q, p_map)))
    inv = invert(T_map_base)
    print('T_base_map = invert(T_map_base): t = (%.2f, %.2f), theta = %.0f deg' % pose(inv))
    print('  negating t alone gives t = (%.2f, %.2f), %.0f cm off' % (-2.00, -1.50, 100 * math.dist((-2.00, -1.50), pose(inv)[:2])))


if __name__ == '__main__':
    main()