ocelot-desktop/src/main/scala/ocelot/desktop/ui/widget/Plot.scala

package ocelot.desktop.ui.widget

import ocelot.desktop.ColorScheme
import ocelot.desktop.color.Color
import ocelot.desktop.geometry.{Padding2D, Rect2D, Size2D}
import ocelot.desktop.graphics.Graphics
import ocelot.desktop.ui.UiHandler
import ocelot.desktop.ui.event.handlers.{HoverHandler, MouseHandler}
import ocelot.desktop.ui.event.{ClickEvent, MouseEvent}
import ocelot.desktop.ui.widget.Plot.Axis
import ocelot.desktop.ui.widget.contextmenu.ContextMenu
import ocelot.desktop.util.DrawUtils

abstract class Plot extends Widget {
  private var barsMin = new Array[Float](0)
  private var barsMax = new Array[Float](0)

  def xAxis: Axis

  def yAxis: Axis

  def points: (Array[Float], Array[Float])

  def fill: Boolean = false

  def setupContextMenu(menu: ContextMenu): Unit = {}

  def onPanelLeftClick(): Unit = {}

  def onPanelRightClick(): Unit = {
    val menu = new ContextMenu
    setupContextMenu(menu)
    root.get.contextMenus.open(menu)
  }

  private val panel = new Widget with HoverHandler with MouseHandler {
    override protected def receiveClickEvents: Boolean = true

    eventHandlers += {
      case ClickEvent(MouseEvent.Button.Left, _) => onPanelLeftClick()
      case ClickEvent(MouseEvent.Button.Right, _) => onPanelRightClick()
    }

    override def minimumSize: Size2D = Size2D(50, 50)

    override def maximumSize: Size2D = Size2D.Inf

    override def draw(g: Graphics): Unit = {
      resample(bounds)

      val fillColor = ColorScheme("PlotFill")
      val gridColor = ColorScheme("PlotGrid")
      val lineColor = ColorScheme("PlotLine")

      DrawUtils.panel(g, bounds.x - 2, bounds.y - 2, bounds.w + 4, bounds.h + 4)

      for (tickX <- xAxis.ticks) {
        val x = xAxis.mapPos(tickX, bounds.w).floor
        g.line(bounds.x + x, bounds.y, bounds.x + x, bounds.y + bounds.h, 1, gridColor)
      }

      for (tickY <- yAxis.ticks) {
        val y = yAxis.mapPos(tickY, bounds.h).floor + 0.5f
        g.line(bounds.x, bounds.y + y, bounds.x + bounds.w, bounds.y + y, 1, gridColor)
      }

      if (fill) {
        for ((max, x) <- barsMax.zipWithIndex) {
          g.rect(bounds.x + x, bounds.y + max, 1, bounds.h - max, fillColor)
        }
      }

      for (((min, max), x) <- barsMin.zip(barsMax).zipWithIndex) {
        g.rect(bounds.x + x, bounds.y + min, 1, (max - min).max(1), lineColor)
      }
    }
  }

  children :+= new PaddingBox(
    panel,
    Padding2D(
      left = 4,
      right = if (yAxis.ticks.isEmpty || !xAxis.tickLabels) 4 else yAxis.tickWidth * 7 + 6,
      top = 4,
      bottom = if (xAxis.ticks.isEmpty || !xAxis.tickLabels) 4 else 14,
    ),
  )

  override def draw(g: Graphics): Unit = {
    val xAxis = this.xAxis
    val yAxis = this.yAxis
    val pBounds = panel.bounds
    panel.draw(g)

    g.setSmallFont()
    g.background = Color.Transparent
    g.foreground = ColorScheme("PlotText")

    if (xAxis.tickLabels) {
      for (tickX <- xAxis.ticks) {
        val x = xAxis.mapPos(tickX, pBounds.w).floor
        val text = xAxis.tickFormatter(tickX)
        g.text(pBounds.x + x - (text.length * 3.5f).floor, pBounds.y + pBounds.h + 4, text, shrink = 1)
      }
    }

    if (yAxis.tickLabels) {
      for (tickY <- yAxis.ticks) {
        val y = yAxis.mapPos(tickY, pBounds.h).floor + 2
        val text = yAxis.tickFormatter(tickY)
        g.text(pBounds.x + pBounds.w + 4, pBounds.y + y - 4, text, shrink = 1)
      }
    }

    if (panel.mouseOver) {
      g.foreground = ColorScheme("PlotLine")
      val pos = UiHandler.mousePosition
      val x = xAxis.inversePos(pos.x - pBounds.x, pBounds.w)
      val y = yAxis.inversePos(pos.y - pBounds.y, pBounds.h)
      val text = f"${xAxis.hoverFormatter(x)}${xAxis.unit}; ${yAxis.hoverFormatter(y)}${yAxis.unit}"
      g.text(pBounds.x + 1, pBounds.y + 1, text, shrink = 1)
    }

    g.setNormalFont()
  }

  private var mappedX: Array[Float] = new Array(0)
  private var mappedY: Array[Float] = new Array(0)

  private def resample(pBounds: Rect2D): Unit = {
    val xAxis = this.xAxis
    val yAxis = this.yAxis

    val numBars = pBounds.w.toInt
    if (barsMin.length != numBars) {
      barsMin = new Array(numBars)
      barsMax = new Array(numBars)
    }

    val points = this.points
    val numPoints = points._1.length
    assert(points._1.length == points._2.length)
    if (mappedX.length != numPoints) {
      mappedX = new Array(numPoints)
      mappedY = new Array(numPoints)
    }

    var i = 0
    while (i < numPoints) {
      mappedX(i) = xAxis.mapPos(points._1(i), pBounds.w)
      mappedY(i) = yAxis.mapPos(points._2(i), pBounds.h)
      i += 1
    }

    var pointIdx = 0
    var barIdx = 0
    while (barIdx < numBars) {
      var min = Float.PositiveInfinity
      var max = Float.NegativeInfinity

      var count = 0

      while (pointIdx < numPoints && mappedX(pointIdx) <= barIdx) {
        val y = mappedY(pointIdx)
        if (y < min) min = y
        if (y > max) max = y
        pointIdx += 1
        count += 1
      }

      val savedIdx = pointIdx

      while (pointIdx < numPoints && mappedX(pointIdx) - barIdx <= 0.5f) {
        val y = mappedY(pointIdx)
        if (y < min) min = y
        if (y > max) max = y
        pointIdx += 1
        count += 1
      }

      pointIdx = savedIdx

      if (count > 0) {
        barsMin(barIdx) = min
        barsMax(barIdx) = max
      } else {
        val ax = mappedX(pointIdx - 1)
        val ay = mappedY(pointIdx - 1)
        val bx = if (pointIdx < numPoints) mappedX(pointIdx) else ax
        val by = if (pointIdx < numPoints) mappedY(pointIdx) else ay
        val t = (bx - barIdx) / (bx - ax)
        val s = ay * t + by * (1 - t)
        barsMin(barIdx) = s
        barsMax(barIdx) = s
      }

      if (barIdx >= 1) {
        if (barsMax(barIdx - 1) < barsMin(barIdx))
          barsMax(barIdx - 1) = barsMin(barIdx)
        if (barsMin(barIdx - 1) > barsMax(barIdx))
          barsMax(barIdx) = barsMin(barIdx - 1)
      }

      barIdx += 1
    }
  }
}

object Plot {
  case class Axis(min: Float, max: Float, ticks: List[Float] = List(),
                  flip: Boolean = false, logOffset: Option[Float] = None,
                  tickFormatter: Float => String = _.toString,
                  tickLabels: Boolean = true,
                  hoverFormatter: Float => String = _.toString,
                  unit: String = "") {
    lazy val tickWidth: Int = if (ticks.isEmpty || !tickLabels) 0 else ticks.map(tickFormatter(_).length).max

    def mapPos(pos: Float, size: Float): Float = {
      var t = if (logOffset.isDefined) {
        val o = logOffset.get
        val logPos = math.log(pos + o).toFloat
        val logMin = math.log(min + o).toFloat
        val logMax = math.log(max + o).toFloat
        (logPos - logMin) / (logMax - logMin)
      } else {
        (pos - min) / (max - min)
      }

      if (flip) t = 1 - t
      if (t > 1) t = 1
      if (t < 0) t = 0

      t * size
    }

    def inversePos(pos: Float, size: Float): Float = {
      var t = pos / size
      if (flip) t = 1 - t
      if (logOffset.isDefined) {
        val o = logOffset.get
        val logMin = math.log(min + o).toFloat
        val logMax = math.log(max + o).toFloat
        val logPos = t * (logMax - logMin) + logMin
        (math.exp(logPos) - o).toFloat
      } else {
        t * (max - min) + min
      }
    }
  }
}