Problem Solving with Algorithms and Data Structures using Kotlin The Interactive Edition

import java.awt.image.BufferedImage
import javax.swing.JFrame
import javax.swing.WindowConstants

fun simpleQuant(img: BufferedImage) {
    for (row in 0..<img.height) {
        for (col in 0..<img.width) {
            val argb = img.getRGB(col, row)
            var r = (argb shr 16) and 0xff
            var g = (argb shr 8) and 0xff
            var b = argb and 0xff
            r = r / 36 * 36
            g = g / 42 * 42
            b = b / 42 * 42
            img.setRGB(col, row, (r shl 16 or (g shl 8) or b))
        }
    }
}

fun main() {
    print("File name: ")
    val filename = readln()
    val f = JFrame("Simple Quantize")
    f.defaultCloseOperation = WindowConstants.DISPOSE_ON_CLOSE
    val img = ImageComponent(filename)
    simpleQuant(img.image)
    f.add(img)
    f.pack()
    f.isVisible = true
}

import java.awt.Dimension
import java.awt.Graphics
import java.awt.image.BufferedImage
import java.io.File
import javax.imageio.ImageIO
import javax.swing.JComponent

class ImageComponent(filename: String) : JComponent() {
    var image: BufferedImage = ImageIO.read(File(filename))

    override fun paint(g: Graphics) {
        g.drawImage(image, 0, 0, null)
    }

    override fun getPreferredSize(): Dimension {
        return Dimension(
            image.getWidth(null),
            image.getHeight(null)
        )
    }
}

import java.awt.image.BufferedImage

fun octreeQuant(img: BufferedImage) {
    val w = img.width
    val h = img.height

    val ot: Octree = Octree()
    for (row in 0..<h) {
        for (col in 0..<w) {
            val argb = img.getRGB(col, row)
            val r = (argb shr 16) and 0xff
            val g = (argb shr 8) and 0xff
            val b = argb and 0xff
            ot.insert(r, g, b)
        }
    }

    ot.reduce(256)

    for (row in 0..<h) {
        for (col in 0..<w) {
            val argb = img.getRGB(col, row)
            val r = (argb shr 16) and 0xff
            val g = (argb shr 8) and 0xff
            val b = argb and 0xff

            val newRgb = ot.find(r, g, b)!!

            img.setRGB(
                col, row, (newRgb.red shl 16 or (
                        newRgb.green shl 8) or newRgb.blue)
            )
        }
    }
}

data class Pixel(var red: Int, var green: Int, var blue: Int)

class Octree {
    var root: OTNode? = null
    var numLeaves = 0
    var maxLevel = 5
    var allLeaves: MutableList<OTNode> = mutableListOf()

    fun insert(r: Int, g: Int, b: Int) {
        if (this.root == null) {
            this.root = OTNode(outer = this)
        }
        this.root!!.insert(r, g, b, 0, this)
    }

    fun find(r: Int, g: Int, b: Int): Pixel? {
        if (this.root != null) {
            return this.root!!.find(r, g, b, 0)
        } else {
            return null
        }
    }

    fun reduce(maxCubes: Int) {
        System.err.printf(
            "Reducing %d to %d%n",
            this.allLeaves.size, maxCubes
        )
        while (this.allLeaves.size > maxCubes) {
            val smallest = this.findMinCube()
            smallest!!.parent!!.merge()
            this.allLeaves.add(smallest.parent)
            this.numLeaves += 1
        }
    }

    fun findMinCube(): OTNode? {
        var minCount = Int.Companion.MAX_VALUE
        var maxLevel = 0
        var minCube: OTNode? = null

        for (node in allLeaves) {
            if (node.count <= minCount && node.level >= maxLevel) {
                minCube = node
                minCount = node.count
                maxLevel = node.level
            }
        }
        return minCube
    }
}

class OTNode(
    val parent: OTNode? = null,
    val level: Int = 0,
    val outer: Octree? = null
) {
    var red = 0
    var green = 0
    var blue = 0
    var count = 0
    val children: MutableList<OTNode?> = MutableList(8) { null }

fun insert(r: Int, g: Int, b: Int, level: Int, outer: Octree) {
    if (level  this.outer!!.maxLevel) {
        val index = computeIndex(r, g, b, level)
        if (this.children[index] == null) {
            this.children[index] = OTNode(this, level + 1, outer)
        }
        this.children[index]!!.insert(r, g, b, level + 1, outer)
    } else {
        if (this.count == 0) {
            this.outer.numLeaves = this.outer.numLeaves + 1
            this.outer.allLeaves.add(this)
        }
        this.red += r
        this.green += g
        this.blue += b
        this.count += 1
    }
}

fun computeIndex(r: Int, g: Int, b: Int, level: Int): Int {
    val nShift = 8 - level
    val rBits = (r shr (nShift - 2)) and 0x04
    val gBits = (g shr (nShift - 1)) and 0x02
    val bBits = (b shr nShift) and 0x01
    return rBits or gBits or bBits
}

fun find(r: Int, g: Int, b: Int, level: Int): Pixel {
    if (level < this.outer!!.maxLevel) {
        val index = computeIndex(r, g, b, level)
        if (this.children[index] != null) {
            return this.children[index]!!.find(r, g, b, level + 1)
        } else if (this.count > 0) {
            return Pixel(
                this.red / this.count,
                this.green / this.count, this.blue / this.count
            )
        } else {
            throw Exception("No leaf node to represent RGB($r, $g, $b)")
        }
    } else {
        return Pixel(
            this.red / this.count,
            this.green / this.count, this.blue / this.count
        )
    }
}

fun merge() {
    for (child in this.children) {
        if (child != null) {
            if (child.count > 0) {
                val removed = this.outer!!.allLeaves.remove(child)
                this.outer.numLeaves -= 1
            } else {
                println("Recursively merging non-leaf")
                child.merge()
            }

            this.count += child.count
            this.red += child.red
            this.green += child.green
            this.blue += child.blue
        }
    }

    for (i in 0 ..< 8) {
        this.children[i] = null
    }
}