%useLatestDescriptors
%use dataframe
%use kotlin-statistics
%use lets-plot

LetsPlot.getInfo()

val mpgDf = DataFrame.readCSV("https://raw.githubusercontent.com/JetBrains/lets-plot-kotlin/master/docs/examples/data/mpg.csv")
mpgDf.head()


val mpgMap = mpgDf.toMap()

val regModel = (mpgMap["cty"]!! zip mpgMap["hwy"]!!).simpleRegression(
    xSelector = { it.first as Number },
    ySelector = { it.second as Number }
)
val ctyMedian = mpgDf.cty.median()
val hwyMedian = mpgDf.hwy.median()

letsPlot(mpgMap) { x = "cty"; y = "hwy" } +
    geomPoint() +
    geomVLine(xintercept = ctyMedian, color = "#756bb1", linetype = "dashed") +
    geomHLine(yintercept = hwyMedian, color = "#756bb1", linetype = "dashed") +
    geomABLine(slope = regModel.slope, intercept = regModel.intercept, color = "#de2d26") +
    themeMinimal()

fun generateParabolicDataMap(n: Int = 25, a: Double = 1.0): Map<String, List<Double>> {
    val rand = java.util.Random(42)
    val x = List(2 * n + 1) { i -> a * (i - n).toDouble() / n }
    val y = x.map { i -> i * i + rand.nextGaussian() }
    return mapOf("x" to x, "y" to y)
}

val pDataMap = generateParabolicDataMap(a = 3.0)

letsPlot(pDataMap) + geomLine() { x = "x"; y = "y" }

letsPlot(pDataMap) + geomPath() { x = "x"; y = "y" }

letsPlot(pDataMap) + geomStep() { x = "x"; y = "y" }

fun generateArchimedeanDataMap(n: Int = 25, k: Double = 1.0, a: Double = 1.0): Map<String, List<Double>> {
    val phi = List(n) { i -> 2.0 * PI * k * i.toDouble() / (n - 1) }
    val r = phi.map { angle -> (a * angle) / (2.0 * PI) }
    val x = (r zip phi).map { p -> p.first * cos(p.second) }
    val y = (r zip phi).map { p -> p.first * sin(p.second) }
    return mapOf("x" to x, "y" to y)
}

val aDataMap = generateArchimedeanDataMap(n = 200, k = 2.0)

letsPlot(aDataMap) +
    geomLine() { x = "x"; y = "y" } +
    ggtitle("Plotting with geomLine()")

letsPlot(aDataMap) +
    geomPath() { x = "x"; y = "y" } +
    ggtitle("Plotting with geomPath()")

letsPlot(generateArchimedeanDataMap(n = 50)) + geomSegment(xend = 0.0, yend = 0.0) { x = "x"; y = "y" }