%useLatestDescriptors
%use lets-plot

LetsPlot.getInfo()

fun linspace(start: Double, stop: Double, num: Int): DoubleArray {
    return DoubleArray(num) { i -> start + i * (stop - start) / (num - 1) }
}

fun meshgrid(x: DoubleArray, y: DoubleArray): Pair<DoubleArray, DoubleArray> {
    val X = DoubleArray(x.size * y.size) { i -> x[i % x.size] }
    val Y = DoubleArray(x.size * y.size) { i -> y[i / x.size] }
    return Pair(X, Y)
}
  
fun gradient(Z: DoubleArray, d: Double, n: Int): Pair<DoubleArray, DoubleArray> {
    val matrix = Z.toList().chunked(n)
    val dY = mutableListOf<Double>()
    val dX = mutableListOf<Double>()
    for (i in 0 until n) {
        for (j in 0 until n) {
            val y = if (i > 0) (matrix[i][j] - matrix[i - 1][j]) / d else 0.0
            val x = if (j > 0) (matrix[i][j] - matrix[i][j - 1]) / d else 0.0
            dY += y
            dX += x
        }
    }
    return Pair(dY.toDoubleArray(), dX.toDoubleArray())
}

fun getData(n: Int, a: Double, b: Double, f: (DoubleArray, DoubleArray) -> DoubleArray): Map<String, Any> {
    val d = (b - a) / (n - 1)
    
    val xrange = linspace(a, b, n)
    val yrange = linspace(a, b, n)
    
    val (x, y) = meshgrid(xrange, yrange)
    val z = f(x, y)

    val (dY, dX) = gradient(z, d, n)
    val r = DoubleArray(dX.size) { index -> sqrt(dX[index].pow(2) + dY[index].pow(2)) }
    val rMax = r.max()
    val radius = DoubleArray(r.size) { index -> r[index] / rMax * d }
    val angle = DoubleArray(dY.size) { index -> atan2(dY[index], dX[index]) }
    
    return mapOf(
        "x" to x, 
        "y" to y, 
        "z" to z, 
        "radius" to radius, 
        "angle" to angle
    )
}

val data = getData(n = 21, a = -2*PI, b = 2*PI) { xArray, yArray ->
    DoubleArray(xArray.size) { index -> sin(xArray[index]) + cos(yArray[index]) }
}

val p = letsPlot(data) { x = "x"; y = "y"  } + 
    coordFixed() + themeVoid() +
    scaleViridis(listOf("fill", "color"))

p + geomBin2D(stat = Stat.identity) { fill = "z" }

p + geomPoint(size = 1.5) + geomSpoke { angle = "angle"; radius = "radius"; color = "z" }

fun getPlot(pivot: String): org.jetbrains.letsPlot.intern.Plot {
    val a = -3.0
    val b = 3.0 
    val r = 0.75
    val pivotData = getData(
        n = 4, 
        a = a,
        b = b
    ) { xArray, yArray -> DoubleArray(xArray.size) { index -> xArray[index].pow(2) + yArray[index].pow(2) } }
    val title = "pivot=${pivot}" + (" (default)".takeIf { pivot == "tail" } ?: "")
    return ggplot(pivotData) { x = "x"; y = "y" } +
        geomPoint() +
        geomSpoke(radius = r, pivot = pivot) { angle = "angle" } +
        coordFixed() +
        xlim(a - r to b + r) + ylim(a - r to b + r) +
        ggtitle(title) +
        themeVoid() + theme(plotTitle = elementText(hjust = 0.5))
}

gggrid(
    listOf(getPlot("tail"), getPlot("mid"), getPlot("tip")),
    ncol=3
)