import Foundation

struct Coord : Hashable, CustomStringConvertible {
    let (x, y): (Int, Int)
    var description: String { return "(\(x), \(y))" }
}

struct Queue<T> {
    var arr: [T?]
    var front: Int = 0
    var back: Int = 0
    init(size: Int) {
        arr = Array(repeating: nil, count: size)
    }
    mutating func push(_ element: T) {
        arr[back] = element
        back += 1
    }
    mutating func pop() -> T? {
        if front < back {
            front += 1
            return arr[front-1]
        }
        return nil
    }
}

struct Maze : CustomStringConvertible {
    let (w, h): (Int, Int)
    let walls: Set<Coord>
    let (start, end): (Coord, Coord)
    let cheatMax: Int
    var description: String {
        var s = Array(repeating: Array(repeating: "  ", count: w), count: h)
        walls.forEach { wall in
            s[wall.y][wall.x] = "██"
        }
        path.forEach { cell, cost in
            s[cell.y][cell.x] = String(format: "%2d", cost%100)
        }
        let hdr = "   " + (0..<w).map { " \($0 % 10)" }.joined() + "\n"
        return hdr +
            s.enumerated().map { i, c in
                String(format: "%2d ", i) + c.joined() + "\n"
            }.joined() + "\(start), \(end)"
    }
    var path: [Coord: Int] {
        var cost: [Coord: Int] = [start: 0]
        var cell = start
        var prev = start
        var currentCost = 0
        while cell != end {
            let next = neighbors(of: cell)
                .filter(isPath)
                .filter { $0 != prev }[0]
            currentCost += 1
            cost[next] = currentCost 
            prev = cell
            cell = next
        }
        return cost
    }

    func valid(_ cell: Coord) -> Bool {
        return cell.x >= 0 && cell.x < w && cell.y >= 0 && cell.y < h
    }
    func isPath(_ cell: Coord) -> Bool {
        return !walls.contains(cell)
    }
    func isWall(_ cell: Coord) -> Bool {
        return walls.contains(cell)
    }

    func neighbors(of cell: Coord) -> [Coord] {
        return [(-1, 0), (1, 0), (0, -1), (0, 1)]
            .map { Coord(x: cell.x + $0, y: cell.y + $1) }.filter(valid)
    }

    /*
    func cheats(from entryCell: Coord) -> [Coord: Int] {
        var q = Queue<(Coord, Int)>(size: cheatMax*cheatMax*2*2)
        q.push((entryCell, 0))
        var seen: Set<Coord> = [entryCell]
        var jumps: [Coord: Int] = [:]
        while let (cell, dist) = q.pop() {
            neighbors(of: cell).filter { !seen.contains($0) }.forEach { nb in
                seen.insert(nb)
                if isPath(nb) {
                    if dist + 1 <= cheatMax {
                        //print("\(cell) -> \(nb) . \(dist + 1)")
                        jumps[nb] = min(jumps[nb] ?? Int.max, dist + 1)
                    }
                } else if dist < cheatMax {
                    //print("\(cell) -> \(nb) # \(dist + 1)")
                    q.push((nb, dist + 1))
               }
            }
        }
        return jumps
    }
    */
    func cheats(from entry: Coord) -> [Coord: Int] {
        let jumpList: [(Int, Int)] = (2...cheatMax).flatMap { dist in
                (0..<dist).flatMap { d in
                    let asdf: [(Int, Int)] = [
                        (d, dist-d), (dist-d, -d), (-d, d-dist), (d-dist, d)
                    ]
                    return asdf
                }
            }
        let jumpCoords: [Coord] = jumpList
            .map { dx, dy in Coord(x: entry.x + dx, y: entry.y + dy) }
            .filter(valid).filter(isPath)
        var res: [Coord: Int] = [:]
        jumpCoords.forEach { end in
            res[end] = abs(end.x - entry.x) + abs(end.y - entry.y)
        }
        return res
    }

    func cheats() -> [(Coord, Coord, Int, Int)] {
        var cost: [Coord: Int] = [start: 0]
        var jumps: [(Coord, Coord, Int)] = []
        var cell = start
        var prev = start
        var currentCost = 0
        while cell != end {
            // list potential cheats from current cell
            let newCheats = cheats(from: cell)
                .filter { dest, _ in !cost.keys.contains(dest) }
                .map { dest, dist in (cell, dest, dist) }
            jumps.append(contentsOf: newCheats)
            // add cost to table
            let next = neighbors(of: cell)
                .filter(isPath)
                .filter { $0 != prev }[0]
            currentCost += 1
            cost[next] = currentCost 
            prev = cell
            cell = next
            //print("looking at \(cell) \(cost[cell]!)")
        }
        return jumps.map { from, to, dist in (from, to, cost[to]! - cost[from]! - dist, dist) }
    }
}

func readInput(_ filePath: String, _ cheatMax: Int) throws -> Maze {
    let content = try String(contentsOfFile: filePath, encoding: .ascii)
    let lines = content.split(separator: "\n")
    let (w, h) = (lines[0].count, lines.count)
    var walls: Set<Coord> = []
    var (start, end) = (Coord(x: 0, y: 0), Coord(x: w-1, y: h-1))
    lines.enumerated().forEach { i, line in
        line.enumerated().forEach { j, cell in
            let coord = Coord(x: j, y: i)
            if cell == "#" {
                walls.insert(coord)
            } else if cell == "S" {
                start = coord
            } else if cell == "E" {
                end = coord
            }
        }
    }
    return Maze(
        w: w, h: h, walls: walls, start: start, end: end, cheatMax: cheatMax
    )
}

let maze = try readInput(CommandLine.arguments[1], Int(CommandLine.arguments[2])!)
print(maze)
// print(maze.cheats(from: maze.start))
let cheats = maze.cheats().filter { _, _, saving, _ in saving > 0 }
let hist = cheats.reduce(into: [:]) { n, t in n[t.2, default: 0] += 1 }
hist.keys.sorted().forEach { print("\(hist[$0]!) - \($0)") }
print(cheats.filter { $0.2 >= 100 }.count)

/*
cheats
    .filter { _, _, saving, _ in saving == 72 }
    .forEach { from, to, save, dist in
        print("\(path[from]!) -> \(path[to]!) = \(dist)(\(save))")
    }
print()
*/