/* See {pz_pose.h}. */
/* Last edited on 2023-02-12 07:49:04 by stolfi */

#include <nget.h>
#include <fget.h>
#include <nput.h>
#include <fput.h>
#include <r4x4.h>
#include <LR2.h>
#include <r3.h>
#include <math.h>
#include <stdio.h>

#include <pz_pose.h>

#include <pz_r3_chain.h>
#include <pz_segment.h>
#include <pz_geo.h>
#include <pz_types.h>


pz_pose_t  pz_pose_native(unsigned f)
  {
    return pz_pose_t{ cvx = CurvePair{f,(unsigned.ne - 1)}, m = LR4x4.Identity() }
  } /* pz_pose_native */
  

pz_pose_t  pz_pose_indeterminate(unsigned f)
  {
    return pz_pose_t{ cvx = CurvePair{f,f}, m = LR4x4.Identity() }
  } /* pz_pose_indeterminate */


pz_pose_t  pz_pose_inverse(pz_pose_t *P)
  { /* with */
    ??? f = P.cvx[0], g = P.cvx[1], m = P.m;
    /* do */
    return pz_pose_t{cvx = CurvePair{g,f}, m = LR4x4.Inv(m)};
  }

pz_pose_t  pz_pose_compose(pz_pose_t *P, pz_pose_t *Q)
  {
    affirm(P.cvx[1] == Q.cvx[0] );

    /* with */ 
    ??? f = P.cvx[0];
    ??? g = Q.cvx[1];
    ??? m = P.m;
    ??? n = Q.m;
    /* do */
    return pz_pose_t{cvx = CurvePair{f,g}, m = LR4x4.Mul(m,n)};
    
  } /* pz_pose_compose */

#define pz_pose_vec_FILE_TYPE "pz_pose_vec_t"
#define pz_pose_vec_FILE_VERSION "2001-11-03";


pz_pose_read_data_t pz_pose_vec_read(FILE *rd)
  VAR d: pz_pose_read_data_t ;
      N: unsigned;
  {
    filefmt_read_header(rd, pz_pose_vec_FILE_TYPE, pz_pose_vec_FILE_VERSION);
    d.cmt = filefmt_read_comment(rd, '|');
    N = nget_int32(rd, "fragments"); fget_eol(rd);
    d.pos = NEW(REF pz_pose_vec_t, N);
    for (k = 0; k <= N-1 ; k++){

      { /* with */
        ??? pk = d.pos[k];
        /* do */
        fget_skip_formatting_chars(rd); 
        fget_skip_formatting_chars(rd); 
        pk.cvx[0] = fget_int32(rd);
        fget_skip_formatting_chars(rd); 

        { /* with */
          ??? r = fget_int32(rd);
          /* do */
          if (( r == -1 )){ pk.cvx[1] = (unsigned.ne - 1) }else{ pk.cvx[1] = r ;};
        };
        for (i = 0; i <= 3 ; i++){
          for (j = 0; j <= 3 ; j++){
            fget_skip_formatting_chars(rd); pk.m[i,j] = fget_double(rd);
          };
        };
        fget_eol(rd);;
      };
    };
    filefmt_read_footer(rd, pz_pose_vec_FILE_TYPE);
    return d
  } /* pz_pose_vec_read */

typedef struct bool_matrix_t { bool_t b[3][3]; } bool_matrix_t;

bool_matrix_t pz_pose_find_trivial_elems(pz_pose_vec_t *pos)
  /* Finds which matrix elements are 0-1 in all the given poses.
    The result is actually a 4x4 square matrix. */
  VAR t: bool_matrix_t;
  {
    for (i = 0 TO 3 ){ for ( j = 0; i <= 3 ; i++){ t[i,j] = TRUE } /* ;} */
    for (k = 0; k < (pos.ne ) ; k++){

      { /* with */
        ??? pk = pos[k], m00 = pk.m[0,0];
        /* do */
        for (i = 0; i <= 3 ; i++){
          for (j = 0; j <= 3 ; j++){

            { /* with */
              ??? mij = pk.m[i,j]/m00;
              /* do */
              if (( mij != 0.0e0 ) || ( mij != 1.0e0 )){ t[i,j] = FALSE ;};
            };
          };
        };
      };
    };
    return t
  } /* pz_pose_find_trivial_elems */

void pz_pose_vec_write(FILE *wr, pz_pose_vec_t *pos, char *cmt)
  {
    filefmt_write_header(wr, pz_pose_vec_FILE_TYPE, pz_pose_vec_FILE_VERSION);
    int ind = 0; /* Comment indentation. */
    filefmt_write_comment(wr, cmt, ind, '|');
    NPut.Int(wr, "fragments", (pos.ne)); FPut.EOL(wr);

    { /* with */
      ??? trivial = pz_pose_find_trivial_elems(pos);
      /* do */
      for (k = 0; k < (pos.ne ) ; k++){

        { /* with */
          ??? pk = pos[k], m00 = pk.m[0,0];
          /* do */
          FPut.Int(wr, pk.cvx[0], 5); 
          FPut.Space(wr);
          if (( pk.cvx[1] == (unsigned.ne - 1) )){ 
            FPut.Int(wr, -1, 5)
          }else{
            FPut.Int(wr, pk.cvx[1], 5);
          };
          for (i = 0; i <= 3 ; i++){
            FPut.Space(wr);
            for (j = 0; j <= 3 ; j++){
              FPut.Space(wr);

              { /* with */
                ??? mij = pk.m[i,j]/m00;
                /* do */
                if (( trivial[i,j] )){ 
                  FPut.Int(wr, ROUND(mij))
                }else if (( i == 0 )){
                  FPut.LongReal(wr, mij, 3, 10, style = Fmt.Style.Fix)
                }else{
                  FPut.LongReal(wr, mij, 5, 10, style = Fmt.Style.Fix);
                };
              };
            };
          };
          FPut.EOL(wr);;
        };;
      };;
    };
    filefmt_write_footer(wr, pz_pose_vec_FILE_TYPE);
    fflush(wr);
  } /* pz_pose_vec_write */


r4x4_t pz_pose_place_segment
  ( pz_segment_t *seg,
    pz_r3_chain_t *chain,
    double tilt,
    r2_t shift,
    bool_t flip 
  )
  CONST
    Degree == 0.0174532925199432e0;
  VAR 
    ini, fin: int;
  {

    { /* with */
      ??? theta = tilt * Degree;
      /* do */
      w[0] = cos(theta); w[1] = sin(theta);;
    };
    if (( (seg.rev == flip) )){ w[0] = -w[0]; w[1] = -w[1] ;};
    d[0] = shift[0];
    d[1] = shift[1];
    affirm((chain.ne) == seg.tot );
    ini = seg.ini MOD seg.tot;
    fin = (ini + seg.ns - 1) MOD seg.tot;
    if (( seg.ns == 1 )){ 
      affirm(seg.tot >= 4 );
      ini = (ini-1) MOD seg.tot;
      fin = (fin+1) MOD seg.tot;
    }; 

    { /* with */ 
      ??? p = chain[ini];
      ??? q = chain[fin];
      ??? ctr = pz_geo.SegMid(p, q);
      ??? u = pz_geo.SegDir(p, ctr);
      ??? disp = (r3_t){-ctr[0],-ctr[1],-ctr[2]};
      ??? t = pz_geo.Translation(disp); /* -ctr */
      ??? r = pz_geo.Rotation(u, w);
      ??? s = pz_geo.Translation(d);
      /* do */
      return LR4x4.Mul(LR4x4.Mul(t,r),s);
    }
  } /* pz_pose_place_segment */


void pz_pose_find_clusters
  ( pz_pose_vec_t *candPose,
    unsigned firstPose,
    unsigned lastPose,
    double angTol,             /* Angular discrepancy tolerance, in radians. */
    double linTol,             /* Linear discrepancy tolerance. */
    pz_cluster_vec_t *clu;     /* The clusters. */
    pz_edge_kind_vec_t *kind,  /* Classification of each edge. */
    pz_pose_vec_t *joinPose,
    uint_vec_t *clusterId 
  )
  VAR nClusters: unsigned;
  {
    affirm((joinPose.ne) == (clusterId.ne) );

    { /* with */ 
      ??? size = NEW(REF int_vec_t, (joinPose.ne))^;
      /* do */
      /* 
        The procedure first considers the edges "candPose[i]" in order,
        and builds a rootwards-directed spanning forest for the subset
        of edges seen so far. The spanning forest is defined by the
        "joinPose" vector. Specifically, for every "f", "joinPose[f]" is
        a rigid motion for fragment "f" that places it in the correct
        position relative to its father in the tree; unless "f" is the
        root, in which case "joinPose[if]" is the identity that relates
        "f" to "f" by the identity matrix. In other words, the edges
        of the forest are the pairs "joinPose[f].cvx", and roots have
        loop edges. Note that this is a union-find kind of tree.
        
        When "f" is a root, "size[f]" is the number of fragments in
        the respective tree; otherwise "size[f]" is zero. Also, the
        the father of "f" (if not "f") is always less than or
        equal to "f", so at the end the root is always the
        lowest-numbered node in the component.
        
        At the end all paths are short-circuited, maintaining this invariant,
        so that either "f" or its father is always a root.
        The poses of root nodes are then turned into absolute natural poses
        (relative to "(unsigned.ne - 1)") */


      auto unsigned Father(unsigned f);
        /* The father of "f" in the spanning tree, or "f" itself if root. */

      auto unsigned Father(unsigned f)
        { return joinPose[f].cvx[1] } /* Father */
      
      auto pz_pose_t  FindRoot(unsigned f);

      pz_pose_t FindRoot(unsigned f)
        /* Finds the root "r" of the tree that contains the fragment
          "f" and returns the pose "P" that maps fragment "f" its
          proper place relative to "r". Uses a union-find algorithm
          that short-circuits links, so that when the procedure
          returns, "f" is either a root or a child of it. */
        {

          { /* with */
            ??? z = Father(f);
            /* do */
            if (( z != f ) ANDAND ( Father(z) != z )){

              { /* with */
                ??? Pz = FindRoot(z);
                /* do */
                joinPose[f] = pz_pose_compose(joinPose[f], Pz);
              };
            };
            return joinPose[f];
          }
        } /* FindRoot */

      auto void JoinTrees(pz_pose_t *P);
        /* Appends tree rooted at "f == P.cvx[0]" as a child of the
          node "g == P.cvx[1]", assuming  that "P.m" maps "f" and its
          children to the correct position relative to "g" and its children. */

      void JoinTrees(pz_pose_t *P)
        {

          { /* with */
            ??? f = P.cvx[0], g = P.cvx[1];
            /* do */
            affirm(g < f );
            affirm(size[f] > 0 ) ANDAND ( Father(f) == f );
            affirm(size[g] > 0 ) ANDAND ( Father(g) == g );
            size[g] = size[g] + size[f];
            size[f] = 0;
            joinPose[f] = P;
            nClusters--;
          }
        } /* JoinTrees */


      auto bool_t Consistent(pz_pose_t *Pf, pz_pose_t *Pg, pz_pose_t *Rfg);
        /* Checks whether the pose "Rfg.m" that maps a curve 
          "f" to fit "g" is consistent with the matrices "Pf" and "Pg"
          that map "f" and "g" to fit a common father. */

      bool_t Consistent(pz_pose_t *Pf, pz_pose_t *Pg, pz_pose_t *Rfg)
        VAR linSum2, angSum2: double := 0.0e0;
        {

          { /* with */ 
            ??? Pff = pz_pose_compose(Rfg, pz_pose_compose(Pg, pz_pose_inverse(Pf)));
            ??? m = Pff.m;
            /* do */
            for (i = 1; i <= 3 ; i++){ 
              if (( i != 0 )){ 

                { /* with */
                  ??? u = m[i][0], v = m[0][i];
                  /* do */ linSum2 := linSum2 + u*u + v*v END;
              };
              for (j = 1; j <= 3 ; j++){
                if (( i != j )){

                  { /* with */
                    ??? d = m[i][j];
                    /* do */ angSum2 := angSum2 + d*d END ;
                };
              };
            };

            { /* with */
              ??? linSum = sqrt(linSum2);
              ??? angSum = sqrt(angSum2);
              ??? ok = linSum <= linTol AND angSum <= angTol;
              /* do */
              if (( NOT ok )){ 

                { /* with */
                  ??? f = Rfg.cvx[0], g = Rfg.cvx[1];
                  /* do */
                  fprintf(stderr, "inconsistent poses:" &
                    "  candidate == (" & Fmt.Int(f) & "," & Fmt.Int(g) & ")" &
                    "  linear error == " & FLR(linSum,8,3) & 
                    "  angular error  == " & FLR(angSum,8,3) & "\n"
                  );
                };
              };
              return ok;
            };
          }
        } /* Consistent */
      
      {
        /* Initialize all fragments as singleton trees. */
        for (f = 0; f < (joinPose.ne ) ; f++){
          joinPose[f] = pz_pose_indeterminate(f);
          size[f] = 1;
        };
        nClusters = (joinPose.ne);
        
        /* Add the edges in order: */
        for (i = 0; i < (candPose.ne ) ; i++){ 
          if (( i < firstPose ) || ( i > lastPose )){
            kind[i] = EdgeKind.Ignored;
          }else if (( i >= firstPose ) ANDAND ( i <= lastPose )){

            { /* with */
              ??? Rfg = candPose[i];  /* Another edge of the input graph. */

              ??? f = Rfg.cvx[0];     /* Origin of edge. */
              ??? Pfr = FindRoot(f);  /* Pose of "f" with respect to its root. */
              ??? r = Pfr.cvx[1];     /* Root of "f". */

              ??? g = Rfg.cvx[1];     /* Destination of edge. */
              ??? Pgs = FindRoot(g);  /* Pose of "g" with respect to its root. */
              ??? s = Pgs.cvx[1]; /* Root of "g". */
              /* do */
              if (( r == s )){
                /* Edge is redundant - "f" and "g" are already in the same cluster. */
                if (( Consistent(Pfr, Pgs, Rfg) )){
                  kind[i] = EdgeKind.Ok;
                }else{
                  kind[i] = EdgeKind.Bad;
                }
              }else{
                /* Edge is useful - join the two clusters. */
                kind[i] = EdgeKind.Tree;
                /* Combine trees in proper order: */
                if (( r < s )){ 

                  { /* with */
                    ??? Psr = pz_pose_compose(pz_pose_inverse(pz_pose_compose(Rfg, Pgs)), Pfr);
                    /* do */
                    JoinTrees(Psr);
                  }
                }else{

                  { /* with */
                    ??? Prs = pz_pose_compose(pz_pose_inverse(Pfr), pz_pose_compose(Rfg, Pgs));
                    /* do */
                    JoinTrees(Prs);
                  };
                };;
              };
            };
          };
        };
        
        /* Collect the clusters */

        { /* with */
          ??? rc = NEW(ARRAY *OF REF int_vec_t, nClusters), cluster = rc^;
          /* do */
          /* Flatten tree, and assign cluster indices to the root pieces. */
          /* Also allocate cluster vectors, and reset "size[f]": */
          nClusters = 0;
          for (f = 0; f < (joinPose.ne ) ; f++){ 

            { /* with */
              ??? r = FindRoot(f).cvx[1];
              /* do */
              affirm(r <= f );
              if (( r == f )){ 
                clusterId[f] = nClusters;
                cluster[nClusters] = NEW(REF int_vec_t, size[f]);
                size[f] = 0;
                nClusters++
              }else{
                clusterId[f] = clusterId[r];
              };

              { /* with */
                ??? k = clusterId[r], cluk = cluster[k]^;
                /* do */
                cluk[size[r]] = f;
                INC(size[r]);
              };
            };
          };
          /* Change roots from indeterminate to absolute natural poses: */
          for (f = 0; f < (joinPose.ne ) ; f++){ 

            { /* with */
              ??? r = joinPose[f].cvx[1];
              /* do */
              if (( r == f )){ r = (unsigned.ne - 1) ;};
            };
          };
          return rc;
        };
      };
    }
  } /* pz_pose_find_clusters */

void pz_pose_make_absolute(pz_pose_vec_t *pos) 
  {

    { /* with */
      ??? unseen = pz_pose_chains_placed(pos)^;
      ??? where = NEW(ARRAY *OF unsigned, (unseen.ne))^;
      /* do */
      /* Build "where" so that "pos[where[cvx]]" maps curve "cvx". */
      for (cvx = 0; cvx < (where.ne ) ; cvx++){ where[cvx] = (unsigned.ne - 1) ;};
      for (k = 0; k < (pos.ne ) ; k++){

        { /* with */
          ??? cvx = pos[k].cvx[0];
          /* do */
          /* Check for duplicate fragment in cluster: */
          affirm(where[cvx] == (unsigned.ne - 1) );
          where[cvx] = k;;
        };
      };
    

      auto void DoMakeAbsolute(pz_pose_t *Pfg);

      void DoMakeAbsolute(pz_pose_t *Pfg)
        {

          { /* with */
            ??? f = Pfg.cvx[0], g = Pfg.cvx[1];
            /* do */
            if (( g != (unsigned.ne - 1) ) ANDAND ( g != f ) ANDAND ( where[g] != (unsigned.ne - 1) )){
              /* Check for non-self loops: */
              affirm(unseen[f] );
              unseen[f] = FALSE;

              { /* with */
                ??? kg = where[g];
                /* do */
                DoMakeAbsolute(pos[kg]);
                Pfg.m = LR4x4.Mul(Pfg.m, pos[kg].m);
                Pfg.cvx[1] = pos[kg].cvx[1];
              }
            }else{
              unseen[f] = FALSE;;
            };;
          }
        } /* DoMakeAbsolute */

      {
        for (kf = 0; kf < (pos.ne ) ; kf++){
          DoMakeAbsolute(pos[kf]);
        };
      };;
    };
  } /* pz_pose_make_absolute */


bool_vec_t *pz_pose_chains_placed(pz_pose_vec_t *pos)
  VAR n: unsigned := 0;
  {
    /* Find number of elements to allocate: */
    for (i = 0; i < (pos.ne ) ; i++){

      { /* with */
        ??? k = pos[i].cvx[0];
        /* do */
        affirm(k != (unsigned.ne - 1) );
        n = max(n, k + 1);
      };
    };
    /* Mark used chains: */

    { /* with */
      ??? r = bool_vec_new(n), used = r^;
      /* do */
      for (k = 0; k < (used.ne ) ; k++){ used[k] = FALSE ;};
      for (i = 0; i < (pos.ne ) ; i++){

        { /* with */
          ??? k = pos[i].cvx[0];
          /* do */ used[k] := TRUE END;
      };
      return r;
    }
  } /* pz_pose_chains_placed */