#! /usr/bin/python3
# Last edited on 2021-03-30 04:04:44 by jstolfi

# Computes the power spectrum of the coordinates of 
# a closed curve.

import curve_filter
import rn
import rmxn
import sys
from math import sqrt, sin, cos, exp, log, floor, ceil, inf, nan, pi
import re

def main():
  pts_file = sys.argv[1] # File name of input points, complete, rel to "tests/".
  oprefix = sys.argv[2]  # Prefix of output files, rel to "tests/", sans "-r{N}", extension.
  
  # Get the raw outlines:
  raw = read_contour(pts_file)
  nr = len(raw)
  
  # Subsample them with cubic interpolation to prevent circle shrinkage:
  dbl = subsample_cubic(raw) 
  nd = len(dbl)
  
  # Resample them at equal length steps:
  ns = 2*nr
  for irep in range(3):
    pts = dbl if irep == 0 else resample_by_length(pts, ns)
    pts_file = ("%s-r%d-rsp.txt" % (oprefix, irep))
    write_pairs(pts_file, pts)
    rms = compute_rms_spectrum(pts)
    rms_file = ("%s-r%d-rms.txt" % (oprefix, irep))
    write_pairs(rms_file, rms)
  return 0
  # ----------------------------------------------------------------------

def subsample_cubic(raw):
  # Subsamples a polygon with vertices {raw} by inserting a point
  # between every two vertices, using bicubic interpolation 
  # with index as argument.
  
  nr = len(raw)
  ns = 2*nr
  sub = []
  for kb in range(nr):
    # Take vertex {kb}:
    pb = raw[kb]
    sub.append(pb)
    # Interpolate between {kb} and {kb+1}:
    ka = (kb-1) % nr; pa = raw[ka]
    kc = (kb+1) % nr; pc = raw[kc]
    kd = (kb+2) % nr; pd = raw[kd]
    q = rn.mix(0.125, rn.add(pa, pd), 0.375, rn.add(pb, pc))
    sub.append(q)
    
  return sub
  # ----------------------------------------------------------------------

def resample_by_length(raw, ns):
  # Resamples a polygon with vertices {raw}
  # by computing {ns} equally spaced steps along its perimeter.
  
  nr = len(raw)
  
  # Compute the total length of the polygon up to each vertex:
  L = [0]  # {L[k]} is the length up to {raw[k]}, for {k} in {0..nr}.
  for k in range(nr):
    dL = rn.dist(raw[(k + 1) % nr], raw[k])
    L.append(L[k] + dL)
  Ltot = L[nr]
  
  # Now interpolate {raw} linearly based on {L}:
  pts = [] # The interpolated samples are {pts[i]} for {i} in {0..ns-1} 
  k0 = 0 # Last index of {L} used.
  for i in range(ns):
    Li = Ltot*(i/ns)
    assert L[k0] <= Li
    while k0 < nr and L[k0 + 1] < Li: k0 = k0 + 1
    assert k0 < nr
    s = (Li - L[k0])/(L[k0 + 1] - L[k0])
    k1 = (k0 + 1) % nr
    pti = rn.mix(1-s, raw[k0], s, raw[k1])
    pts.append(pti)
  return pts
  # ----------------------------------------------------------------------

def compute_rms_spectrum(pts):
  # Receives a list of raw outline points as returned by {read_contour}.
  # Computes the Fourier power spectrum of the {x} and {y} coordiates.
  # Writes it 

  np = len(pts)
  
  bar, area = curve_filter.barycenter(pts)
  sys.stderr.write("area = %.2f  bar = ( %.3f %.3f )\n" % (area, bar[0],bar[1]))
  
  # Compute the Fourier transform of each component:
  fou = curve_filter.fourier(pts, bar, area)
  assert len(fou) == np

  # Convert to a power epectrum and write
  pwr = curve_filter.spectrum(fou)
  ne = len(pwr)
  
  # Convert to rms value
  rms = []
  for k in range(ne):
    ek = pwr[k]
    rmsx = sqrt(ek[0]/np)
    rmsy = sqrt(ek[1]/np)
    rms.append((rmsx,rmsy))
  return rms
  # ----------------------------------------------------------------------
  
def write_pairs(fname, dat):
  # Writes a list {dat} of tuples of floats to file {fname}.
  wr = open(fname, "w")
  nd = len(dat)
  for k in range(nd):
    wr.write("%5d" % k)
    for v in dat[k]:
      wr.write(" %18.12f" % v)
    wr.write("\n")
  wr.close()
  return
  # ----------------------------------------------------------------------
  
def read_contour(fname):
  # Read coordinates of ponts on outline -- two per line, X and Y,
  # as integer multiples of the pixel length. Ignores blank lines.
  #
  # The ouline is not necessarily closed, meaning that the last point 
  # need not be equal to the first one.  However, the two should be 
  # at most 1 pixel apart.
  #
  # Returns the list of pairs {(x[i],y[i]).
  #
  sys.stderr.write("reading pts outline from %s\n" % fname)
  rd = open(fname, 'r', encoding="ISO-8859-1")
  np = 0
  pts = []
  for line in rd:
    line = line.strip()
    fld = re.split(r'[ ]+', line)
    # sys.stderr.write("fld = %s\n" % str(fld))
    if len(fld) == 1 and fld[0] == '':
      # Blank line
      pass
    elif len(fld) == 3:
      ir = int(fld[0])
      xr = float(fld[1])
      yr = float(fld[2])
      assert ir == np
      pts.append((xr, yr))
      np += 1
    else:
      assert False, "bad line format = [%s]" % line
  sys.stderr.write("read %d points\n" % np)
  if rn.dist(pts[0], pts[np-1]) < 0.5:
    sys.stderr.write("outline is closed\n")
    sys.stderr.write("  pts[0] = ( %.9f %.9f )\n" % (pts[0][0], pts[0][1]))
    sys.stderr.write("  pts[%d] =  ( %.9f %.9f )\n" % (np-1,pts[np-1][0],pts[np-1][1]))
    del pts[np-1]
    sys.stderr.write("removing the last point\n")
    np = np - 1
  assert np == len(pts)
  return pts
  # ----------------------------------------------------------------------

main()