# Module to smooth curves and compute local curvature.
# Last edited on 2021-03-28 16:39:37 by jstolfi

import lsq
import rn
import rmxn
import sys
import math

def fit_circle(p, q, R_ini):
  # Fits a quadratic function to the raw outline points in the window
  # {q[0..]}, assumed to be centered at {p} of the raw outline. Namely,
  # returns point {o}, unit vector {u}, and curvature {cvt} such that
  # 
  #   q[j][0] ~= o[0] + u[0]*x[j] + v[0]*h(x[j],1/cvt)
  #   q[j][1] ~= o[1] + u[1]*x[j] + v[1]*h(x[j],1/cvt)
  #
  # where {x[j]} is {q[j] - p} projected in the direction {u},
  # {v} is the unit vector 90 degrees ccw from {u},
  # {h(x,R)} is {R - sqrt(R^2 - x^2)}.
  #
  # Requires a guess {R_ini} for the radius of curvature.
  
  nq = len(q)
  u = initial_tangent(p, q)
  R_ref = rn.dist(q[0], q[nq-1])
  R_max = 1000*R_ref
  R_min = 0.1*R_ref
  R = min(R_max, max(R_min, R_ini)) # Initial guess for the radius.
  cvt = 1/R
  niter = 20 # Iterations of adjustment.
  for k in range(niter):
    o, u_new, cvt_new = fit_circle_aux(q, p, u, cvt)
    show_fit("  ", o, u_new, cvt_new)
    R_new = max(R_min, min(R_max, 1/cvt_new))
    if abs(R_new - R) < 0.01: 
      sys.stderr.write("=== converged ===\n")
      break
    u = u_new; R = R_new; cvt = cvt_new
  return o, u_new, cvt_new
  # ----------------------------------------------------------------------

def show_fit(pref, o, u, cvt):
  sys.stderr.write(pref)
  sys.stderr.write("o = ( %10.2f %10.2f )" % (o[0], o[1]))
  sys.stderr.write(" u = ( %+10.9f %+10.9f )" % (u[0],u[1]))
  sys.stderr.write(" cvt = %+10.9f R = %14.9f" % (cvt, 1/cvt))
  sys.stderr.write("\n")

def initial_tangent(p, q):
  # Computes an initial guess for the tangent direction {u}.
  nq = len(q)
  u, dw = rn.dir(rn.sub(q[nq-1], q[0]))
  return u
  # ----------------------------------------------------------------------

def local_coords(p, q, u):
  # Computes the coordinates of the window points {q[i]}
  # in the {p,u,v} coordinate system, where {v} is {u} rotated
  # 90 degrees CCW.  Assumes that {u} has unit length.

  nq = len(q)
  v = (-u[1],u[0])
  D = []
  for i in range(nq):
    pqi = rn.sub(q[i], p)
    xi = rn.dot(pqi, u)
    yi = rn.dot(pqi, v)
    D.append((xi,yi))
  return D
  # ----------------------------------------------------------------------

def fit_circle_aux(q, p, u, cvt):
  # Fit an arc of a circle to the points in {q}.
  # The circle will have center on the line
  # through {p} that is perpendicular to {u}.
  # 
  # Requires a guesses {u} for the tangent vector and
  # {cvt} for the local curvature.
  # Returns fitted point {o}, adjusted tangent vector {u_new},
  # adjusted curvature {cvt_new}.
  
  debug = False
  debug_mat = False
  debug_plot = False
  
  nq = len(q)
  R = 1/cvt
  # Data to be fitted:
  D = local_coords(p, q, u)
  wrad = min(abs(D[0][0]), abs(D[nq-1][0]))  # Window radius.
  wrad_max = 0.999*R  # To ensure that {h(x,R)} is valid. 
  if wrad > wrad_max:
    sys.stderr.write("weight window radius reduced from %.9f to %.9f\n" % (wrad,wrad_max))
    wrad = wrad_max
  wt = make_weight_table(D, wrad)

  v = (-u[1],u[0])       # Inward normal direction.
  G = make_basis_matrix(D, wrad, R, wt)
  M = lsq.mat(G, wt)
  B = lsq.rhs(D, G, wt)
  A = lsq.solve(M, B)
  if debug_mat:
    sys.stderr.write("M =\n")
    rmxn.print(sys.stderr, "  | ", M, "%14.9f", " ", " |\n")
    sys.stderr.write("\n")

    sys.stderr.write("B =\n")
    rmxn.print(sys.stderr, "  | ", B, "%14.9f", " ", " |\n")
    sys.stderr.write("\n")

    sys.stderr.write("A =\n")
    rmxn.print(sys.stderr, "  | ", A, "%14.9f", " ", " |\n")
    sys.stderr.write("\n")
  
  # Parse the coefficients {A} into {o,u,cvt_new}.
  # Note that column {j} of {G} is scaled by {1/wrad^j}.
  o_new = rn.mix(1.0, p, +A[0][1], v)
  u_new = rn.mix(1.0, u, +A[1][1]/wrad, v)
  cvt_new = A[2][1]/(wrad**2)*cvt
  if debug:
    ang = math.atan2(u[1],u[0]) # Curent tangent angle.
    u_new,mu_new = rn.dir(u_new)
    ang_new = math.atan2(u_new[1],u_new[0])
    dang = ang_new - ang
    dang = (dang + math.pi) % (2*math.pi) - math.pi
    sys.stderr.write("ang old = %14.9f" % ang)
    sys.stderr.write(" new = %14.9f (len = %14.9f)" % (ang_new,mu_new))
    sys.stderr.write(" change = %14.9f\n" % dang)
    tilt = A[2][0] # Tilt of parabola's axis.
    sys.stderr.write("cvt old = %14.9f new = %14.9f\n" % (cvt, cvt_new))
    sys.stderr.write("R old = %14.9f new = %14.9f\n" % (1/cvt, 1/cvt_new))
    sys.stderr.write("ignored residues:\n")
    sys.stderr.write("linear for u[0] = %14.9f\n" % (A[1][0]/wrad))
    sys.stderr.write("quadratic for u[0] = %14.9f\n" % (A[2][0]/wrad**2))
    sys.stderr.write("cubic = %14.9f %14.9f\n" % (A[3][0]/wrad**3,A[3][1]/wrad**3))
    
  if debug_plot:
    wr = open("tests/out/fit.txt","w")
    for i in range(nq):
      if wt[i] > 0:
        wr.write("%6d  %16.9f  %16.9f" % (i, q[i][0],q[i][1]))
        for j in range(2):
          wr.write("  %16.9f" % D[i][j])
          Fij = 0
          for k in range(4):
            Fijk =  G[i][k]*A[k][j]
            wr.write(" %16.9f" % Fijk)
            Fij += Fijk
          wr.write(" %16.9f" % Fij)
        wr.write("\n")
    wr.close()
  return o_new, u_new, cvt_new
  # ----------------------------------------------------------------------
  
def make_basis_matrix(D, wrad, R, wt):
  # Creates the basis matric for the fitting. 
  #
  # It has 4 columns. In principle, column 0 and 1 are {1} and {x[i]},
  # where {x[i] = D[i][0]}. Column 2 is {R - sqrt(R^2 - x[i]^2)}. Column
  # 3 is {x[i]**3} except that it is made orthogonal to column 1.
  # 
  # However, in order to keep the values under 1.0, each column {j} is
  # scaled by {1/wrad^j}.
  
  nq = len(D)
  G = []
  for i in range(nq):
    xi = D[i][0]
    if wt[i] > 0:
      assert abs(xi) < R, "bad weights"
      Gi0 = 1
      Gi1 = xi/wrad
      Gi2 = (R - math.sqrt(R*R - xi*xi))/(wrad*wrad)
      Gi3 = (xi/wrad)**3
      G.append( [Gi0, Gi1, Gi2, Gi3])
    else:
      G.append( [0,0,0,0] )

  # Make the cubic term orthogonal to the linear term:
  ortho_col(G,3,1,wt)
  
  return G
  # ----------------------------------------------------------------------
  
def ortho_col(G, j0, j1, wt):
  # Makes column {j0} of {G} orthogonal to column {j1}.
  
  # Compute dot products:
  nq = len(G)
  dot01 = 0
  dot11 = 0
  for i in range(nq):
    dot01 += G[i][j0]*G[i][j1]*wt[i]
    dot11 += G[i][j1]**2*wt[i]
  assert dot11 > 1.0e-5
  
  # Make orthogonal, compute norm:
  s = dot01/dot11
  for i in range(nq):
    G[i][j0] = G[i][j0] - s*G[i][j1]
  
  return
  # ----------------------------------------------------------------------
  
def make_weight_table(D, wrad):
  # Returns a table with the weight of each data point {q[i]} in the windo.
  # it is a decaying function of the absolute local tangent coordinate {x[i]=D[i][0]} from the
  # window center {p} to {q[i]} along the current tangent direction {u}. 
  
  nq = len(D)
  wt = [ lsq_weight(D[i][0], wrad) for i in range(nq) ]
  
  # Normalize to unit sum:
  wtot = 0
  for i in range(nq): wtot += wt[i]
  for i in range(nq): wt[i] = wt[i]/wtot
  
  return wt
  # ----------------------------------------------------------------------

def lsq_weight(xk, wrad):
  if abs(xk) > wrad:
    w = 0
  else:
    # Gaussian:
    # z = 5*xk/wrad
    # w = math.exp(-z*z/2)
    # Hahn:
    a = math.pi*xk/wrad
    w = 0.5*(1+math.cos(a))
  return w
  
def test():
  
  # Create the test data:
  ctr = (500,500)
  R = 400
  cvt = 1/R
  ang = math.pi/3
  u = (-math.sin(ang), +math.cos(ang))
  v = (-u[1],+u[0])
  o = rn.mix(1.0, ctr, -R, v)
  nq = 401; assert (nq%2) == 1
  dang = nq/R
  cubic = 0
  noise = 3
  q = []
  for k in range(nq):
    s = (k+0.5)/nq - 0.5
    ak = ang + dang*s
    dk = cubic*s*s*s
    xk = ctr[0] + (R+dk)*math.cos(ak) + 0.5*noise*math.cos(17*k*k+7)
    yk = ctr[1] + (R+dk)*math.sin(ak) + 0.5*noise*math.cos(41*k*k*k-23)
    q.append((xk,yk))

  p = q[nq//2]
  
  # Fit the circle:
  
  R_ini = 3*R
  o1, u1, cvt1 = fit_circle(p, q, R_ini)
  v1 = (-u1[1],+u1[0])
  
  err_o = rn.dist(o, o1)
  
  err_u = rn.dist(u, u1)
  
  R1 = 1/cvt1
  err_R = R1 - R
  
  ctr1 = rn.mix(1.0, o1, R1, v1)
  err_ctr = rn.dist(ctr, ctr1)
  
  sys.stderr.write("------------------------------------------------------\n")
  sys.stderr.write("o = ( %.14f %.14f ) o1 = ( %.14f %.14f ) err = %.14f\n" \
    % (o[0],o[1], o1[0],o1[1], err_o))
  sys.stderr.write("u = ( %.14f %.14f ) u1 = ( %.14f %.14f ) err = %.14f\n" \
    % (u[0],u[1], u1[0],u1[1], err_u))
  sys.stderr.write("R = %.14f R1 = %.14f err = %.14f\n" %(R, R1, err_R))
  sys.stderr.write("ctr = ( %.14f %.14f ) ctr1 = ( %.14f %.14f ) err = %.14f\n" \
    % (ctr[0],ctr[1],ctr1[0],ctr1[1],err_ctr))
  return
  # ----------------------------------------------------------------------