import FFT from 'fft.js';

import type { DataReIm } from '../types/index.ts';

import { zeroShift } from './zeroShift.ts';

export interface ReimFFTOptions {
  inverse?: boolean;
  applyZeroShift?: boolean;
  /**
   * Write the result back into the input arrays instead of allocating new ones.
   * @default false
   */
  inPlace?: boolean;
}

/**
 * ReimFFT.
 * @param data - complex spectrum
 * @param options - options.
 * @returns FFT of complex spectrum.
 */
export function reimFFT(
  data: DataReIm,
  options: ReimFFTOptions = {},
): DataReIm<Float64Array> {
  const { inverse = false, applyZeroShift = false, inPlace = false } = options;

  const { re, im } = data;
  const size = re.length;
  const csize = size << 1;

  let complexArray = new Float64Array(csize);
  for (let i = 0; i < csize; i += 2) {
    complexArray[i] = re[i >>> 1];
    complexArray[i + 1] = im[i >>> 1];
  }

  const fft = new FFT(size);
  let output = new Float64Array(csize);
  if (inverse) {
    if (applyZeroShift) complexArray = zeroShift(complexArray, true);
    fft.inverseTransform(output, complexArray);
  } else {
    fft.transform(output, complexArray);
    if (applyZeroShift) output = zeroShift(output);
  }

  if (inPlace) {
    for (let i = 0; i < csize; i += 2) {
      re[i >>> 1] = output[i];
      im[i >>> 1] = output[i + 1];
    }
    return data as DataReIm<Float64Array>;
  }

  const newRe = new Float64Array(size);
  const newIm = new Float64Array(size);
  for (let i = 0; i < csize; i += 2) {
    newRe[i >>> 1] = output[i];
    newIm[i >>> 1] = output[i + 1];
  }

  return { re: newRe, im: newIm };
}