%useLatestDescriptors
%use lets-plot

import kotlin.math.*

// This example was found at: 
// https://jakevdp.github.io/PythonDataScienceHandbook/04.04-density-and-contour-plots.html

fun meshgridPoints(x:List<Double>, y:List<Double>):List<Pair<Double, Double>> {
  val xSer = y.flatMap { x }
  var yInd = -1
  val ySer = y.flatMap { 
      yInd++
      List(x.size) {y[yInd]} 
  }
  return xSer.zip(ySer)
}

val gridPoints = meshgridPoints(
    generateSequence(0.0, {it + 6.0/50} ).takeWhile { it <= 5.0 }.toList(),
    generateSequence(0.0, {it + 6.0/40} ).takeWhile { it <= 5.0 }.toList())


fun f(x:Double, y:Double): Double {
    return sin(x).pow(10) + cos(10 + y * x) * cos(x)
}

val X = gridPoints.map { it.first }
val Y = gridPoints.map { it.second }
val Z = gridPoints.map { f(it.first, it.second) }

val p = letsPlot {x=X; y=Y; fill=Z}
p + geomTile() + scaleFillHue()

// Lets use greyscale and also remove colorbar.
val p1 = p + geomTile(alpha=.5) + scaleFillGrey(1, .2, guide="none")
p1

// Add contours
p1 + geomContour(color="red") {z=Z}

// Set contour color by level
p1 + geomContour {z=Z; color="..level.."} + scaleColorGradient(low="dark_green", high="yellow")

val p2 = letsPlot {x=X; y=Y}

// Filled polygons are not working very well in this case
p2 + geomPolygon(stat=Stat.contour {z=Z}) {fill="..level.."}

// 'contourf' geom works a lot better:
p2 + geomContourFilled {z=Z; fill="..level.."}