package scalaExercises.adventOfCode2022
package Day16

import scala.io.Source
import scala.collection.mutable
import scala.collection.immutable._
import scala.util.control.Breaks._
import scala.util.Random

val problemInput = Source.fromFile("/Users/mik/IdeaProjects/scala-exercises/src/scalaExercises/adventOfCode2022/Day16/input.txt")
val testSmall = Source.fromString(
  """Valve AA has flow rate=0; tunnels lead to valves DD, II, BB
    |Valve BB has flow rate=13; tunnels lead to valves CC, AA
    |Valve CC has flow rate=2; tunnels lead to valves DD, BB
    |Valve DD has flow rate=20; tunnels lead to valves CC, AA, EE
    |Valve EE has flow rate=3; tunnels lead to valves FF, DD
    |Valve FF has flow rate=0; tunnels lead to valves EE, GG
    |Valve GG has flow rate=0; tunnels lead to valves FF, HH
    |Valve HH has flow rate=22; tunnel leads to valve GG
    |Valve II has flow rate=0; tunnels lead to valves AA, JJ
    |Valve JJ has flow rate=21; tunnel leads to valve II
    |""".stripMargin)

val valveRE = """Valve ([A-Z]{2}) has flow rate=(\d+); tunnels? leads? to valves? ([,A-Z ]+)""".r.anchored

class Valve(val flowRate : Int, var open : Boolean, val neighbors : Seq[String])

object State {
  val graph: mutable.Map[String, Valve] = mutable.Map()
  val paths: mutable.Map[String, mutable.Map[String, Seq[String]]] = mutable.Map()

  def process(line: String): Unit = line match {
    case valveRE(id, fr, tunnels) => graph.update(id, new Valve(fr.toInt, false, tunnels.split(", ")))
  }

  def findPaths(start: String): Unit = {
    if (!paths.contains(start)) paths.update(start, mutable.Map(start -> Seq[String]()))
    val heads: mutable.Queue[String] = mutable.Queue(start)
    while (heads.nonEmpty) {
      breakable {
        val head = heads.dequeue()
        val nbrs = graph(head).neighbors.filter(nb => !paths(start).contains(nb))
        for (nb <- nbrs) {
          paths(start).update(nb, paths(start)(head).appended(nb))
          heads.enqueue(nb)
        }
      }
    }
  }


  sealed trait Action
  case class Move(destination: String) extends Action
  case class Open() extends Action
  case class Minute(choice: Action, pressureRelease: Int)

  var currentPosition = "AA"
  def nextPosition(cp: String) = State.graph.filter((i, v) => !v.open).maxBy((i, v) => v.flowRate - State.paths(cp)(i).size)._1
  def releases() = State.graph.filter((i, v) => v.open).values.map(v => v.flowRate).sum
  val log: mutable.Queue[Minute] = mutable.Queue()

  def simulate(path: Seq[String]): Seq[Minute] = {
    val graph : mutable.Map[String, Valve] = this.graph.clone()
    graph.foreach((i,v) => v.open = false)
    val log: mutable.Queue[Minute] = mutable.Queue()
    var currentPosition = "AA"
    for (next <- path) {
      val release = releases()
      log.enqueueAll(
        paths(currentPosition)(next).map(p => Minute(Move(p), release))
      )
      log.enqueue(Minute(Open(), release))
      graph(next).open = true
      currentPosition = next
    }
    val release = releases()
    (log.dequeueAll(_ => true) ++ Seq.fill(30)(Minute(Move(currentPosition), release))).take(30).toSeq
  }

  def time(path : Seq[String]) : Int = {
    path.foldLeft[(String,Int)](("AA",0))((fv,t) => (t, fv._2 + paths(fv._1)(t).size+1))._2
  }

  def init(source : Source) = {
    source.getLines().foreach(process)
    valves = graph.filter((i,v) => v.flowRate > 0).map((i,v) => i).toSeq
    findPaths("AA")
    valves.foreach(findPaths)
  }

  var maxEnergy : Double = Double.NegativeInfinity
  var valves : Seq[String] = Seq()

  object simulatedAnnealing {
    def swap(i : Int, j : Int) : Seq[String] =
      valves.updated(i, valves(j)).updated(j, valves(i))

    def flip(ix : Int) : Seq[String] = swap(ix, ix+1)

    def deltaTime(newValves : Seq[String]) : Int =
      time(valves) - time(newValves)

    def deltaRelease(newValves : Seq[String]) : Int =
      simulate(valves).map(m => m.pressureRelease).sum -
        simulate(newValves).map(m => m.pressureRelease).sum

    def energy(path : Seq[String]) : Double =
      simulate(path).take(30).map(m => m.pressureRelease).sum.toDouble

    val T0 : Double = 100.0
    val dT : Double = -0.01
    val Tlo : Double= 0.01
    val Tsteps : Int = 5
    def temperature(k : Int) : Double =
      math.max(Tlo, T0 + (k/Tsteps) * dT)

    def accept(newValves : Seq[String], k : Int) : Double =
      if(energy(newValves) > energy(valves)) 1.0
      else math.exp(-(energy(valves)-energy(newValves))/temperature(k))

    def step(k : Int) : Seq[String] = {
      val i = Random.between(0, valves.size)
      val j = Random.between(0, valves.size)
      val newValves = swap(i,j)

      maxEnergy = math.max(maxEnergy, energy(valves))

      val P = accept(newValves, k)
      if(k % 1000 == 0)
        println(s"""Swap ${valves.mkString("-")} to ${newValves.mkString("-")}
  k: ${k/1000},  T: ${temperature(k)},  E: ${energy(valves)} to ${energy(newValves)},  P: $P
        """)
      if(math.random() < P) newValves
      else valves
    }
  }
}


// To use this:
// import scalaExercises.adventOfCode2022.Day16._
// State.init(problemInput)
// (1 to 150_000).toList.foreach(k => State.valves = State.simulatedAnnealing.step(k))
// State.maxEnergy