#! /usr/bin/python3
# Last edited on 2020-03-08 18:56:20 by jstolfi 

# This program reads a file of Vicon Lock+ marker coordinates, then
# computes and writes a file with positions and orientations of the 
# various body parts.

# The input must have one data frame per line. Each line has 62 fields,
# comprising frame and sub-frame indices (integers) followed by the
# coordinate triplets (in mm) of {M=20} markers that were affixed to
# specific points of the subject's skin.

# An input comment "# sampling_freq ..." is copied to the output. Three
# CSV-like header lines are written after it. Other comment or blank
# lines in the input are discarded.

# The ouput will also have one data frame per line, that describes the
# position, orientation, and attitude of a matchstick human figure.

# For the representation of the latter, see
# {libnmwk/nmwk_compute_abs_matchstick_parms.py}. Here it suffices to
# know that it has {N} parts, and that the position and orientaton of
# each part is defined by a /placement/ data block with three coordinate
# triplets: a point (in mm) and two unit vectors (adimensional).

# Thus each output line wil have the frame and sub-frame indices
# (integers) followed by {9*N} numbers.

import sys, math, re, rn, rmxn
import nmwk_compute_abs_matchstick_parms as cpm

def main_loop():
  # Main loop: reads each input frame, computes the part placements
  # from the markers, writes the output frame.
  nframes = 0 # Frame count.
  while True:
    iframe, jframe, mk = read_frame();
    if iframe == None: break
    nframes = nframes + 1
    rp = cpm.compute_placements(mk)
    if nframes == 1:
      write_csv_headers(rp)
    write_frame(iframe, jframe, rp)
  sys.stderr.write("%d frames processed.\n" % nframes)
  #---------------------------------------------------------------------
  
def read_frame():
  # Tries to reads one data frame from {stdin}. 
  # If it succeeds, returns the frame and sub-frame indices
  # {iframe,jframe} (non-negative integers) and a tuple
  # of {M=19} marker coordinate triples (in mm). 
  # 
  # Any blank or comment lines are ignored, except the line
  # "sampling_freq..." that is copied to {stdout}.
  # 
  # If there are no more data frames in {stdin}, returns
  # {None,None,None}.
  
  M = 20; # Expected number of markers in input data frames.
  while True:
    lin = sys.stdin.readline()
    if lin == "":
      # End of file:
      return None, None, None
    if re.match(r'[ ]*[#][ ]*sampling_freq', lin) != None:
      # Copy to output and ignore:
      sys.stdout.write(lin)
    elif re.match(r'[ ]*([#]|$)', lin) != None:
      # Comment or blank line, ignore:
      pass
    else:
      # Data frame, parse and return:
      lin = lin.strip()
      fld = re.split(r'[ ]+', lin)
      if len(fld) != 2 + 3*M:
        sys.stderr.write("** bad number of fields = %d" % len(fld))
        assert False
      iframe = parse_int(fld[0], "iframe", 0, None)
      jframe = parse_int(fld[1], "jframe", 0, None)
      mk = [None]*M; # The marher coords.
      for k in range(M):
        X = parse_float(fld[2 + 3*k], ("X%d" % k), -4000.0, +4000.0)
        Y = parse_float(fld[3 + 3*k], ("Y%d" % k), -4000.0, +4000.0)
        Z = parse_float(fld[4 + 3*k], ("Z%d" % k), -4000.0, +4000.0)
        mk[k] = (X,Y,Z)
      return iframe, jframe, mk
  #---------------------------------------------------------------------
  
def parse_int(s, name, vmin, vmax):
  # Parses the string {s} as an integer. If it succeeds, checks whether
  # it is at least {vmin} and at most {vmax}. If either is {None}, assumed
  # infinity. If it fails, prints an error message with the {name} and
  # aborts.
  try:
    v = int(s)
    if vmin != None and v < vmin:
      sys.stderr.write("** field %s = %s too small (min %d)" % (name, s, vmin))
      assert False
    if vmax != None and v > vmax:
      sys.stderr.write("** field %s = %s too big (max %d)" % (name, s, vmax))
      assert False
    return v
  except ValueError:
    sys.stderr.write("** field %s = %s is not a valid int" % (name, s))
    assert False
  #---------------------------------------------------------------------
  
def parse_float(s, name, vmin, vmax):
  # Parses the string {s} as a float. If it succeeds, checks whether
  # it is at least {vmin} and at most {vmax}. If either is {None}, assumed
  # infinity. If it fails, prints an error message with the {name} and
  # aborts.
  try:
    v = float(s)
    if vmin != None and v < vmin:
      sys.stderr.write("** field %s = %s too small (min %.8f)" % (name, s, vmin))
      assert False
    if vmax != None and v > vmax:
      sys.stderr.write("** field %s = %s too big (max %.8f)" % (name, s, vmax))
      assert False
    return v
  except ValueError:
    sys.stderr.write("** field %s = %s is not a valid float" % (name, s))
    assert False
  #---------------------------------------------------------------------
  
def write_frame(iframe, jframe, rp):
  # Writes a data frame to {stdout}.  
  # 
  # The args {iframe,jframe} should be ints. The arg {rp} should be a list
  # of {N=19} part placements. Each part placement should be a tuple of
  # the form {(name,parent,L,O,X,Y,Z}, where {name} and {parent} are strings,
  # {L} is the nominal part size, {O} is the position
  # of the part's origin (in mm) and {X,Y,Z} are unit axis vectors.
  # 
  # The function writes to {stdout} {iframe,jfrmae} and the {O,X,Y} data
  # of all parts, as a single line.
  N = 19 # Expeceted number of parts.
  assert len(rp) == N
  sys.stdout.write("%6d %6d" % (iframe, jframe))
  for k in range(N):
    place = rp[k];
    assert len(place) == 7;
    L = place[2]
    O = place[3];
    X = place[4];
    Y = place[5];
    sys.stdout.write("   %6.1f" % L)
    sys.stdout.write("  %5.0f %5.0f %5.0f" % (O[0], O[1], O[2]))
    sys.stdout.write("  %+8.6f %+8.6f %+8.6f" % (X[0], X[1], X[2]))
    sys.stdout.write("  %+8.6f %+8.6f %+8.6f" % (Y[0], Y[1], Y[2]))
  sys.stdout.write("\n")
  #---------------------------------------------------------------------
  
def write_csv_headers(rp):
  # Writes two CSV-like column header lines to {stdout}.  
  # 
  # Assumes that the line has integer indices {iframe,jframe} 
  # then {N=19} part placement blocks. Each block should be 
  # 10 numbers {L,Ox,Oy,Oz,Xx,Xy,Xz,Yx,Yy,Yz}.
  
  N = 19 # Expeceted number of parts.
  assert len(rp) == N
  # First line with part names:
  sys.stdout.write("# ,,") # {iframe,jframe}
  for k in range(N):
    place = rp[k];
    assert len(place) == 7;
    name = place[0]
    sys.stdout.write("%s,,,,,,,,,," % name)
  sys.stdout.write("\n")
  # Second line with field names:
  sys.stdout.write("# Frame,Sub Frame,") # {iframe,jframe}
  for k in range(N):
    sys.stdout.write("Size,Ox,Oy,Oz,Xx,Xy,Xz,Yx,Yy,Yz,")
  sys.stdout.write("\n")
  # Third line with measurement units:
  sys.stdout.write("# ,,") # {iframe,jframe}
  for k in range(N):
    sys.stdout.write("mm,mm,mm,mm,,,,,,,")
  sys.stdout.write("\n")
  #---------------------------------------------------------------------

main_loop()