package ocelot.desktop.ui.widget

import ocelot.desktop.geometry.Size2D
import ocelot.desktop.ui.UiHandler
import ocelot.desktop.ui.widget.contextmenu.{ContextMenu, ContextMenuEntry, ContextMenuSubmenu}
import org.jtransforms.fft.FloatFFT_1D
import totoro.ocelot.brain.Settings

import scala.collection.immutable.ArraySeq
import scala.collection.mutable

class Oscilloscope(isTiny: Boolean = false) extends Widget {
  private val queue = new mutable.ArrayDeque[(Array[Float], Int)]

  private val sampleRate = Settings.get.soundCardSampleRate
  private val wavePeriods = 4
  private var fftSize = 8192
  private var smoothing: Option[Float] = None

  private var fft: FloatFFT_1D = _
  private var fftBuffer0: Array[Float] = _
  private var fftBuffer1: Array[Float] = _

  private var spectrumPointsX: Array[Float] = _
  private var spectrumPointsY: Array[Float] = _
  private var spectrumPointsPrevY: Array[Float] = _

  private var wavePointsX: Array[Float] = _
  private var wavePointsY: Array[Float] = _

  setFFTSize(8192)

  private class SpectrumPlot extends Plot {
    private lazy val ticksLarge = List(50f, 100f, 200f, 300f, 500f, 1000f, 2000f, 3000f, 5000f, 10000f, 20000f)
    private lazy val ticksSmall = List(50f, 200f, 500f, 2000f, 5000f, 20000f)

    private lazy val xAxes = List(ticksLarge, ticksSmall).map(ticks => {
      Plot.Axis(
        0,
        22000,
        ticks,
        tickLabels = !isTiny,
        tickFormatter = x => if (x >= 1000) f"${x / 1000}%.0fk" else f"$x%.0f",
        hoverFormatter = x => if (x >= 1000) f"${x / 1000}%.2fk" else f"$x%.1f",
        logOffset = Some(120),
        unit = "Hz",
      )
    })

    override def xAxis: Plot.Axis = if (width >= 600) xAxes.head else xAxes.last

    override lazy val yAxis: Plot.Axis = Plot.Axis(
      -150,
      10,
      ticks = List(0f, -20f, -40f, -60f, -80f, -100f, -120f, -140f),
      tickLabels = !isTiny,
      tickFormatter = y => f"$y%+.0f",
      hoverFormatter = y => f"$y%+.1f",
      flip = true,
      unit = "dB",
    )

    override def points: (Array[Float], Array[Float]) = (spectrumPointsX, spectrumPointsY)

    override def fill: Boolean = true

    override def setupContextMenu(menu: ContextMenu): Unit = {
      addFFTSizeEntry(menu)
      addSmoothingEntry(menu)

      super.setupContextMenu(menu)
    }

    override def onPanelLeftClick(): Unit = updateMode(i => if (i == 0) 2 else 0)
  }

  private class WavePlot extends Plot {
    override lazy val xAxis: Plot.Axis = Plot.Axis(0, wavePeriods,
      ticks = List(0.5f, 1f, 1.5f, 2f, 2.5f, 3f, 3.5f),
      tickLabels = !isTiny,
      tickFormatter = x => f"${2 * x}%.0fπ",
      hoverFormatter = x => f"${2 * x}%.1fπ")

    override lazy val yAxis: Plot.Axis = Plot.Axis(-1.1f, 1.1f,
      ticks = List(-1f, -0.5f, 0f, 0.5f, 1f),
      tickLabels = !isTiny,
      tickFormatter = y => f"$y%+.1f",
      hoverFormatter = y => f"$y%+.2f",
      flip = true)

    override def points: (Array[Float], Array[Float]) = (wavePointsX, wavePointsY)

    override def setupContextMenu(menu: ContextMenu): Unit = {
      addFFTSizeEntry(menu)
      super.setupContextMenu(menu)
    }

    override def onPanelLeftClick(): Unit = updateMode(i => if (i == 1) 2 else 1)
  }

  private val modes: List[() => Widget] = List(
    () => new SpectrumPlot,
    () => new WavePlot,
    () => {
      new Widget {
        children :+= new SpectrumPlot
        children :+= new WavePlot
      }
    },
  )

  private var modeIdx = 2
  children :+= modes(modeIdx)()

  private def updateMode(f: Int => Int): Unit = {
    modeIdx = f(modeIdx)
    children = ArraySeq(modes(modeIdx)())
  }

  override def minimumSize: Size2D = if (isTiny) Size2D(100, 68) else Size2D(500, 120)

  private def addFFTSizeEntry(menu: ContextMenu): Unit = {
    menu.addEntry(new ContextMenuSubmenu("FFT Size") {
      for (v <- List(512, 1024, 2048, 4096, 8192, 16384, 32768)) {
        val dot = if (v == fftSize) '•' else ' '
        addEntry(ContextMenuEntry(s"$v".padTo(14, ' ') + dot) { setFFTSize(v) })
      }
    })
  }

  private def addSmoothingEntry(menu: ContextMenu): Unit = {
    menu.addEntry(new ContextMenuSubmenu("Smoothing (LP)") {
      for (v <- List(None, Some(0.1f), Some(0.3f), Some(0.5f), Some(0.7f), Some(0.9f))) {
        val dot = if (v == smoothing) '•' else ' '
        val msg = if (v.isEmpty) "None" else s"${v.get}"
        addEntry(ContextMenuEntry(msg.padTo(14, ' ') + dot) { smoothing = v })
      }
    })
  }

  private def setFFTSize(fftSize: Int): Unit = {
    this.fftSize = fftSize
    fft = new FloatFFT_1D(fftSize)
    fftBuffer0 = new Array(fftSize)
    fftBuffer1 = new Array(fftSize)
    spectrumPointsX = new Array(fftSize)
    spectrumPointsY = new Array(fftSize)
    spectrumPointsPrevY = new Array(fftSize)
    wavePointsX = new Array(fftSize)
    wavePointsY = new Array(fftSize)
  }

  def enqueue(samples: Array[Float]): Unit = {
    queue.synchronized {
      queue.append((samples, 0))
    }
  }

  override def update(): Unit = {
    super.update()

    for (i <- fftBuffer0.indices)
      fftBuffer0(i) = 0f

    queue.synchronized {
      var offset = 0
      var bufIdx = 0
      while (offset < fftSize && bufIdx < queue.length) {
        var (buf, bufPos) = queue(bufIdx)
        val take = (fftSize - offset).min(buf.length - bufPos)
        for (i <- bufPos until bufPos + take) {
          fftBuffer0(offset) = buf(i)
          offset += 1
        }

        bufPos += take
        if (bufPos < buf.length) {
          bufIdx += 1
        }
      }

      val advance = (UiHandler.dt * sampleRate).toInt

      offset = 0
      while (offset < advance && queue.nonEmpty) {
        var (buf, bufPos) = queue.removeHead()
        val take = (advance - offset).min(buf.length - bufPos)
        offset += take
        bufPos += take
        if (bufPos < buf.length) {
          queue.prepend((buf, bufPos))
        }
      }
    }

    fftBuffer0.copyToArray(fftBuffer1)

    // apply Hann window

    for (i <- fftBuffer0.indices) {
      val w = math.sin(math.Pi * i / fftSize)
      fftBuffer0(i) *= (w * w).toFloat
    }

    fft.realForward(fftBuffer0)

    var maxFreq = 0f
    var maxFreqMag = 0f
    var maxFreqPhase = 0f

    for (i <- 0 until fftSize / 2) {
      val re = fftBuffer0(2 * i)
      val im = if (i == 0) 0f else fftBuffer0(2 * i + 1)
      val mag = math.sqrt(re * re + im * im).toFloat
      val freq = i.toFloat * sampleRate / fftSize
      if (mag > maxFreqMag || (maxFreq < 25 && freq >= 25)) {
        maxFreq = freq
        maxFreqMag = mag
        maxFreqPhase = math.atan2(im, re).toFloat
      }
      spectrumPointsX(i) = freq
      spectrumPointsY(i) = magToDb(mag).max(-150)
    }

    val wavelength = sampleRate / maxFreq
    val offset = (wavelength * (2.0 - maxFreqPhase / math.Pi) * 0.5).toInt
    val length = (wavelength * wavePeriods).min(fftSize).toInt
    var samples = fftBuffer1.slice(offset, offset + length)

    if (samples.isEmpty) samples = Array.fill(length)(0f)

    wavePointsX = samples.indices.map(i => wavePeriods * i.toFloat / length).toArray
    wavePointsY = samples

    for (a <- smoothing) {
      applySmoothing(a, spectrumPointsPrevY, spectrumPointsY)
    }

    val tmp = spectrumPointsPrevY
    spectrumPointsPrevY = spectrumPointsY
    spectrumPointsY = tmp
  }

  private def magToDb(mag: Float): Float = {
    20 * math.log10(2 * mag / fftSize * math.sqrt(8.0 / 3)).toFloat
  }

  private def applySmoothing(a: Float, prevArr: Array[Float], curArr: Array[Float]): Unit = {
    prevArr.iterator.zip(curArr.iterator).zipWithIndex.foreach {
      case ((prev, cur), i) =>
        curArr(i) = prev * a + cur * (1 - a)
    }
  }
}
