/** This file contains utilities for working with Fourier transforms, including
    plotting the transform and figuring out what frequency components a
    specific element of the Fourier transform represents. */

/** These functions are intended to help you figure out what frequency or
    period each component in the Fourier transform corresponds to. */

/** Calculate the frequency of the specified index given:
    
    [ft, totalPeriod, index]

    where:
      ft: the transformed sequence
      totalPeriod:  the total period of the initial sequence (e.g. 1 year)
      index:  the index for which we want to calculate the index.

    Note that this gives the frequency.  If you want the *angular* frequency
    (omega) such as to pass to a sin or cos function, you need to multiply this
    frequency by 2 pi.
*/
frequencyFromTotalPeriod[ft, totalPeriod, index] :=
{
   if index == 0  // DC component
      return 0 * 1/totalPeriod  // Get the units right?

   return 1/periodFromTotalPeriod[ft, totalPeriod, index]
}


/** Calculate the period of the specified index given:
    
    [ft, totalPeriod, index]

    where:
      ft: the transformed sequence
      totalPeriod:  the total period of the initial sequence (e.g. 1 year)
      index:  the index for which we want to calculate the index.
*/
periodFromTotalPeriod[ft, totalPeriod, index] :=
{
   N = length[ft]

   if index == 0  // Period is infinite.
      return undef

   if index <= N/2
      return totalPeriod / index
   else
      return -totalPeriod / (N - index)     
}


/** Calculate the period of the specified index given:
    
    [ft, samplePeriod, index]

    where:
      ft: the transformed sequence
      samplePeriod:  the period of a single sample (e.g. 10 ms)
      index:  the index for which we want to calculate the index.
*/
periodFromSamplePeriod[ft, samplePeriod, index] :=
{
   N = length[ft]

   totalPeriod = samplePeriod * N  // Is this valid with padded FFTs?

   println["totalPeriod is $totalPeriod, samplePeriod is $samplePeriod"]

   return periodFromTotalPeriod[ft, totalPeriod, index]
}

/** Calculate the frequency of the specified index given:
    
    [ft, samplePeriod, index]

    where:
      ft: the transformed sequence
      samplePeriod:  the period of each sample (e.g. 10 ms)
      index:  the index for which we want to calculate the index.

    Note that this gives the frequency.  If you want the *angular* frequency
    (omega) such as to pass to a sin or cos function, you need to multiply this
    frequency by 2 pi.
*/
frequencyFromSamplePeriod[ft, samplePeriod, index] :=
{
   if index == 0  // DC component
      return 0 * 1/samplePeriod  // Get the units right?

   println["length[ft] is " + length[ft]]

   return 1/periodFromSamplePeriod[ft, samplePeriod, index]
}

/** Plots a Fourier Transform as a magnitude-phase diagram. */
plotFourier[ft, showDC=true, centered=true, realOnly=false] :=
{
   N = length[ft]
   g = new graphics

   start = 0
   if showDC==false
      start = 1
   
   end = N-1
   if realOnly
      end = N div 2

   offset = 0
   modulus = N
   if centered
      offset = N div 2

   dc = abs[ft@0]
   if dc == 0
      dc = 1
   for i = start to end
   {
      v = ft@i
      mag = abs[v]
      if realOnly
         arg = 0     // Real transform has complex conjugates at i and N-i
      else
         arg = arctan[Im[v], Re[v]]

      x = (i+offset) mod modulus
      transparency = clamp[mag/dc, 0 , 1]^(1/4)
//      println["mag = $mag, dc = $dc, transparency = $transparency"]
      g.color[0, 1, 0, transparency]
      g.fillRectSides[x-1/2, -arg, x+1/2, 0]
      g.color[0, 0, 0, .7]
      g.fillRectSides[x-1/2, -Re[mag], x+1/2, 0]
   }

   return g
}


/** Adds a sinewave to the specified array.  If the array is undef, this
    creates a new array and returns it.

    TODO:  Specify phase?
*/
addSinewaveRaw[f, amplitude, startIndex, endIndex, array=undef] :=
{
   if array == undef
      array = new array[[endIndex+1],0]

   len = endIndex-startIndex
   for k=0 to len
   {
      prev = array@(k+startIndex)
      if prev == undef
         prev = 0
      array@(k+startIndex) = prev + amplitude sin[2 pi f k]
   }

   return array
}