import { dbToLinear, createScanPoints, distance3, findPeak, scanIndexToPosition, TAU } from "./math.js";
import { binsForBand, extractStftBins, getFrameCount } from "./stft.js";

function throwIfAborted(signal) {
  if (signal?.aborted) {
    const error = new Error("계산이 취소되었습니다.");
    error.name = "AbortError";
    throw error;
  }
}

function yieldToEventLoop() {
  return new Promise((resolve) => setTimeout(resolve, 0));
}

export function createDistanceMatrix(arrayDefinition, scan) {
  const points = createScanPoints(scan);
  const microphoneCount = arrayDefinition.microphones.length;
  const pointCount = scan.columns * scan.rows;
  const distances = new Float32Array(pointCount * microphoneCount);
  for (let pointIndex = 0; pointIndex < pointCount; pointIndex += 1) {
    const pointOffset = pointIndex * 4;
    const point = [points[pointOffset], points[pointOffset + 1], points[pointOffset + 2]];
    for (let microphoneIndex = 0; microphoneIndex < microphoneCount; microphoneIndex += 1) {
      distances[pointIndex * microphoneCount + microphoneIndex] = distance3(
        point,
        arrayDefinition.microphones[microphoneIndex].position,
      );
    }
  }
  return distances;
}

export async function computeBinMap(options) {
  const {
    bin,
    spectrum,
    frameCount,
    arrayDefinition,
    pointCount,
    distances,
    sampleRate,
    fftSize,
    speedOfSound,
    signal,
    pointBatchSize,
  } = options;
  const microphoneCount = arrayDefinition.microphones.length;
  const output = new Float32Array(pointCount);
  const steeringReal = new Float64Array(microphoneCount);
  const steeringImaginary = new Float64Array(microphoneCount);
  const angularWaveNumber = (TAU * bin * sampleRate) / (fftSize * speedOfSound);
  const microphoneScale = 1 / microphoneCount;

  for (let pointIndex = 0; pointIndex < pointCount; pointIndex += 1) {
    throwIfAborted(signal);
    const distanceOffset = pointIndex * microphoneCount;
    for (let microphoneIndex = 0; microphoneIndex < microphoneCount; microphoneIndex += 1) {
      const microphone = arrayDefinition.microphones[microphoneIndex];
      const calibrationPhase = ((microphone.phaseOffsetDeg ?? 0) * Math.PI) / 180;
      const phase = angularWaveNumber * distances[distanceOffset + microphoneIndex] - calibrationPhase;
      const calibrationGain = dbToLinear(-(microphone.sensitivityDb ?? 0));
      steeringReal[microphoneIndex] = Math.cos(phase) * calibrationGain;
      steeringImaginary[microphoneIndex] = Math.sin(phase) * calibrationGain;
    }

    let accumulatedPower = 0;
    for (let frameIndex = 0; frameIndex < frameCount; frameIndex += 1) {
      const spectrumOffset = frameIndex * microphoneCount;
      let sumReal = 0;
      let sumImaginary = 0;
      for (let microphoneIndex = 0; microphoneIndex < microphoneCount; microphoneIndex += 1) {
        const spectrumIndex = spectrumOffset + microphoneIndex;
        const inputReal = spectrum.real[spectrumIndex];
        const inputImaginary = spectrum.imaginary[spectrumIndex];
        const phaseReal = steeringReal[microphoneIndex];
        const phaseImaginary = steeringImaginary[microphoneIndex];
        sumReal += inputReal * phaseReal - inputImaginary * phaseImaginary;
        sumImaginary += inputReal * phaseImaginary + inputImaginary * phaseReal;
      }
      accumulatedPower += (sumReal * sumReal + sumImaginary * sumImaginary) * microphoneScale * microphoneScale;
    }
    output[pointIndex] = accumulatedPower / frameCount;
    if ((pointIndex + 1) % pointBatchSize === 0) await yieldToEventLoop();
  }
  return output;
}

export function estimateDasWorkload(options) {
  const { channels, sampleRate, arrayDefinition, experiment } = options;
  const { fftSize, hopSize } = experiment.stft;
  const bandBins = experiment.frequencyBands.map((band) => binsForBand(band, sampleRate, fftSize));
  const selectedBins = [...new Set(bandBins.flat())].sort((a, b) => a - b);
  const frameCount = getFrameCount(channels[0].length, fftSize, hopSize);
  const microphoneCount = arrayDefinition.microphones.length;
  const pointCount = experiment.scan.columns * experiment.scan.rows;
  const bandCount = experiment.frequencyBands.length;
  const distanceBytes = pointCount * microphoneCount * Float32Array.BYTES_PER_ELEMENT;
  const outputBytes = bandCount * pointCount * Float32Array.BYTES_PER_ELEMENT * 2;
  const binMapBytes = pointCount * Float32Array.BYTES_PER_ELEMENT;
  const spectrumBytesPerBin = frameCount * microphoneCount * Float32Array.BYTES_PER_ELEMENT * 2;
  const fftWorkspaceBytes = fftSize * Float64Array.BYTES_PER_ELEMENT * 2;
  const baseWorkingBytes = distanceBytes + outputBytes + binMapBytes + fftWorkspaceBytes;
  const complexMacCount = pointCount * microphoneCount * frameCount * selectedBins.length;
  return {
    bandBins,
    selectedBins,
    frameCount,
    microphoneCount,
    pointCount,
    distanceBytes,
    outputBytes,
    spectrumBytesPerBin,
    baseWorkingBytes,
    complexMacCount,
  };
}

export function finalizeDasBands(rawMaps, workload, experiment) {
  return experiment.frequencyBands.map((band, bandIndex) => {
    const rawMap = rawMaps[bandIndex];
    for (let pointIndex = 0; pointIndex < workload.pointCount; pointIndex += 1) {
      rawMap[pointIndex] /= workload.bandBins[bandIndex].length;
    }
    const rawPeak = findPeak(rawMap);
    const displayMap = Float32Array.from(rawMap);
    let displayScale = "power";
    if (experiment.beamforming.normalization === "peak" && rawPeak.value > 0) {
      displayScale = "normalized";
      for (let pointIndex = 0; pointIndex < workload.pointCount; pointIndex += 1) {
        displayMap[pointIndex] /= rawPeak.value;
      }
    }
    return {
      id: band.id,
      color: band.color,
      bins: workload.bandBins[bandIndex],
      rawMap,
      displayMap,
      displayScale,
      peak: {
        index: rawPeak.index,
        position: scanIndexToPosition(rawPeak.index, experiment.scan),
        rawPower: rawPeak.value,
      },
    };
  });
}

export async function computeDasMap(options) {
  const {
    channels,
    sampleRate,
    arrayDefinition,
    scene,
    experiment,
    signal,
    onProgress,
    maximumWorkingBytes = 512 * 1024 * 1024,
    maximumComplexMacs = 1_000_000_000,
    maximumFftButterflies = 1_000_000_000,
    maximumSpectrumBatchBins = Number.MAX_SAFE_INTEGER,
    pointBatchSize = 64,
    stftTransformBatchSize = 32,
  } = options;
  if (channels.length !== arrayDefinition.microphones.length) {
    throw new Error(`채널 수 ${channels.length}와 마이크 수 ${arrayDefinition.microphones.length}가 다릅니다.`);
  }
  if (!Number.isInteger(pointBatchSize) || pointBatchSize < 1) throw new Error("pointBatchSize는 1 이상의 정수여야 합니다.");
  const workload = estimateDasWorkload({ channels, sampleRate, arrayDefinition, experiment });
  if (workload.complexMacCount > maximumComplexMacs) {
    throw new Error(
      `CPU DAS 예상 복소 누산 ${workload.complexMacCount.toLocaleString()}회가 ` +
        `현재 한도 ${maximumComplexMacs.toLocaleString()}회를 넘습니다. grid, duration 또는 대역폭을 줄이세요.`,
    );
  }
  const availableSpectrumBytes = maximumWorkingBytes - workload.baseWorkingBytes;
  const batchCapacity = Math.floor(availableSpectrumBytes / workload.spectrumBytesPerBin);
  if (batchCapacity < 1) {
    throw new Error(
      `CPU DAS 기본 working set ${(workload.baseWorkingBytes / 1024 ** 2).toFixed(1)} MiB가 ` +
        `현재 한도 ${(maximumWorkingBytes / 1024 ** 2).toFixed(1)} MiB를 넘습니다.`,
    );
  }
  const batchSize = Math.max(
    1,
    Math.min(workload.selectedBins.length, batchCapacity, Math.max(1, maximumSpectrumBatchBins)),
  );
  const spectrumBatchCount = Math.ceil(workload.selectedBins.length / batchSize);
  const fftButterflyCount =
    spectrumBatchCount *
    workload.frameCount *
    workload.microphoneCount *
    (experiment.stft.fftSize / 2) *
    Math.log2(experiment.stft.fftSize);
  if (fftButterflyCount > maximumFftButterflies) {
    throw new Error(
      `CPU STFT 예상 butterfly ${Math.ceil(fftButterflyCount).toLocaleString()}회가 ` +
        `현재 한도 ${maximumFftButterflies.toLocaleString()}회를 넘습니다. duration, FFT 크기 또는 대역폭을 줄이세요.`,
    );
  }
  const distances = createDistanceMatrix(arrayDefinition, experiment.scan);
  const rawMaps = experiment.frequencyBands.map(() => new Float32Array(workload.pointCount));
  const bandIndexesByBin = new Map();
  workload.bandBins.forEach((bins, bandIndex) => {
    bins.forEach((bin) => {
      if (!bandIndexesByBin.has(bin)) bandIndexesByBin.set(bin, []);
      bandIndexesByBin.get(bin).push(bandIndex);
    });
  });

  let completedBins = 0;
  for (let batchStart = 0; batchStart < workload.selectedBins.length; batchStart += batchSize) {
    throwIfAborted(signal);
    await yieldToEventLoop();
    throwIfAborted(signal);
    const batchBins = workload.selectedBins.slice(batchStart, batchStart + batchSize);
    const stft = await extractStftBins(channels, {
      fftSize: experiment.stft.fftSize,
      hopSize: experiment.stft.hopSize,
      bins: batchBins,
      signal,
      transformBatchSize: stftTransformBatchSize,
      onProgress: ({ completedTransforms, totalTransforms }) =>
        onProgress?.({
          stage: "stft",
          completedBins,
          totalBins: workload.selectedBins.length,
          completedTransforms,
          totalTransforms,
          spectrumBatch: Math.floor(batchStart / batchSize) + 1,
          spectrumBatchCount,
        }),
    });
    await yieldToEventLoop();
    throwIfAborted(signal);

    for (const bin of batchBins) {
      const binMap = await computeBinMap({
        bin,
        spectrum: stft.spectra.get(bin),
        frameCount: stft.frameCount,
        arrayDefinition,
        pointCount: workload.pointCount,
        distances,
        sampleRate,
        fftSize: experiment.stft.fftSize,
        speedOfSound: scene.speedOfSound,
        signal,
        pointBatchSize,
      });
      for (const bandIndex of bandIndexesByBin.get(bin)) {
        const rawMap = rawMaps[bandIndex];
        for (let pointIndex = 0; pointIndex < workload.pointCount; pointIndex += 1) {
          rawMap[pointIndex] += binMap[pointIndex];
        }
      }
      completedBins += 1;
      onProgress?.({ stage: "beamforming", completedBins, totalBins: workload.selectedBins.length });
    }
  }

  const bands = finalizeDasBands(rawMaps, workload, experiment);

  return {
    backend: "cpu",
    frameCount: workload.frameCount,
    microphoneCount: workload.microphoneCount,
    pointCount: workload.pointCount,
    workload,
    fftButterflyCount,
    bands,
  };
}