#! /usr/bin/python3
# Last edited on 2024-07-27 00:35:52 by stolfi

# Test program for module {affine}

import affine
import trafo
import sys
import rn
import rmxn
from math import sqrt, sin, cos, floor, ceil, inf, nan, pi
    
def check_map(A, name,d,Adir,bdir,Ainv,binv,kinv):
  # Checks whether the elements of {A} are {Adir,bdir,Ainv,binv} and
  # its dimension is {d} or less.
  if name == None: name = "??"
  sys.stderr.write("checking %s...\n" % name)
  name_act = trafo.get_name(A)
  assert name_act == name, "wrong name"
    
  dim_act = affine.dimension(A)
  assert dim_act == d, "wrong dimension"
  kinv_act, Aobj_act = affine.unpack(A)
  assert kinv_act == kinv
  
  Adir_act = affine.lin_matrix(A)
  check_mats(name, "dir", Adir, Adir_act, d, 1.0e-10)
  
  bdir_act = affine.disp_vector(A)
  check_vecs(name, "dir", bdir, bdir_act, d, 1.0e-10)
  
  Ainv_act = affine.lin_matrix(affine.inv(A))
  check_mats(name, "inv", Ainv, Ainv_act, d, 1.0e-10)
  
  binv_act = affine.disp_vector(affine.inv(A))
  check_vecs(name, "inv", binv, binv_act, d, 1.0e-10)

def check_mats(name, mtag, M_exp, M_act, d, tol):
  if dist_mats(M_exp, M_act, d) > tol:
    sys.stderr.write("%s.%s_exp = %s\n" % (name, mtag, str(M_exp)))
    sys.stderr.write("%s.%s_act = %s\n" % (name, mtag, str(M_act)))
    assert False, "matrices don't match"
  # ----------------------------------------------------------------------
    
def check_vecs(name, vtag, v_exp, v_act, d, tol):
  if dist_vecs(v_exp, v_act, d) > tol:
    sys.stderr.write("%s.%s_exp = %s\n" % (name, vtag, str(v_exp)))
    sys.stderr.write("%s.%s_act = %s\n" % (name, vtag, str(v_act)))
    assert False, "vectors don't match"
  # ----------------------------------------------------------------------
  
def dist_mats(M, N, d):
  # Compares {d} by {d} matrices {M,N}. If either is {None},
  # uses the identity matrix.
  if M == None: M = rmxn.ident_matrix(d,d)
  if N == None: N = rmxn.ident_matrix(d,d)
  err = sqrt(rmxn.norm_sqr(rmxn.sub(M, N)))
  return err
  
def dist_vecs(u, v, d):
  # Compares {d}-vectors. If either is {None}, uses the all-zeros vector.
  if u == None: u = (0,)*d
  if v == None: v = (0,)*d
  err = rn.dist(u,v)
  return err
  

sys.stderr.write("--- testing {make,dimension,unpack,lin_matrix,disp_vector} ----------------\n")

name0 = "A0"
d0 = 3
theta0 = pi/3
st0 = sin(theta0)
ct0 = cos(theta0)
Adir0 = ((ct0, st0, 0), (-st0, ct0, 0), (0,0,-2))
bdir0 = (1, 2, 3)
Ainv0 = ((ct0, -st0, 0), (st0, ct0, 0), (0,0,-0.5))
binv0 = rmxn.map_row(rn.scale(-1, bdir0), Ainv0)
A0 = affine.make(name0, d0,Adir0,bdir0,Ainv0,binv0)
kinv0 = +1

check_map(A0, name0,d0,Adir0,bdir0,Ainv0,binv0,kinv0)

sys.stderr.write("--- testing {inv} ----------------\n")

A1 = affine.inv(A0)
name1 = name0 + "'"
kinv1 = -kinv0

check_map(A1, name1,d0,Ainv0,binv0,Adir0,bdir0,kinv1)

sys.stderr.write("--- testing {None} as identity ----------------\n")

name2 = "A2"
d2 = 3
Adir2 = None
bdir2 = (0,0,0)
Ainv2 = ((1,0,0),(0,1,0),(0,0,1))
binv2 = None
A2 = affine.make(name2, d2,Adir2,bdir2,Ainv2,binv2)
kinv2 = +1

check_map(A2, name2,d2,Ainv2,binv2,Adir2,bdir2,kinv2)

sys.stderr.write("--- testing {apply} ----------------\n")

p0 = (2,4,6)
q0 = affine.apply(p0,A0);
q0_exp = rn.add(rmxn.map_row(p0,Adir0),bdir0)
assert rn.dist(q0,q0_exp) < 1.0e-10

p1 = affine.apply(q0,A1);
assert rn.dist(p1,p0) < 1.0e-10
   
sys.stderr.write("--- testing {compose} ----------------\n")

A3 = affine.compose(A0, A0)

d3 = d0
name3 = "??"
Adir3 = rmxn.mul(Adir0,Adir0)
bdir3 = rn.add(rmxn.map_row(bdir0, Adir0), bdir0)
Ainv3 = rmxn.mul(Ainv0,Ainv0)
binv3 = rn.add(rmxn.map_row(binv0, Ainv0), binv0)
kinv3 = +1

check_map(A3, name3,d3,Adir3,bdir3,Ainv3,binv3,kinv3)

A4 = affine.compose(A0, A1)

d4 = d0
name4 = "??"
Adir4 = rmxn.ident_matrix(d4,d4)
bdir4 = (0, 0, 0)
Ainv4 = rmxn.ident_matrix(d4,d4)
binv4 = (0, 0, 0)
kinv4 = +1

check_map(A4, name4,d4,Adir4,bdir4,Ainv4,binv4,kinv4)