use "collections"
use "files"
use "itertools"

// A basic two-dimensional point class
class val Point is Hashable
  let x: USize
  let y: USize

  new val create(_x: USize, _y: USize) =>
    x = _x
    y = _y

  // because we want to use this as a key in a hashmap
  fun hash(): USize =>
    // hey this is probably bad, huh?
    x.hash() + y.hash()

  fun eq(other: Point): Bool =>
    (x == other.x) and (y == other.y)

  fun ne(other: Point): Bool => not eq(other)

class Layer
  let pixels: Map[Point, U8]
  let w: USize
  let h: USize

  new create(pixels': Map[Point, U8], w': USize, h': USize) =>
    pixels = pixels'
    w = w'
    h = h'

  fun ref string(): String =>
    var buf = String()
    for y in Range(0, h) do
      for x in Range(0, w) do
        try
          match pixels(Point(x, y))?
          | '0' => buf.push(' ')
          | '1' => buf.push('X')
          | '2' => buf.push('?')
          end
        else
          buf.push('.')
        end
      end
      buf.push('\n')
    end
    buf.clone()

  fun ref count(n: U8): USize =>
    Iter[U8](pixels.values()).filter({(x) => x == n}).count()

class Image
  var layers: Array[Layer]
  new create(layers': Array[Layer]) =>
    layers = layers'

  fun ref string(): String =>
    var buf = ""
    buf = buf.add("image [\n")
    for (i, l) in Iter[Layer](layers.values()).enum[USize]() do
      buf = buf.add("  layer ")
      buf = buf.add(i.string())
      buf = buf.add(": ")
      buf = buf.add(l.string())
      buf = buf.add("\n")
    end
    buf = buf.add("]")
    buf.clone()

  fun ref composite(): Layer =>
    var result: Map[Point, U8] = Map[Point, U8]()
    for layer in layers.values() do
      for (p, v) in layer.pixels.pairs() do
        if result.get_or_else(p, '2') == '2' then
          if v != '2' then result(p) = v end
        end
      end
    end
    Layer(result, 25, 6)

actor Main
  fun mk_image(env: Env, width: USize, height: USize, file: File): Image =>
    var layers = Array[Layer]()
    var pixels = Map[Point, U8]()
    try
      while true do
        for y in Range(0, height) do
          for x in Range(0, width) do
            pixels(Point(x, y)) = file.read(1)(0)?
          end
        end
        layers.push(Layer(pixels = Map[Point, U8](), width, height))
      end
    end
    Image(layers)

  new create(env: Env) =>
    let caps = recover val FileCaps.>set(FileRead).>set(FileStat) end
    let target_file = try env.args(1)? else "input.txt" end
    try
      with
        file = OpenFile(FilePath(env.root as AmbientAuth, target_file, caps)?) as File
      do
        let w: USize = 25
        let h: USize = 6
        let image = mk_image(env, w, h, file)
        env.out.print("image is " + image.string())
        var min: (USize | None) = None
        var min_idx: (USize | None) = None
        for (idx, layer) in Iter[Layer](image.layers.values()).enum() do
          let zeroes = layer.count('0')
          let is_min = try min as USize > zeroes else true end
          env.out.print("layer " + idx.string() + ": " + zeroes.string())
          if is_min then
            min = zeroes
            min_idx = idx
          end
        end
        env.out.print("min layer " + min_idx.string() + " had " + min.string() + " zeroes")
        try
          let tgt = image.layers(min_idx as USize)?
          env.out.print("  answer: " + (tgt.count('1') * tgt.count('2')).string())
        end
        let composite = image.composite()
        env.out.print("composite:\n" + composite.string())
      end
    end


    env.out.print(".")