from math import exp, log

scores = [1000 + log(i) for i in range(1, 5)]
values = [10.0, 20.0, 30.0, 40.0]
maximum = max(scores)
weights = [exp(score - maximum) for score in scores]
dense = sum(w * v for w, v in zip(weights, values)) / sum(weights)

def streaming(tile_size):
    running_max = float("-inf")
    denominator = 0.0
    numerator = 0.0
    for start in range(0, len(scores), tile_size):
        tile_scores = scores[start:start + tile_size]
        tile_values = values[start:start + tile_size]
        new_max = max(running_max, max(tile_scores))
        correction = exp(running_max - new_max)
        local_weights = [exp(s - new_max) for s in tile_scores]
        denominator = denominator * correction + sum(local_weights)
        numerator = numerator * correction + sum(
            w * v for w, v in zip(local_weights, tile_values))
        running_max = new_max
    return numerator / denominator

for tile_size in (1, 2, 4):
    result = streaming(tile_size)
    assert abs(result - dense) < 1e-10
    print(f"tile={tile_size}: output={result:.6f}")
assert abs(dense - 30.0) < 1e-10
length, block = 4096, 64
print(f"score_entries: full_matrix={length * length}, one_row_tile={block}")
