Screeps_Body_Composer.ts

// Code by xTwisteDx aka MKaulfers (https://www.github.com/mkaulfers)

// Example usage
const someComp = new BodyComp(300)  // assuming 300 energy available
    .work(1, 3, 0.2)                // at least 1 work part, up to 3, with a ratio of 20%
    .move(1, undefined, 0.5)        // at least 1 move part, up to 2, with a ratio of 50%
    .carry(1, 2)                    // at least 1 carry part, up to 2

const someBody: BodyPartConstant[] = someComp.body()
console.log(someBody)

// -----------------------------------------------------------------------------------------------------------

/**
 * Utilitiy class for creating a body composition for creeps.
 * Chainable methods are provided for adding part configurations.
 * Each part configuration specifies the type of body part, the minimum
 * and maximum number of parts, and an optional ratio to control the
 * balance between parts. The ratio is the proportion of the part type
 * to the total number of parts in the body. The class calculates the
 * body composition based on the energy limit provided in the constructor.
 */
export class BodyComp {
    private energyExpenseLimit: number
    private partsConfig: BodyPartConfig[] = []
    private bodyParts: BodyPartConstant[] = []

    /**
     * @returns - Returns the body composition as an array of body parts
     * sorted by the default sorting function. An optional sorting function
     * can be provided to sort the body parts in a custom order.
     * @param sortFn - Optional sorting function for custom sorting
     * @returns - Returns the sorted body parts array
     * @example
     *
     * const body = someComp.body((a, b) => {
     *    return a === WORK ? -1 : b === WORK ? 1 : 0
     * })
     * console.log(body)  // [WORK, MOVE, MOVE, CARRY]
     */
    get body() {
        this.calculateBody()
        return (sortFn: (a: BodyPartConstant, b: BodyPartConstant) => number = this.defaultSorting) => {
            return [...this.bodyParts].sort(sortFn)
        }
    }

    get cost() {
        return this.body().reduce((cost, part) => cost + BODYPART_COST[part], 0)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    attack(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(ATTACK, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    carry(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(CARRY, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    claim(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(CLAIM, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    heal(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(HEAL, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    move(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(MOVE, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    rangedAttack(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(RANGED_ATTACK, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    tough(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(TOUGH, min, max, ratio)
    }

    /**
     * @param min - Minimum number of parts
     * @param max - Maximum number of parts
     * @param ratio - Optional ratio to control the balance between parts
     * @returns - Returns the BodyComp instance for method chaining
     */
    work(min: number = 1, max: number = Infinity, ratio?: number) {
        return this.addPartConfig(WORK, min, max, ratio)
    }

    private addPartConfig(type: BodyPartConstant, min: number, max: number, ratio?: number) {
        this.partsConfig.push({ type, min, max, ratio })
        return this
    }

    // Default sorting function
    private defaultSorting(a: BodyPartConstant, b: BodyPartConstant): number {
        const order = [TOUGH, MOVE, WORK, CARRY, ATTACK, RANGED_ATTACK, HEAL, CLAIM]
        return order.indexOf(a) - order.indexOf(b)
    }

    private calculateBody() {
        // Calculate the energy cost of each body part type
        const partCosts: { [part: string]: number } = {
            [ATTACK]: BODYPART_COST[ATTACK],
            [CARRY]: BODYPART_COST[CARRY],
            [CLAIM]: BODYPART_COST[CLAIM],
            [HEAL]: BODYPART_COST[HEAL],
            [MOVE]: BODYPART_COST[MOVE],
            [RANGED_ATTACK]: BODYPART_COST[RANGED_ATTACK],
            [TOUGH]: BODYPART_COST[TOUGH],
            [WORK]: BODYPART_COST[WORK],
        }

        let remainingEnergy = this.energyExpenseLimit
        let totalParts = 0

        // First, satisfy the minimum requirements
        for (const config of this.partsConfig) {
            const cost = partCosts[config.type]
            const maxParts = Math.min(config.max, Math.floor(remainingEnergy / cost))

            for (let i = 0; i < config.min && remainingEnergy >= cost; i++) {
                this.bodyParts.push(config.type)
                remainingEnergy -= cost
                totalParts++
            }
        }

        // Apply ratios across all part types
        while (remainingEnergy > 0 && totalParts < 50) {  // Screeps has a max body part limit of 50
            const partOptions = this.partsConfig.filter(config => {
                const cost = partCosts[config.type]
                const partCount = this.bodyParts.filter(part => part === config.type).length
                const withinMax = partCount < config.max
                const ratioMet = !config.ratio || (partCount + 1) / (totalParts + 1) <= config.ratio  // Check if adding maintains ratio

                return remainingEnergy >= cost && withinMax && ratioMet
            })

            if (partOptions.length === 0) break

            // Add parts based on the first available option that meets criteria
            for (const config of partOptions) {
                const cost = partCosts[config.type]
                if (remainingEnergy >= cost) {
                    this.bodyParts.push(config.type)
                    remainingEnergy -= cost
                    totalParts++
                    break  // Add one part per loop iteration
                }
            }
        }
    }

    constructor(energyExpenseLimit: number) {
        this.energyExpenseLimit = energyExpenseLimit
    }
}