import java.util.*

class Graph {

    val vertices: MutableList<Cities> = mutableListOf()
    val edges: MutableList<Edge> = mutableListOf()
    val adjacencyList: MutableMap<Cities, MutableList<Pair<Cities, Int>>> = mutableMapOf()

    fun addVertex(vararg names: Cities) {
        names.forEach {
            if(!vertices.contains(it))
                vertices.add(it)
        }
    }

    fun addEdge(weight: Int, vertex1: Cities, vertex2: Cities = vertex1) {
        if(!vertices.contains(vertex1) || !vertices.contains(vertex2)) {
            error("""
                Nie znaleziono jednego lub więcej z podanych wierzchołków!
                [vertex1: $vertex1, vertex2: $vertex2, weight: $weight]
            """.trimIndent())
        }
        edges.add(Edge(weight, vertex1, vertex2))
    }

    fun showEdges() {
        edges.forEach(::println)
    }

    fun createAdjacencyList() {
        adjacencyList.clear()
        edges.forEach { edge ->
                adjacencyList.getOrPut(edge.vertex1, ::mutableListOf).add(Pair(edge.vertex2, edge.weight))
                adjacencyList.getOrPut(edge.vertex2, ::mutableListOf).add(Pair(edge.vertex1, edge.weight))
        }
    }
    fun kruskalMST(){
        val sorted = edges.toMutableList().also { it.sortBy { it.weight } }

        val mst = mutableListOf<Edge>()
        val sets = vertices.map { listOf(it).toMutableList() }.toMutableList()

        fun find(vertex: Cities): MutableList<Cities>? {
            sets.forEach {
                val result = it.contains(vertex)
                if(result)
                    return it
            }
            return null
        }

        sorted.forEach { edge ->
            val set1 = find(edge.vertex1)
            val set2 = find(edge.vertex2)
            if(set1 != set2) {
                mst.add(edge)
                set1!!.addAll(set2!!)
                sets.remove(set2)
            }
        }
        mst.forEach {
            println(it)
        }
    }
    fun primMST()
    {
        class CityPair(val city: Cities, val weight: Int) : Comparable<CityPair> {
            override fun compareTo(other: CityPair): Int {
                if (weight < other.weight)
                    return -1
                if (weight == other.weight)
                    return 0
                return 1
            }

            operator fun component1() = city
            operator fun component2() = weight
        }
        class MSTProps (
            var parent: Cities = Cities.NOT_FOUND,
            var weight: Int = Int.MAX_VALUE,
            var visited: Boolean = false
        ) {
            override fun toString() = "parent: $parent, dist: $weight, inMST?: $visited"
        }

        val mst = vertices.map { it to MSTProps() }.toMap().toMutableMap()
        val queue = PriorityQueue<CityPair>()
        val src = vertices.random()

        queue.add(CityPair(src, 0))
        mst[src]?.weight = 0

        println("Starting from: $src")

        while (!queue.isEmpty()) {

            val parent = queue.remove().city

            mst[parent]?.visited = true

            adjacencyList[parent]?.forEach { (child, weight) ->

                if (!mst[child]!!.visited && mst[child]!!.weight > weight) {
                    mst[child]?.weight = weight
                    queue.add(CityPair(child, weight))
                    mst[child]?.parent = parent
                }
            }

        }

        var distanceSum = 0
        mst.forEach { (city, info) ->
            println("city: $city, info := [$info]")
            distanceSum += info.weight
        }
        println("Total distance: $distanceSum")
    }
    fun primsAlgorithmOld() {

        createAdjacencyList()
        class CityPair(val first: Cities, val second: Int): Comparable<CityPair> {
            override fun compareTo(other: CityPair): Int {
                if(second < other.second)
                    return -1
                if(second == other.second)
                    return 0
                return 1
            }
            operator fun component1() = first
            operator fun component2() = second
        }
        val pq: Queue<CityPair> = PriorityQueue()

        val src = vertices.random()
        println("Starting from: $src")
        val key: MutableMap<Cities, Int> = vertices.map { it to Int.MAX_VALUE }.toMap().toMutableMap()
        val parent: MutableMap<Cities, Cities> = vertices.map { it to Cities.NOT_FOUND }.toMap().toMutableMap()
        val inMST: MutableMap<Cities, Boolean> = vertices.map { it to false }.toMap().toMutableMap()

        pq.add(CityPair(src, 0))

        key[src] = 0

        while(!pq.isEmpty()) {

            val (u, w) = pq.remove()

            inMST[u] = true

            adjacencyList[u]?.forEach {
                val (v, weight) = it

                if(inMST[v] == false && key[v]!! > weight) {
                    key[v] = weight
                    pq.add(CityPair(v, weight))
                    parent[v] = u
                }
            }
        }

        parent.forEach { (city, parent) ->
            println("city: $city, parent $parent")
        }

    }
}