quality.py

"""
https://wiki.factorio.com/Quality#Quality_modules
When a machine produces an item, quality modules give it Q percent chance of improving the item
"""
from enum import Enum
import numpy as np
from collections import Counter

np.set_printoptions(formatter={'float': lambda x: f"{x:10.4g}"})

class ModuleTier(Enum):
    """
    Tiers of modules
    """
    TIER1 = 0
    TIER2 = 1
    TIER3 = 2

class QualityIndex(Enum):
    """
    Quality tiers
    """
    COMMON = 0
    UNCOMMON = 1
    RARE = 2
    EPIC = 3
    LEGENDARY = 4

    def __lt__(self, other):
        return self.value < other.value

_QUALITY = {
    # In order of QualityIndex
    ModuleTier.TIER1: [0.01, 0.013, 0.016, 0.019, 0.025],
    ModuleTier.TIER2: [0.02, 0.026, 0.032, 0.038, 0.05],
    ModuleTier.TIER3: [0.025, 0.032, 0.04, 0.047, 0.062]
}

_PRODUCTIVITY = {
    # In order of QualityIndex
    ModuleTier.TIER1: [0.04, 0.05, 0.06, 0.07, 0.1],
    ModuleTier.TIER2: [0.06, 0.07, 0.09, 0.11, 0.15],
    ModuleTier.TIER3: [0.10, 0.13, 0.16, 0.19, 0.25]
}

class ModuleKind(Enum):
    """
    Different kinds of modules
    """
    QUALITY = _QUALITY
    PRODUCTIVITY = _PRODUCTIVITY

input_quality = QualityIndex.COMMON

def select_modules(kind: ModuleKind, tier: ModuleTier, quality: QualityIndex, n: int):
    """
    Helper for getting a list of modules of a certain kind/tier/quality
    """
    mapping = kind.value
    module = mapping[tier][quality.value]
    return [module] * n

def calculate_percents(base_prod, quality_modules, prod_modules):
    """
    Calculate a probability matrix given a base productivity bonus,
    a set of quality modules, and a set of productivity modules
    """
    Q = sum(quality_modules)
    P = sum(prod_modules) + (base_prod * 0.01)

    result = np.zeros((5, 5))

    for i in QualityIndex:
        # i is the starting quality
        for j in range(i.value, QualityIndex.LEGENDARY.value + 1):
            # j is output quality
            if i.value == j:
                out = (1 + P) * (1 - Q)
            else:
                i_shifted = j - i.value
                i0 = i_shifted
                if j == QualityIndex.LEGENDARY.value:
                    q = (Q * (1 / (10**(i0 - 1))))
                else:
                    q = (Q * (9 / (10**(i0))))
                out = (1 + P) * q
            result[i.value, j] = out

    result[QualityIndex.LEGENDARY.value, QualityIndex.LEGENDARY.value] = (1 + P)
    return result

def recycle(quality_modules):
    """
    Calculates probability of getting a certain quality of item, note that
    legendary items aren't recycled
    """
    results = calculate_percents(-75, quality_modules, [])
    results[-1, :] = 0
    return results


def build_transition_matrix(q_mods_craft, q_mods_recycle, prod_modules):
    """
    Build a transition matrix, i.e. multiplying the result of this function
    on itself implies 2 crafting steps and 2 recycling steps
    """
    # Crafting submatrix (first 5 rows, last 5 columns)
    crafting_submatrix = calculate_percents(50, q_mods_craft, prod_modules)
    
    # Recycling submatrix (last 5 rows, first 5 columns)
    recycling_submatrix = recycle(q_mods_recycle)
    
    # Initialize the full 10x10 transition matrix with zeros
    transition_matrix = np.zeros((10, 10))
    
    # First 5 rows, last 5 columns: Crafting transitions
    transition_matrix[:5, 5:] = crafting_submatrix
    
    # Last 5 rows, first 5 columns: Recycling transitions
    transition_matrix[5:, :5] = recycling_submatrix

    # Legendary products are the goal
    transition_matrix[-1, -1] = 1
    
    return transition_matrix

def test_recycle():
    print("test recycle")
    qrmods = select_modules(ModuleKind.QUALITY, ModuleTier.TIER2, QualityIndex.LEGENDARY, 5)
    rdist = recycle(qrmods)[1]
    assert np.allclose(
        rdist,
        [0, 0.1875, 0.05625, 0.005625, 0.000625]
    )
    print(rdist)

def test_quality():
    print("test quality")
    # test just q
    qmods = select_modules(ModuleKind.QUALITY, ModuleTier.TIER3, QualityIndex.LEGENDARY, 4)
    distribution = calculate_percents(0, qmods, [])[0]
    assert np.allclose(
        distribution,
        [0.752, 0.2232, 0.02232, 0.002232, 0.000248]
    )
    print(distribution)

def test_matrix():
    qrmods = select_modules(ModuleKind.QUALITY, ModuleTier.TIER2, QualityIndex.LEGENDARY, 5)
    M = build_transition_matrix(qrmods, qrmods, [])
    M_real = np.array([
        # crafting matrix
        [0,     0,      0,          0,          0,          1.125,  0.3375, 0.03375,0.003375,   0.000375],
        [0,     0,      0,          0,          0,          0,      1.125,  0.3375, 0.03375,    0.00375],
        [0,     0,      0,          0,          0,          0,      0,      1.125,  0.3375,     0.0375],
        [0,     0,      0,          0,          0,          0,      0,      0,      1.125,      0.375],
        [0,     0,      0,          0,          0,          0,      0,      0,      0,          1.5],
        # recycling matrix
        [0.1875,0.05625,0.005625,   0.0005625,  0.0000625,  0,      0,      0,      0,          0],
        [0,     0.1875, 0.05625,    0.005625,   0.000625,   0,      0,      0,      0,          0],
        [0,     0,      0.1875,     0.05625,    0.00625,    0,      0,      0,      0,          0],
        [0,     0,      0,          0.1875,     0.0625,     0,      0,      0,      0,          0],
        [0,     0,      0,          0,          0,          0,      0,      0,      0,          1]
    ])
    assert(np.allclose(M, M_real))

test_recycle()
test_quality()
test_matrix()


def required_inputs_for_legendary(transition_matrix, initial_quality):
    """
    First, calculate a converged state for the transition matrix,
    i.e. given one set of inputs, there should not be a significant difference between
    applying x recycling/crafting steps, and x+1 reyclcing/crafting steps

    Then we can calculate, given the expected output, how many input sets are required
    """
    state_vector = np.zeros(10)
    state_vector[initial_quality.value] = 1  # Start at the given initial quality
    s = state_vector

    steps_to_converge = 0
    old_s = s.copy()
    old_s[9] = 1
    while abs(old_s[9] - s[9]) > 0.0001:
        transition_matrix = transition_matrix @ transition_matrix
        old_s[9] = s[9]
        s = state_vector @ transition_matrix
        steps_to_converge += 1

    return 1 / s[9]

qrmods = select_modules(ModuleKind.QUALITY, ModuleTier.TIER3, QualityIndex.LEGENDARY, 4)
qmods = select_modules(ModuleKind.QUALITY, ModuleTier.TIER3, QualityIndex.LEGENDARY, 1)
pmods = select_modules(ModuleKind.PRODUCTIVITY, ModuleTier.TIER3, QualityIndex.LEGENDARY, 3)
T = build_transition_matrix(qmods, qrmods, pmods)
input_sets = required_inputs_for_legendary(T, QualityIndex.COMMON)
print(input_sets)