from math import factorial, log10
from komorebi_math.sequences import fubini_sequence
from komorebi_browser import show_plot, export_file

last, row = 18, 5
if not (type(last) is int and 0 <= last <= 18 and type(row) is int and 0 <= row <= last):
    raise ValueError("Use last=0..18 and row=0..last.")
S = [[0]*(last+1) for _ in range(last+1)]
S[0][0] = 1
for n in range(1, last+1):
    for k in range(1, n+1):
        S[n][k] = k*S[n-1][k] + S[n-1][k-1]
weighted = [[factorial(k)*S[n][k] for k in range(n+1)] for n in range(last+1)]
if tuple(map(sum, weighted)) != fubini_sequence().first(last+1):
    raise ArithmeticError("The Stirling sums disagree with the independent library recurrence.")
ks = list(range(row+1))
show_plot([
    {"type": "bar", "x": ks, "y": S[row][:row+1],
     "name": "Unordered S(n,k)", "marker": {"color": "#0067a5"},
     "hovertemplate": "k = %{x}<br>Count = %{y}<extra>Unordered blocks</extra>"},
    {"type": "bar", "x": ks, "y": weighted[row],
     "name": "Ordered k! S(n,k)", "marker": {"color": "#b54b00", "pattern": {"shape": "/"}},
     "hovertemplate": "k = %{x}<br>Count = %{y}<extra>Ordered blocks</extra>"}],
    {"title": {"text": f"What factorial weighting changes (n={row})"}, "barmode": "group",
     "xaxis": {"title": {"text": "Number of blocks k"}, "dtick": 1},
     "yaxis": {"title": {"text": "Count"}, "rangemode": "tozero"}},
    description=f"Blue bars count unordered partitions. Striped orange bars count ordered partitions. These are ordinary counts, not logarithms. For this row the totals are {sum(S[row])} and {sum(weighted[row])}. Exact values are in the CSV below.")
show_plot([{"type": "heatmap", "x": list(range(last+1)), "y": list(range(last+1)),
    "z": [[log10(1+weighted[n][k]) if k <= n else None for n in range(last+1)] for k in range(last+1)],
    "colorscale": "Cividis", "zmin": 0, "hoverongaps": False,
    "colorbar": {"orientation": "h", "y": 1.12, "yanchor": "bottom", "len": 0.95, "thickness": 16, "nticks": 4, "title": {"text": "log10 (1 + count)", "side": "top"}},
    "hovertemplate": "n = %{x}<br>k = %{y}<br>log10 (1 + count) = %{z:.3f}<extra></extra>"}],
    {"title": {"text": "Where the factorial-weighted contributions sit"},
     "xaxis": {"title": {"text": "Set size n"}, "dtick": 3},
     "yaxis": {"title": {"text": "Block count k"}, "dtick": 3}, "margin": {"t": 120}},
    description="Each square is one integer pair (n, k). Read a column to inspect a row of Stirling contributions. Brighter squares mean larger k! S(n,k), using log10(1 + count). The white triangle has k greater than n and is outside the domain. The data table gives the plotted values.")
points = [(n, k, weighted[n][k]) for n in range(last+1) for k in range(n+1) if weighted[n][k]]
show_plot([{"type": "scatter3d", "mode": "markers", "name": "Nonzero contributions",
    "x": [n for n,k,v in points], "y": [k for n,k,v in points],
    "z": [log10(1+v) for n,k,v in points],
    "hovertemplate": "n = %{x}<br>k = %{y}<br>log10 (1 + count) = %{z:.3f}<extra></extra>",
    "marker": {"size": 5, "color": [log10(1+v) for n,k,v in points], "colorscale": "Cividis", "cmin": 0}}],
    {"title": {"text": "The contributions in 3D"},
     "scene": {"xaxis": {"title": {"text": "n"}},
               "yaxis": {"title": {"text": "k"}},
               "zaxis": {"title": {"text": "Log size"}}}},
    description="Axes: n is set size, k is block count, and height is log10(1 + k! S(n,k)). This shows the nonzero entries from the 2D map above. Each exact row sum is a Fubini number. The CSV also includes zero contributions.")
csv = "n,k,stirling,k_factorial,weighted_contribution\n"
csv += "".join(f"{n},{k},{S[n][k]},{factorial(k)},{weighted[n][k]}\n" for n in range(last+1) for k in range(n+1))
export_file("factorial-stirling-exact.csv", csv)
print(f"Stirling row n={row}:", S[row][:row+1])
print("Factorial-weighted row:", weighted[row])
print("Bell row sum:", sum(S[row]))
print("Fubini row sum:", sum(weighted[row]))
print(f"All {last+1} weighted row sums match Komorebi Math.")
print(f"Exported exact contributions through n={last}.")