import math
import matplotlib.pyplot as plt

from functools import cache



def collatz_step(n: int) -> tuple[int, int]:
    n = 3 * n + 1
    v2 = (n & -n).bit_length() - 1
    n >>= v2
    return n, v2


def S_k(seed: int, k: int) -> int:
    
    @cache
    def get_2adic_valuation_list(seed: int) -> list[int]:
        result = []
        n = seed
        while True:
            n, v2 = collatz_step(n)
            result.append(v2)
            if n == 1:
                break
        
        # print("-" * 20)
        # print(f"The longest E-streak length for N = {seed}")
        # print(f"L = {max(result)}")
        # print("-" * 20)
        
        return result
    
    valuations = get_2adic_valuation_list(seed)
    k = min(k, len(valuations))
    partial_sum = sum(valuations[:k-1])
    return partial_sum


def Δ_k(n, k) -> float:
    return S_k(n, k) - k*math.log2(3)

S_k(27, math.inf)

seed = 27

k_values = []
Δ_k_values = []

for k in range(100):
    k_values.append(k)
    Δ_k_values.append(Δ_k(seed, k))

plt.plot(k_values, Δ_k_values, "ko-", markersize=4)
plt.grid(True, linestyle=':', alpha=0.9)
plt.tight_layout()
plt.show()