Flowchart.scala

package hod.btlg.flowchart

import java.awt.{Color, Dimension, Graphics, GridLayout}
import javax.swing.*
import scala.collection.mutable
import scala.collection.mutable.{ArrayBuffer, ListBuffer}
import scala.util.Random

object DemoGraph {

  val random = new Random(123445)
  var counter = 0

  def biasedRandomInt(max: Int): Int = {
    val floatRandom = random.nextFloat()
    val biasedRandom = Math.sqrt(floatRandom)
    (biasedRandom * max).toInt + 1
  }

  def buildRandomTree(depth: Int): Node = {
    if (depth == 0) {
      val node = Node(counter)
      counter += 1
      node
    } else {
      val numChildren = biasedRandomInt(2)
      val childDepth = Random.nextInt(depth / 2 + 1) + depth / 2
      val children = List.fill(numChildren)(buildRandomTree(childDepth))
      val node = Node(counter, children)
      counter += 1
      node
    }
  }

  val root = buildRandomTree(12)

}

case class Connection(startNode: Node, endNode: Node)

case class GraphModel(allNodes: List[Node], connections: List[Connection])

case class Node(nodeId: Int, children: List[Node] = Nil, var x: Double = 0, var y: Double = 0, var visible: Boolean = false) {
  lazy val leafCount: Int = if (children.isEmpty) 1 else children.iterator.map(_.leafCount).sum
}

object GraphModel {

  def init(root: Node): GraphModel = {
    val allNodes = new ListBuffer[Node]()
    val connections = new ListBuffer[Connection]()
    val queue = mutable.Queue[(Node, Double, Double, Double)]() // (node, xMin, xMax, y)
    val verticalSpacing = 250
    val horizontalSpacing = root.leafCount * 30

    queue.enqueue((root, 0, horizontalSpacing, 0)) // root at level 0, y=0, x in range [0, horizontalSpacing]

    while (queue.nonEmpty) {
      val (node, xMin, xMax, y) = queue.dequeue()

      node.x = (xMax + xMin) / 2.0
      node.y = y

      allNodes += node

      val sliceWidth = (xMax - xMin) / node.leafCount
      var childXMin = xMin

      for (child <- node.children) {
        val childXMax = childXMin + sliceWidth * child.leafCount

        queue.enqueue((child, childXMin, childXMax, y + verticalSpacing))
        connections += Connection(node, child)

        childXMin = childXMax
      }
    }
    println(s"graph size ${allNodes.size}")
    GraphModel(allNodes.toList, connections.toList)
  }
}


class GraphBuilder(root: Node) extends Iterator[GraphModel] {
  private val nodes = mutable.ArrayBuffer(root)
  private var currentState: Option[GraphModel] = Some(GraphModel.init(createVisibleVersion(root)))

  root.visible = true // root node is visible initially

  private def createVisibleVersion(node: Node): Node = {
    val visibleChildren = node.children.filter(_.visible)
    node.copy(children = visibleChildren.map(createVisibleVersion), visible = node.visible)
  }

  override def hasNext: Boolean = nodes.nonEmpty

  override def next(): GraphModel = currentState match {
    case None => throw new NoSuchElementException("No more nodes to add.")
    case Some(state) =>
      if (nodes.isEmpty) {
        currentState = None
      } else {
        // Randomly select a parent node and remove it from the buffer
        val randomIndex = scala.util.Random.nextInt(nodes.length)
        val parentNode = nodes.remove(randomIndex)

        // Separate child nodes into visible and invisible lists
        val (visibleChildren, invisibleChildren) = parentNode.children.partition(_.visible)

        // If there are invisible children, make them visible and add them to the queue
        if (invisibleChildren.nonEmpty) {
          for (child <- invisibleChildren) {
            // The chosen node starts at the location of its parent.
            child.x = parentNode.x
            child.y = parentNode.y

            child.visible = true // make this child visible
            nodes += child
          }
        }

        // Generate a new graph state
        currentState = Some(GraphModel.init(createVisibleVersion(root)))
      }

      state
  }
}

case class IntermediateGraphModel(before: GraphModel, after: GraphModel) {
  def toGraphModel(percentage: Double): GraphModel = {
    val beforeNodeMap = before.allNodes.map(node => node.nodeId -> node).toMap
    val afterNodeMap = after.allNodes.map(node => node.nodeId -> node).toMap

    def findBeforeNode(nodeId: Int): Node = {
      beforeNodeMap.getOrElse(nodeId, findParent(afterNodeMap(nodeId)).getOrElse(afterNodeMap(nodeId)))
    }

    def createIntermediateNode(afterNode: Node): Node = {
      val beforeNode = findBeforeNode(afterNode.nodeId)
      // Calculate the intermediate coordinates based on the progress percentage
      val x = beforeNode.x + (afterNode.x - beforeNode.x) * percentage
      val y = beforeNode.y + (afterNode.y - beforeNode.y) * percentage
      val children = afterNode.children.map(createIntermediateNode)

      Node(afterNode.nodeId, children, x, y, afterNode.visible)
    }

    val intermediateNodes = after.allNodes.map(createIntermediateNode)
    val intermediateConnections: List[Connection] = for {
      afterConn <- after.connections
      startNode <- intermediateNodes.find(_.nodeId == afterConn.startNode.nodeId)
      endNode <- intermediateNodes.find(_.nodeId == afterConn.endNode.nodeId)
    } yield Connection(startNode, endNode)

    GraphModel(intermediateNodes, intermediateConnections)
  }

  private def findParent(node: Node): Option[Node] = {
    (for {
      connection <- after.connections
      if connection.endNode.nodeId == node.nodeId
      parent <- before.allNodes.find(_.nodeId == connection.startNode.nodeId)
    } yield parent).headOption
  }
}


class FlowChartPanel(root: Node) extends JPanel {
  private val graphBuilder = new GraphBuilder(root)
  private var graphModel: GraphModel = graphBuilder.next()
  private var intermediateGraphModel: Option[IntermediateGraphModel] = None
  private var intermediateStep: Int = 0
  private val totalIntermediateSteps: Int = 100

  override def getPreferredSize: Dimension = Dimension(500, 500)

  override def paintComponent(g: Graphics): Unit = {
    super.paintComponent(g)

    val modelToDraw = intermediateGraphModel.map(_.toGraphModel(intermediateStep.toDouble / totalIntermediateSteps)).getOrElse(graphModel)
    //val modelToDraw = graphModel

    g.setColor(Color.BLACK)
    g.fillRect(0, 0, getWidth, getHeight)

    val minX = modelToDraw.allNodes.minBy(_.x).x
    val maxX = modelToDraw.allNodes.maxBy(_.x).x
    val minY = modelToDraw.allNodes.minBy(_.y).y
    val maxY = modelToDraw.allNodes.maxBy(_.y).y

    val graphWidth = maxX - minX
    val graphHeight = maxY - minY

    val panelWidth = getWidth.toDouble
    val panelHeight = getHeight.toDouble

    val scaleX = panelWidth / graphWidth
    val scaleY = panelHeight / graphHeight

    val scale = Math.min(scaleX, scaleY)

    // Drawing connections
    for (connection <- modelToDraw.connections) {
      g.setColor(Color.WHITE)
      g.drawLine(
        ((connection.startNode.x - minX) * scale).toInt,
        ((connection.startNode.y - minY) * scale).toInt,
        ((connection.endNode.x - minX) * scale).toInt,
        ((connection.endNode.y - minY) * scale).toInt)
    }

    // Drawing nodes
    for (node <- modelToDraw.allNodes) {
      val r = (15 * scale).toInt
      val x = ((node.x - minX) * scale).toInt
      val y = ((node.y - minY) * scale).toInt

      g.setColor(Color.RED)
      g.fillOval(x - r, y - r, 2 * r, 2 * r)
      g.setColor(Color.WHITE)
      g.drawString(node.nodeId.toString, x - (r / 2), y + (r / 2))
    }
  }

  def toNextStep(): Unit = {
    if (intermediateGraphModel.isDefined && intermediateStep < totalIntermediateSteps) {
      intermediateStep += 1
      repaint()
    } else if (graphBuilder.hasNext) {
      intermediateStep = 0
      val nextModel = graphBuilder.next()
      intermediateGraphModel = Some(IntermediateGraphModel(graphModel, nextModel))
      graphModel = nextModel
      repaint()
    }
  }
}

@main def showDemo(): Unit = {
  val panel = new FlowChartPanel(DemoGraph.root)
  val frame = new JFrame("FlowChart")

  frame.setDefaultCloseOperation(WindowConstants.EXIT_ON_CLOSE)
  frame.getContentPane.setLayout(GridLayout(1, 1))
  frame.getContentPane.add(panel)
  frame.pack()
  frame.setVisible(true)

  // Create a timer that triggers the toNextStep method every 100 milliseconds (10 frames per second)
  val timer = new Timer(1, e => {
    panel.toNextStep()
  })
  timer.start()

}