import com.raquo.laminar.api.L.{*, given}
import narr.*
import org.scalajs.dom
import org.scalajs.dom.document
import scala.collection.mutable.ArrayBuffer
import scala.math.{cos, sin}
import slash.matrix.Matrix
import slash.matrix.decomposition.Eigen

object CircularLawApp:
  val canvas = canvasTag(widthAttr := 600, heightAttr := 600)
  val ctx = canvas.ref.getContext("2d").asInstanceOf[dom.CanvasRenderingContext2D]
  
  val (eigvalSeries, runTimeReport) = randEigvalSeries[60]()
  val timeStepState = Var("0")
  
  def draw(timeStep: String): Unit =
    // center and normalize the coordinate system
    val width = canvas.ref.width
    val height = canvas.ref.height
    ctx.setTransform(1d, 0d, 0d, -1d, 0.5*width, 0.5*height)
    
    // clear the previous frame
    ctx.clearRect(-0.5*width, -0.5*width, width, height)
    
    // find the resolution
    val rDisp: Double = 1.5
    val res = width / (2*rDisp)
    
    // draw the eigenvalues
    val eigvals = eigvalSeries(timeStep.toInt)
    for n <- 0 to eigvals(0).length-1 do
      ctx.beginPath()
      ctx.arc(
        res * eigvals(0)(n),
        res * eigvals(1)(n),
        3d,
        0d, 2*math.Pi
      )
      ctx.fill()
  
  def complexEigenvalues[N <: Int](mat: Matrix[N, N])(using ValueOf[N]): (NArray[Double], NArray[Double]) =
    val eigen = Eigen(mat)
    (
      eigen.realEigenvalues.asInstanceOf[NArray[Double]],
      eigen.imaginaryEigenvalues.asInstanceOf[NArray[Double]]
    )
  
  def randEigvalSeries[N <: Int]()(using ValueOf[N]): (ArrayBuffer[(NArray[Double], NArray[Double])], String) =
    // start timing
    val startTime = System.currentTimeMillis()
    
    // initialize the random matrix step
    val dim: Int = valueOf[N]
    var randMat = new Matrix[N, N](
      NArray.tabulate(dim*dim)(k => (math.E*k*k) % 2 - 1)
    ).times(math.sqrt(3d / dim))
    
    // initialize the rotation step
    val timeRes = 100
    val maxFreq = 4
    val rotStep = Matrix.identity[N, N]
    for n <- 0 to dim by 2 do
      val ang = math.Pi * (n % maxFreq) / timeRes
      val cos_ang = cos(ang)
      val sin_ang = sin(ang)
      rotStep(n, n) = cos_ang
      rotStep(n+1, n) = sin_ang
      rotStep(n, n+1) = -sin_ang
      rotStep(n+1, n+1) = cos_ang
    
    // find the eigenvalues
    val eigvalSeries = ArrayBuffer(complexEigenvalues(randMat))
    for _ <- 1 to timeRes-1 do
      randMat = rotStep * randMat
      eigvalSeries += complexEigenvalues(randMat)
    
    // finish timing
    val runTime = System.currentTimeMillis() - startTime
    (eigvalSeries, runTime.toString() + " ms")
  
  def main(args: Array[String]): Unit =
    ctx.fillStyle = "white"
    
    lazy val app = div(
      idAttr := "app",
      div(runTimeReport),
      canvas,
      input(
        typ := "range",
        maxAttr := (eigvalSeries.length-1).toString,
        controlled(
          value <-- timeStepState.signal,
          onInput.mapToValue --> timeStepState.writer
        ),
        timeStepState.signal --> draw
      )
    )
    renderOnDomContentLoaded(document.body, app)