Fast2048.swift

import Foundation

struct FastRandomGenerator {
    static var shared = FastRandomGenerator()
    
    private var seed: UInt32

    init() {
        seed = UInt32.random(in: 1...500000000)
    }

    mutating func nextRandom() -> UInt32 {
        seed ^= (seed << 13)
        seed ^= (seed >> 17)
        seed ^= (seed << 5)
        return seed
    }
    
    static func generate(range: UInt32) -> UInt32 {
        let rand = FastRandomGenerator.shared.nextRandom()
        return UInt32((UInt64(rand) * UInt64(range)) >> 32)
    }
}

extension String {
    func paddedZeros(length: Int) -> String {
        if count == length {
            return self
        }
        return ([Character](repeating: "0", count: length - count) + Array(self)).map({ String($0) }).joined()
    }
}

extension Int {
    func paddedString() -> String {
        var str = Array(String(self == 0 ? 0 : (1 << self))).map({ String($0) })
        str = [String](repeating: " ", count: 4 - str.count) + str
        return str.joined()
    }
}

typealias Board = UInt64
typealias Row = UInt16

let ROW1: UInt64 = 0xFFFF000000000000
let ROW2: UInt64 = 0xFFFF00000000
let ROW3: UInt64 = 0xFFFF0000
let ROW4: UInt64 = 0xFFFF

struct Utils2048 {
    var rowLeftTable: [Row]
    var rowRightTable: [Row]
    var scoreTable: [Double]
    var winTable: [Bool]

    static let shared = Utils2048()

    private init() {
        rowLeftTable = [Row](repeating: 0, count: 65536)
        rowRightTable = [Row](repeating: 0, count: 65536)
        scoreTable = [Double](repeating: 0, count: 65536)
        winTable = [Bool](repeating: false, count: 65536)
        initializeTables()
    }

    mutating func initializeTables() {
        print("Initializing tables...")
        for row in 0..<UInt16.max {
            var line = [(row & 0xF000) >> 12,
                        (row & 0xF00) >> 8,
                        (row & 0xF0) >> 4,
                        (row & 0xF) >> 0]
            winTable[Int(row)] = line.contains(11)

            let scoreValues = line.map { v -> UInt in
                let v = UInt(v)
                return v == 0 ? 0 : ((1 << v) * (v - 1))
            }
            scoreTable[Int(row)] = Double(scoreValues.reduce(0, +))

            var newLine: [UInt16] = [0, 0, 0, 0]
            var j = 0
            var previous: UInt16? = nil
            for i in 0..<4 {
                if line[i] != 0 {
                    if previous == nil {
                        previous = line[i]
                    } else {
                        if previous == line[i] {
                            newLine[j] = line[i] + 1
                            j += 1
                            previous = nil
                        } else {
                            newLine[j] = previous!
                            j += 1
                            previous = line[i]
                        }
                    }
                }
            }
            if previous != nil {
                newLine[j] = previous!
            }
            line = newLine

            var row = row
            var result = (line[0] << 12) | (line[1] << 8) | (line[2] << 4) | (line[3] << 0)

            rowLeftTable[Int(row)] = result
            Utils2048.reverse(row: &result)
            Utils2048.reverse(row: &row)
            rowRightTable[Int(row)] = result
        }
    }

    static func reverse(row: inout Row) {
        row = (row >> 12) | ((row >> 4) & 0x00F0) | ((row << 4) & 0x0F00) | (row << 12)
    }

    static func line(row: Row) -> [Row] {
        let line = [(row & 0xF000) >> 12,
                    (row & 0xF00) >> 8,
                    (row & 0xF0) >> 4,
                    (row & 0xF) >> 0]
        return line
    }

    static func row(line: [UInt16]) -> Row {
        return (line[0] << 12) | (line[1] << 8) | (line[2] << 4) | (line[3] << 0)
    }

    static func rows(board: Board) -> [Row] {
        let rows = [
            Row(board >> 48),
            Row((board >> 32) & 0xFFFF),
            Row((board >> 16) & 0xFFFF),
            Row((board >>  0) & 0xFFFF)
        ]
        return rows
    }

    static func lines(board: Board) -> [[Row]] {
        let lines = rows(board: board).map({ line(row: $0) })
        return lines
    }

    static func transpose(board x: inout Board) {
        let a1 = x & 0xF0F00F0FF0F00F0F
        let a2 = x & 0x0000F0F00000F0F0
        let a3 = x & 0x0F0F00000F0F0000
        let a = a1 | (a2 << 12) | (a3 >> 12)
        let b1 = a & 0xFF00FF0000FF00FF
        let b2 = a & 0x00FF00FF00000000
        let b3 = a & 0x00000000FF00FF00
        x = b1 | (b2 >> 24) | (b3 << 24)
    }

    static func emptyTiles(board x: Board) -> Int {
        var x = x
        x |= (x >> 2) & 0x3333333333333333
        x |= (x >> 1)
        x = ~x & 0x1111111111111111
        x += x >> 32
        x += x >> 16
        x += x >>  8
        x += x >>  4
        return Int(x & 0xf)
    }
    
    static func insertRandomTile(board: Board, _ tiles: Int) -> Board {
        let chosenTile = FastRandomGenerator.generate(range: UInt32(tiles))
        var t: UInt32 = 0
        var offset: UInt32 = 0
        var temp = board
        while temp != 0 {
            if temp & 0xF == 0 {
                if t == chosenTile {
                    break
                }
                t += 1
            }
            temp >>= 4
            offset += 4
        }
        if t != chosenTile {
            offset += (chosenTile - t) << 2
        }
        var insert = 1
        if FastRandomGenerator.shared.nextRandom() < 429496729 {
            insert = 2
        }
        return board | Board(insert << offset)
    }
    
    static func canMoveLeft(board: Board) -> Bool {
        var row = Int((board >> 0) & 0xFFFF)
        if row != Utils2048.shared.rowLeftTable[row] {
            return true
        }
        row = Int((board >> 16) & 0xFFFF)
        if row != Utils2048.shared.rowLeftTable[row] {
            return true
        }
        row = Int((board >> 32) & 0xFFFF)
        if row != Utils2048.shared.rowLeftTable[row] {
            return true
        }
        row = Int((board >> 48) & 0xFFFF)
        if row != Utils2048.shared.rowLeftTable[row] {
            return true
        }
        return false
    }
    
    static func canMoveRight(board: Board) -> Bool {
        var row = Int((board >> 0) & 0xFFFF)
        if row != Utils2048.shared.rowRightTable[row] {
            return true
        }
        row = Int((board >> 16) & 0xFFFF)
        if row != Utils2048.shared.rowRightTable[row] {
            return true
        }
        row = Int((board >> 32) & 0xFFFF)
        if row != Utils2048.shared.rowRightTable[row] {
            return true
        }
        row = Int((board >> 48) & 0xFFFF)
        if row != Utils2048.shared.rowRightTable[row] {
            return true
        }
        return false
    }
    
    static func canMoveUp(board: Board) -> Bool {
        var board = board
        transpose(board: &board)
        return canMoveLeft(board: board)
    }
    
    static func canMoveDown(board: Board) -> Bool {
        var board = board
        transpose(board: &board)
        return canMoveRight(board: board)
    }
    
    static func gameWon(board: Board) -> Bool {
        Utils2048.shared.winTable[Int((board >> 0) & 0xFFFF)] ||
        Utils2048.shared.winTable[Int((board >> 16) & 0xFFFF)] ||
        Utils2048.shared.winTable[Int((board >> 32) & 0xFFFF)] ||
        Utils2048.shared.winTable[Int((board >> 48) & 0xFFFF)]
    }

    static func printBoard(board: Board) {
        let board = lines(board: board)
        for r in 0..<4 {
            for c in 0..<4 {
                let powerVal = board[r][c]
                print(powerVal == 0 ? 0 : (1 << powerVal), terminator: "|")
            }
            print("")
        }
        print("")
    }
}

struct Game2048 {
    var board: Board

    var score: Double {
        Utils2048.shared.scoreTable[Int((board >> 0) & 0xFFFF)] +
        Utils2048.shared.scoreTable[Int((board >> 16) & 0xFFFF)] +
        Utils2048.shared.scoreTable[Int((board >> 32) & 0xFFFF)] +
        Utils2048.shared.scoreTable[Int((board >> 48) & 0xFFFF)]
    }

    var gameWon: Bool {
        return Utils2048.gameWon(board: board)
    }

    var gameOver: Bool {
        if gameWon {
            return true
        }
        if Utils2048.emptyTiles(board: board) == 0 {
            if Utils2048.canMoveLeft(board: board) {
                return false
            } else if Utils2048.canMoveRight(board: board) {
                return false
            } else if Utils2048.canMoveUp(board: board) {
                return false
            } else if Utils2048.canMoveDown(board: board) {
                return false
            }
            return true
        }
        return false
    }

    enum Direction {
        case up, down, left, right

        var string: String {
            switch self {
            case .up:
                return "up"
            case .down:
                return "down"
            case .left:
                return "left"
            case .right:
                return "right"
            }
        }
    }

    private mutating func moveHorizontal(lookup: [Row]) {
        let row1 = Board(lookup[Int((board & ROW1) >> 48)])
        let row2 = Board(lookup[Int((board & ROW2) >> 32)])
        let row3 = Board(lookup[Int((board & ROW3) >> 16)])
        let row4 = Board(lookup[Int((board & ROW4) >>  0)])
        board = (row1 << 48) | (row2 << 32) | (row3 << 16) | (row4 << 0)
    }

    private mutating func moveVertical(lookup: [Row]) {
        Utils2048.transpose(board: &board)
        moveHorizontal(lookup: lookup)
        Utils2048.transpose(board: &board)
    }

    mutating func moveLeft() {
        moveHorizontal(lookup: Utils2048.shared.rowLeftTable)
    }

    mutating func moveRight() {
        moveHorizontal(lookup: Utils2048.shared.rowRightTable)
    }

    mutating func moveUp() {
        moveVertical(lookup: Utils2048.shared.rowLeftTable)
    }

    mutating func moveDown() {
        moveVertical(lookup: Utils2048.shared.rowRightTable)
    }

    mutating func move(direction: Direction) {
        switch direction {
        case .up:
            moveUp()
        case .down:
            moveDown()
        case .left:
            moveLeft()
        case .right:
            moveRight()
        }
    }

    mutating func play(direction: Direction) {
        move(direction: direction)
        board = Utils2048.insertRandomTile(board: board, Utils2048.emptyTiles(board: board))
    }

    func print() {
        let board = Utils2048.lines(board: self.board).map({ $0.map({ Int($0) }) })
        var temp = ""
        temp += [String](repeating: "-", count: "|    |    |    |    |".count).joined()
        temp += "\n"
        temp += "|\(board[0][0].paddedString())|\(board[0][1].paddedString())|\(board[0][2].paddedString())|\(board[0][3].paddedString())|"
        temp += "\n"
        temp += [String](repeating: "-", count: "|    |    |    |    |".count).joined()
        temp += "\n"
        temp += "|\(board[1][0].paddedString())|\(board[1][1].paddedString())|\(board[1][2].paddedString())|\(board[1][3].paddedString())|"
        temp += "\n"
        temp += [String](repeating: "-", count: "|    |    |    |    |".count).joined()
        temp += "\n"
        temp += "|\(board[2][0].paddedString())|\(board[2][1].paddedString())|\(board[2][2].paddedString())|\(board[2][3].paddedString())|"
        temp += "\n"
        temp += [String](repeating: "-", count: "|    |    |    |    |".count).joined()
        temp += "\n"
        temp += "|\(board[3][0].paddedString())|\(board[3][1].paddedString())|\(board[3][2].paddedString())|\(board[3][3].paddedString())|"
        temp += "\n"
        temp += [String](repeating: "-", count: "|    |    |    |    |".count).joined()
        temp += "\n"
        Swift.print(temp)
    }

    func copy() -> Game2048 {
        return Game2048(board: board)
    }
}

extension Game2048 {
    init() {
        board = Utils2048.insertRandomTile(board: Utils2048.insertRandomTile(board: 0, 16), 15)
    }

    init(board: [[Int]]) {
        let str = board.flatMap({ $0 }).map({ $0 == 0 ? "0000" : String(Int(log(Float($0)) / log(2)), radix: 2).paddedZeros(length: 4) }).joined()
        self.board = Board(str, radix: 2)!
    }
}

let directionsMap: [Game2048.Direction] = [
    .left,
    .right,
    .up,
    .down
]

func monteCarloSearch(state: Game2048) -> Game2048.Direction {
    func performRollout(state: Game2048, direction: Game2048.Direction) -> Double {
        var newState = state.copy() // Copy the initial state.
        newState.play(direction: direction) // Move in the direction to be scored.
        while !newState.gameOver { // Repeat until game is over.
            let left = Utils2048.canMoveLeft(board: newState.board) ? 1 : 0
            let right = Utils2048.canMoveRight(board: newState.board) ? 1 : 0
            let up = Utils2048.canMoveUp(board: newState.board) ? 1 : 0
            let down = Utils2048.canMoveDown(board: newState.board) ? 1 : 0
            let valids = left + right + up + down
            let chosen = Int(FastRandomGenerator.generate(range: UInt32(valids)))
            var t = 0
            if left != 0 {
                if t == chosen {
                    newState.play(direction: .left)
                    continue
                } else {
                    t += 1
                }
            }
            if right != 0 {
                if t == chosen {
                    newState.play(direction: .right)
                    continue
                } else {
                    t += 1
                }
            }
            if up != 0 {
                if t == chosen {
                    newState.play(direction: .up)
                    continue
                } else {
                    t += 1
                }
            }
            if down != 0 {
                if t == chosen {
                    newState.play(direction: .down)
                    continue
                }
            }
        }
        return newState.score
    }

    let left = Utils2048.canMoveLeft(board: state.board) ? 1 : 0
    let right = Utils2048.canMoveRight(board: state.board) ? 1 : 0
    let up = Utils2048.canMoveUp(board: state.board) ? 1 : 0
    let down = Utils2048.canMoveDown(board: state.board) ? 1 : 0
    let total = left + right + up + down

    if total == 1 {
        if left != 0 {
            return .left
        } else if right != 0 {
            return .right
        } else if up != 0 {
            return .up
        } else if down != 0 {
            return .down
        }
    }

    let maxTime: Double = 1 / 9600
    var scoreUp = 0.0
    var scoreDown = 0.0
    var scoreLeft = 0.0
    var scoreRight = 0.0

    DispatchQueue.concurrentPerform(iterations: 4) { dirIndex in
        let direction: Game2048.Direction
        switch dirIndex {
        case 0:
            if left == 0 {
                return
            }
            direction = .left
        case 1:
            if right == 0 {
                return
            }
            direction = .right
        case 2:
            if up == 0 {
                return
            }
            direction = .up
        case 3:
            if down == 0 {
                return
            }
            direction = .down
        default:
            fatalError()
        }
        var totalScore: Double = 0
        var rollouts = 0
        let timeBarrier = DispatchTime.now() + maxTime
        repeat {
            rollouts += 1
            let score = performRollout(state: state, direction: direction)
            totalScore += score
        } while DispatchTime.now() <= timeBarrier
        let avgScore = totalScore / Double(rollouts)
        switch direction {
        case .up:
            scoreUp = avgScore
        case .down:
            scoreDown = avgScore
        case .left:
            scoreLeft = avgScore
        case .right:
            scoreRight = avgScore
        }
    }
    
    var max = scoreUp
    var best = Game2048.Direction.up
    if scoreDown > max {
        max = scoreDown
        best = .down
    }
    if scoreLeft > max {
        max = scoreLeft
        best = .left
    }
    if scoreRight > max {
        max = scoreRight
        best = .right
    }
    return best
}