LoudnessMeasurer.kt

package com.keyboardr.bluejay.audio

import com.keyboardr.bluejay.util.toFraction
import java.nio.ByteBuffer
import kotlin.math.*

class LoudnessMeasurer(private val channelCount: Int, sampleRate: Float, totalFrames: Int) {

  private class LoudnessMeasurerChannel(var bufferFrames: Int,
                                        /** How many frames are needed for a gating block. Will correspond to 400ms
                                        of audio at initialization, and 100ms after the first block (75% overlap
                                        as specified in the 2011 revision of BS1770). */
                                        var neededFrames: Int,
                                        blocksCapacity: Int,
                                        overviewCapacity: Int) {
    val inputBuffer = FloatArray(bufferFrames)

    val bufferPre = FloatArray(bufferFrames)
    val bufferPost = DoubleArray(bufferFrames)
    var bufferIndex = 0

    val scratch = DoubleArray(bufferFrames + 2)
    val scratch2 = DoubleArray(bufferFrames + 2)

    val filterState = Array(2) { DoubleArray(4) }

    var samplePeak = 0f

    val blocks = ArrayList<Double>(blocksCapacity)
    val overview = ArrayList<Double>(overviewCapacity)
  }

  private val samplesIn100ms: Int = (sampleRate / 10).roundToInt()
  private val samplesIn400ms: Int = samplesIn100ms * 4

  private val channels = Array(channelCount) {
    LoudnessMeasurerChannel(samplesIn400ms, samplesIn400ms,
        10 * ceil(totalFrames / sampleRate).toInt(),
        100 * ceil(totalFrames / sampleRate).toInt())
  }

  private val coefficients = Array(2) { DoubleArray(5) }

  init {
    var f0 = 1681.974450955533
    val gain = 3.999843853973347
    var q = 0.7071752369554196

    var k = tan(PI * f0 / sampleRate)
    val vh = 10.0.pow(gain / 20.0)
    val vb = vh.pow(0.4996667741545416)

    val a0 = 1.0 + k / q + k * k

    val pb = doubleArrayOf(
        (vh + vb * k / q + k * k) / a0,
        2.0 * (k * k - vh) / a0,
        (vh - vb * k / q + k * k) / a0)
    val pa = doubleArrayOf(
        1.0,
        2.0 * (k * k - 1.0) / a0,
        (1.0 - k / q + k * k) / a0)


    f0 = 38.13547087602444
    q = 0.5003270373238773
    k = tan(PI * f0 / sampleRate)

    val rb = doubleArrayOf(1.0, -2.0, 1.0)
    val ra = doubleArrayOf(
        1.0,
        2.0 * (k * k - 1.0) / (1.0 + k / q + k * k),
        (1.0 - k / q + k * k) / (1.0 + k / q + k * k))


    coefficients[0][0] = pb[0]
    coefficients[0][1] = pb[1]
    coefficients[0][2] = pb[2]
    coefficients[0][3] = pa[1]
    coefficients[0][4] = pa[2]

    coefficients[1][0] = rb[0]
    coefficients[1][1] = rb[1]
    coefficients[1][2] = rb[2]
    coefficients[1][3] = ra[1]
    coefficients[1][4] = ra[2]
  }

  private fun biquad(inBuffer: DoubleArray, outBuffer: DoubleArray, coef: DoubleArray,
                     state: DoubleArray, frames: Int) {
    inBuffer[0] = state[0]
    inBuffer[1] = state[1]
    outBuffer[0] = state[2]
    outBuffer[1] = state[3]

    deq22(inBuffer, coef, outBuffer, frames)

    state[0] = inBuffer[frames]
    state[1] = inBuffer[frames + 1]
    state[2] = outBuffer[frames]
    state[3] = outBuffer[frames + 1]
  }

  private fun LoudnessMeasurerChannel.filter(src: FloatArray, startIndex: Int, frames: Int) {
    if (frames == 0) return

    val range = src.drop(startIndex).take(frames)
    val max = range.maxMagnitude() ?: 0f

    if (max > samplePeak) {
      samplePeak = max
    }

    src.copyInto(scratch, 2)

    repeat(frames) { i ->
      bufferPre[i + bufferIndex] = scratch[i + 2].toFloat()
    }

    biquad(scratch, scratch2, coefficients[0], filterState[0], frames)
    biquad(scratch2, scratch, coefficients[1], filterState[1], frames)

    repeat(frames) { i ->
      bufferPost[i + bufferIndex] = scratch[i + 2]
    }
  }

  private fun LoudnessMeasurerChannel.calculateOverview(framesPerBlock: Int, blockCount: Int) {
    for (b in (if (overview.isEmpty()) 0 else blockCount - 10) until blockCount) {
      val bufferIndex = (bufferIndex + framesPerBlock * b) % bufferFrames

      // bufferIndex marks the end of the block of interest. If it's 0 then we've either just
      // started or we've wrapped around. If we've just started, then the overview will be empty.
      if (overview.isEmpty() && bufferIndex == 0) {
        continue
      }

      val max = this.accumulateBlock(bufferIndex, framesPerBlock) { acc, startIndex, count ->
        acc.coerceAtLeast(bufferPre.drop(startIndex).take(count).maxMagnitude()!!.toDouble())
      }
      overview.add(max)
    }
  }

  private fun LoudnessMeasurerChannel.calcuclateGatingBlock(framesPerBlock: Int) {
    val sum = this.accumulateBlock(bufferIndex, framesPerBlock) { acc, startIndex, count ->
      if (count == 0) acc
      else acc +
          bufferPost.drop(startIndex).take(count)
              .map { it * it }
              .sum()
    }
    blocks.add(sum)
  }

  private fun LoudnessMeasurerChannel.accumulateBlock(bufferIndex: Int, framesPerBlock: Int,
                                                      accFun: (Double, Int, Int) -> Double): Double {
    var acc = 0.0
    if (bufferIndex < framesPerBlock) {
      // We have wrapped around. Check the first part of the buffer.
      if (bufferIndex > 0) {
        acc = accFun(acc, 0, bufferIndex)
      }
      // Then check the end of the buffer for the remaining frames.
      val i = bufferFrames - (framesPerBlock - bufferIndex)
      val frames = bufferFrames - i
      acc = accFun(acc, i, frames)
    } else {
      acc = accFun(acc, bufferIndex - framesPerBlock, framesPerBlock)
    }
    return acc
  }

  private fun LoudnessMeasurerChannel.process(source: FloatArray, inFrames: Int) {
    var index = 0
    var frames = inFrames

    while (frames > 0) {
      if (frames >= neededFrames) {
        filter(source, index, neededFrames)
        index += neededFrames
        frames -= neededFrames
        bufferIndex += neededFrames

        calcuclateGatingBlock(samplesIn400ms)
        calculateOverview(samplesIn100ms / 10, 40)

        // 100ms are needed for all blocks besides the first one
        neededFrames = samplesIn100ms

        if (bufferIndex == bufferFrames) {
          bufferIndex = 0
        }
      } else {
        filter(source, index, frames)

        bufferIndex += frames
        neededFrames -= frames
        frames = 0
      }
    }
  }

  @Suppress("UsePropertyAccessSyntax")
  fun scanAudioBuffer(byteBuffer: ByteBuffer, bufferChannelCount: Int) {
    var frame = 0
    while (byteBuffer.hasRemaining()) {
      channels.forEachIndexed { ch, channel ->
        channel.inputBuffer[frame] =
            if (ch >= bufferChannelCount) {
              0f
            } else {
              byteBuffer.getShort().toFraction()
            }

        if (bufferChannelCount > channelCount) {
          // Skip channels not tracked by this LoudnessMeasurer
          repeat(bufferChannelCount - channelCount) {
            byteBuffer.getShort()
          }
        }
      }
      frame++
    }
    channels.forEach { 
      it.process(it.inputBuffer, frame + 1)
    }
  }

  fun getOverview() = UByteArray(channels[0].overview.size) { i ->
    val max = channels.map { it.overview[i] }.max()!!
    floor(max * 255.0)
        .toShort()
        .coerceIn(0, 255)
        .toUByte()
  }

  fun getLoudness(): Double {
    val blocks = List(channels[0].blocks.size) { i ->
      (channels
          .map { it.blocks[i] }
          .sum()
          / samplesIn400ms)
    }.filter { it >= ABSOLUTE_GATE_FACTOR }

    if (blocks.isEmpty()) return 0.0

    val relativeThreshold = blocks.average() * RELATIVE_GATE_FACTOR
    val aboveRelativeThreshold = blocks.filter { it >= relativeThreshold }
    if (aboveRelativeThreshold.isEmpty()) return 0.0

    return 10 * log10(aboveRelativeThreshold.average()) - 0.691
  }

  fun getPeak(): Double {
    return channels.map { it.samplePeak }.max()!!.toDouble()
  }

  companion object {
    private val RELATIVE_GATE_FACTOR = 10.0.pow(-1.0)
    private val ABSOLUTE_GATE_FACTOR = 10.0.pow(-7.0691 / 10.0)
  }
}

private fun FloatArray.copyInto(dst: DoubleArray, destinationOffset: Int = 0, startIndex: Int = 0, endIndex: Int = 0) {
  val expandedSource = DoubleArray(size) { this[it].toDouble() }
  expandedSource.copyInto(dst, destinationOffset, startIndex, endIndex)
}

private fun List<Float>.maxMagnitude(): Float? {
  val max = max() ?: return null
  val min = min() ?: return null
  return max(max, -min)
}

private fun deq22(a: DoubleArray, b: DoubleArray, out: DoubleArray, length: Int) {
  for (n in 2 until length + 2) {
    out[n] = (a[n] * b[0]
        + a[n - 1] * b[1]
        + a[n - 2] * b[2]
        - out[n - 1] * b[3]
        - out[n - 2] * b[4])
  }
}