ffts-and-stuff.typ

// Complex numbers are represented as pairs

#let re(x) = x.at(0)
#let im(x) = x.at(1)
#let cmul(x, y) = (re(x)*re(y) - im(x)*im(y), re(x)*im(y) + im(x)*re(y))
#let cadd(x, y) = (re(x) + re(y), im(x) + im(y))
#let csub(x, y) = (re(x) - re(y), im(x) - im(y))
#let cscale(k, x) = (k * re(x), k * im(x))
#let cis(phi) = (calc.cos(phi), calc.sin(phi))
#let cshow(x) = [#re(x) + i #im(x)]

#let rs2cs(v) = v.map(x => (x, 0))

#let fft(v, inv: false) = {
  let n = v.len()
  if n <= 1 {
    v
  } else {
    let half_n = calc.quo(n, 2)
    let evens = range(half_n).map(i => v.at(2*i))
    let odds = range(half_n).map(i => v.at(2*i + 1))

    let evens_fft = fft(evens, inv: inv)
    let odds_fft = fft(odds, inv: inv)

    let sign = if inv { 1 } else { -1 }
    let w_d = cis(sign * 2*calc.pi/n)
    let w = (1, 0)

    let res = (0,)*n

    for k in range(half_n) {
      res.at(k) = cadd(evens_fft.at(k), cmul(w, odds_fft.at(k)))
      res.at(k + half_n) = csub(evens_fft.at(k), cmul(w, odds_fft.at(k)))

      w = cmul(w, w_d)
    }

    if inv {
      for k in range(n) {
        res.at(k) = cscale(1/2, res.at(k))
      }
    }

    res
  }
}

#import "typst-plot/plot.typ": plot
#import "typst-plot/plot-sample.typ": *

#let n = 256
#let half_n = calc.quo(n, 2)

#let make_plot(data, shift: false) = {
  let n = data.len()

  let maximum = calc.max(..data.map(x => re(x)), ..data.map(x => im(x)))
  let minimum = calc.min(..data.map(x => re(x)), ..data.map(x => im(x)))

  if shift {
    plot(
      x-tics: (every: n / 4),
      y-tics: (every: calc.round((maximum - minimum) / 4, digits: 2)),
      (data: range(n).map(i => (i - half_n, re(data.at(calc.rem(i + half_n, n)))))),
      (data: range(n).map(i => (i - half_n, im(data.at(calc.rem(i + half_n, n)))))),
    )
  } else {
    plot(
      x-tics: (every: n / 4),
      y-tics: (every: calc.round((maximum - minimum) / 4, digits: 2)),
      (data: range(n).map(i => i, re(data.at(i)))),
      (data: range(n).map(i => i, im(data.at(i)))),
    )
  }
}

#let a = ((1, 0),)*20 + ((0, 0),)*(n - 39) + ((1, 0),)*19
// #let b = range(n).map(i => (0.15 * calc.cos(40 * 2 * calc.pi * i / n), 0))
// #let b = range(n).map(i => cscale(0.15/2, cis(40 * 2 * calc.pi * i / n)))
// #let q = a.zip(b).map(pair => cadd(pair.at(0), pair.at(1)))
#let q = a
#make_plot(q, shift: true)

#let q_fft = fft(q)
#make_plot(q_fft, shift: true)