5e ability score generation

abilityGen.swift

// Requires Swift 4.2. Compile with `swiftc -O abilityGen.swift`. Run with ./abilityGen <type> > <output.csv>
import Foundation

let costs = [
    8: 0,
    9: 1,
    10: 2,
    11: 3,
    12: 4,
    13: 5,
    14: 7,
    15: 9
]

func randomAbilityScores() -> [Int] {
    let randomScore = { (8..<15).randomElement()! }
    let result = sequence(first: randomScore(), next: { _ in randomScore() })
        .prefix(6)
    return Array(result)
}

func pointBuyCost(for abilityScores: [Int]) -> Int {
    return abilityScores.reduce(0) { $0 + costs[$1]! }
}

enum Direction {
    case back, forth
}

extension RangeReplaceableCollection where Element: Comparable & Numeric {
    
    func neighbors(boundedBy bounds: ClosedRange<Element>) -> [Self] {
        var result: [Self] = []
        
        for (index, value) in zip(self.indices, self) {
            for (neighborIndex, neighborValue) in zip(self.indices, self) where neighborIndex != index && neighborValue < bounds.upperBound  {
                var neighbor = Self()
                neighbor.reserveCapacity(self.count)
                
                let next: Index
                if neighborIndex < index {
                    neighbor.append(contentsOf: self[self.startIndex..<neighborIndex])
                    neighbor.append(neighborValue + 1)
                    neighbor.append(contentsOf: self[self.index(after: neighborIndex)..<index])
                    neighbor.append(value)
                    next = self.index(after: index)
                } else {
                    neighbor.append(contentsOf: self[self.startIndex..<index])
                    neighbor.append(value)
                    neighbor.append(contentsOf: self[self.index(after: index)..<neighborIndex])
                    neighbor.append(neighborValue + 1)
                    next = self.index(after: neighborIndex)
                }
                neighbor.append(contentsOf: self[next..<self.endIndex])
                
                result.append(neighbor)
            }
        }
                
        for (index, value) in zip(self.indices, self) {
            for (neighborIndex, neighborValue) in zip(self.indices, self) where neighborIndex != index && neighborValue > bounds.lowerBound  {
                var neighbor = Self()
                neighbor.reserveCapacity(self.count)
                
                let next: Index
                if neighborIndex < index {
                    neighbor.append(contentsOf: self[self.startIndex..<neighborIndex])
                    neighbor.append(neighborValue - 1)
                    neighbor.append(contentsOf: self[self.index(after: neighborIndex)..<index])
                    neighbor.append(value)
                    next = self.index(after: index)
                } else {
                    neighbor.append(contentsOf: self[self.startIndex..<index])
                    neighbor.append(value)
                    neighbor.append(contentsOf: self[self.index(after: index)..<neighborIndex])
                    neighbor.append(neighborValue - 1)
                    next = self.index(after: neighborIndex)
                }
                neighbor.append(contentsOf: self[next..<self.endIndex])
                
                result.append(neighbor)
            }
        }
        
        return result
    }
    
}

func fitToPointBuy(startingFrom abilityScores: [Int], target: Int) -> [Int] {
    precondition(abilityScores.count == 6)
    var score = pointBuyCost(for: abilityScores)
    guard score != target else {
        return abilityScores
    }
    
    var result = abilityScores
    repeat {
        let candidates = result.neighbors(boundedBy: 8...15)
            .filter {
                let cost = pointBuyCost(for: $0)
                if score < target {
                    return cost >= score
                } else {
                    return cost <= score
                }
            }
        result = candidates.randomElement()!
        score = pointBuyCost(for: result)
    } while score != target

    let indices = result.indices
    
    let big = indices.randomElement()!
    let small = indices.filter { $0 != big }.randomElement()!
    result[big] += 2
    result[small] += 1
    
    return result
}

func randomAbilityScoresPF25e() -> [Int] {
    var result = [Int](repeating: 10, count: 6)
    let indices = result.indices
    
    // Apply background boost
    result[indices.randomElement()!] += 2
    
    // Apply racial boosts
    let big = indices.randomElement()!
    let small = indices.filter { $0 != big }.randomElement()!
    result[big] += 2
    result[small] += 1

    // Apply class boost
    result[indices.randomElement()!] += 2

    // Apply three free boosts
    let freeBoosts = indices.shuffled().dropFirst(3)
    freeBoosts.forEach { result[$0] += 2 }
    
    return result
}

func randomAbilityScoresPF2() -> [Int] {
    var result = [Int](repeating: 10, count: 6)
    let indices = result.indices
    
    // Apply background boosts
    let firstBG = indices.randomElement()!
    let secondBG = indices.filter { $0 != firstBG }.randomElement()!
    result[firstBG] += 2
    result[secondBG] += 2
    
    // Apply racial boosts
    let first = indices.randomElement()!
    let second = indices.filter { $0 != first }.randomElement()!
    let third = indices.filter { $0 != first && $0 != second }.randomElement()!
    result[first] += 2
    result[second] += 2
    result[third] += 2
    result[indices.randomElement()!] -= 2

    // Apply class boost
    result[indices.randomElement()!] += 2

    // Apply four free boosts
    let freeBoosts = indices.shuffled().dropFirst(2)
    freeBoosts.forEach { result[$0] += 2 }
    
    return result
}

func randomAbilityScoresPF2Custom() -> [Int] {
    var result = [Int](repeating: 10, count: 6)
    let indices = result.indices
    
    // Apply background boost
    result[indices.randomElement()!] += 2
    
    // Apply racial boosts
    let big = indices.randomElement()!
    let free = indices.filter { $0 != big }.randomElement()!
    result[big] += 2
    result[free] += 2
    result[indices.randomElement()!] -= 2

    // Apply class boost
    result[indices.randomElement()!] += 2

    // Apply three free boosts
    let freeBoosts = indices.shuffled().dropFirst(3)
    freeBoosts.forEach { result[$0] += 2 }
    
    return result
}

func printUsageAndQuit() -> Never {
    print("./abilityGen 5e|pf2|pf2-5e|custom")
    exit(-1)
}

let genFn: () -> [Int]
guard CommandLine.arguments.count > 1 else {
    printUsageAndQuit()
}
switch CommandLine.arguments[1] {
case "5e":
    genFn = { fitToPointBuy(startingFrom: randomAbilityScores(), target: 27) }
case "pf2-5e":
    genFn = randomAbilityScoresPF25e
case "pf2":
    genFn = randomAbilityScoresPF2
case "custom":
    genFn = randomAbilityScoresPF2Custom
default:
    printUsageAndQuit()
}

print("Strength,Dexterity,Constitution,Intelligence,Wisdom,Charisma")
for _ in 1 ... 1_000_000 {
    let abilityScores = genFn().sorted()
    print(abilityScores[0], abilityScores[1], abilityScores[2], abilityScores[3], abilityScores[4], abilityScores[5], separator: ",")
}