import math, complex, strutils

# Works with floats and complex numbers as input
proc fft[T: float | Complex[float]](x: openarray[T]): seq[Complex[float]] =
  let n = x.len
  if n == 0: return
  result.newSeq(n)  
  if n == 1:
    result[0] = (when T is float: complex(x[0]) else: x[0])
    return
  var evens, odds = newSeq[T]()
  for i, v in x:
    if i mod 2 == 0: evens.add v
    else: odds.add v
  var (even, odd) = (fft(evens), fft(odds))
  let halfn = n div 2
  for k in 0 ..< halfn:
    let a = exp(complex(0.0, -TAU * float(k) / float(n))) * odd[k]
    result[k] = even[k] + a
    result[k + halfn] = even[k] - a


proc ifft(x: var openarray[Complex[float]]): seq[Complex[float]] =
  var n: int = x.len
  for p in 0..<n:
    x[p] = (x[p]).conjugate
  var xi = fft(x)
  for p in 0..<n:
    xi[p] = (xi[p].conjugate) / n.float
  return xi


when isMainModule:
  var ff = fft(@[1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0])
  for i in ff:
    echo formatFloat(abs(i), ffDecimal, 3)
  
  
  var iff = ifft(ff)
  for i in iff:
    echo formatFloat(abs(i), ffDecimal, 3)