# Module to do Fourier analysis and filtering of egg shape.
# Last edited on 2021-03-28 06:11:01 by jstolfi

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

def barycenter(pt):
  # Assumes tha {pt} is a list of points {pt[i]=(x[i],y[i])} of the
  # outline. Assumes that the points are periodic, namely {pt[i+np]==
  # pt[i]} where {np = len(pt)}.
  
  # Returns the barycenter {(bx,by)} and total area {a} of the interior
  # of the figure.
  
  np = len(pt)

  sum_a = 0 # Sum of areas of elem triangles.
  sum_ax = 0; sum_ay = 0 # Sum barycenters of elem triangles wtd by area.
  for ip in range(np):
    xp, yp = pt[ip]
    xq, yq = pt[(ip + 1) % np]
    da =  rn.cross2d((xp,yp), (xq,yq)) # Area of triangle {opq}.
    xb = (xp + xq)/3 # Baricenter X of triangle {opq}.
    yb = (yp + yq)/3 # Baricenter Y of triangle {opq}.
    sum_a += da
    sum_ax += da*xb
    sum_ay += da*yb
  area = sum_a
  bx = sum_ax/area
  by = sum_ay/area
  return (bx, by), area
  # ----------------------------------------------------------------------
  
def fourier(pt, bar, area):
  # Assumes tha {pt} is a list of points {pt[i]=(x[i],y[i])} of the
  # outline. Assumes that the points are periodic namely {pt[i+np]==
  # pt[i]} where {np = len(pt)}.
  #
  # Returns a list {ft} of pairs {(fx[k], fy[k]) with the Fourier transforms
  # of those signals. Namely, {fx[k]} is the Fourier coefficiente frequency {k}
  # of {x[i]}, and analogously for {fy}.
  # 
  # Assumes that {bar=(bx,by)} is the barycenter of the interior of the
  # curve and {area} is its area. Implicitly shifts and rescales the
  # signals so that the total area is {pi} (same as a cicle of radius 1)
  # and the barycenter is at the origin.
  
  scale = sqrt(pi/area)
  np = len(pt)
  
  ft = [None]*2
  for ia in range(2):
    sg = [ scale*(pt[ip][ia] - bar[ia]) for ip in range(np) ]
    ft[ia] = numpy.fft.fft(sg)
    assert len(ft[ia]) == np
  
  fou = [ (ft[0][ip], ft[1][ip]) for ip in range(np)]
  
  return fou
  # ----------------------------------------------------------------------
    
def spectrum(fou):
  # Assumes that {fou} is the Fourier transform of a curve as
  # returned by {fourier} above.
  #
  # Returns the power spectrum of the signals, namely a list of pairs
  # {(ex[k],ey[k])} where {ex[k]} is the sum of squared norms of all coefficients
  # of absolute frequency {k}.  The length is {ceil(nf/2)} where {nf=len(fou)}.
  
  nf = len(fou)
  np = int(ceil(nf/2))
  pow = [None]*np
  for k in range(np):
    pwk = [0,0]
    for ia in range(2):
      fki = fou[k][ia]
      eki = fki*fki.conjugate()
      if k != 0 and nf-k != k:
        fji = fou[nf-k][ia]
        eki += fji*fji.conjugate()
      assert eki.imag == 0
      pwk[ia] = eki.real
    pow[k] = pwk
  return pow
  # ----------------------------------------------------------------------