Generates graph and observations for stats.railway project

package stats.railway

import util.Random

case class Graph(n: Int, observationSize: Int) {

  val vertices = 0 to (n - 1) map { i =>
    (i, (Random.nextInt(2) % 2) + 1)
  }

  val edges = {
    val m = Array.ofDim[Array[Int]](n)
    for (i <- 0 until n) m(i) = Array.ofDim[Int](3)

    for (i <- 0 until n) {
      for (j <- 0 until 3) {
        if (m(i)(j) == 0) {
          var start = Random.nextInt(n)
          while (start != i && m(start)(j) != 0) start = Random.nextInt(n)
          m(start)(j) = i
          m(i)(j) = start
        }
      }
    }

    m
  }

  val observations: Array[Int] = {
    val size = observationSize

    val path = Array.ofDim[(Int, Int)](size)

    for (i <- 0 until size) {
      if (i == 0) path(0) = (Random.nextInt(n), Random.nextInt(3))
      else {
        val (prev, prevEdge) = path(i - 1)
        val next = edges(prev)(prevEdge)
        val (_, switch) = vertices(next)

        val nextEdge = if (prevEdge == 0) switch else 0
        path(i) = (next, nextEdge)
      }
    }

    val observations = path.map(_._2).toArray
    for (i <- 0 until size) {
      val r = Random.nextInt(100)
      if (Random.nextInt(100) < 5) observations(i) = Random.nextInt(3)
    }

    observations
  }

}