import marimo

__generated_with = "0.25.1"
app = marimo.App(width="medium")


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        # Zooming into the Mandelbrot set

        Take a complex number $c$, start from $z = 0$, and repeat
        $z \mapsto z^2 + c$. If $z$ stays bounded forever, $c$ belongs to the
        Mandelbrot set. That one rule draws a boundary with unlimited detail:
        however far you zoom, new structure appears, including small copies of
        the whole set. Mitsuhiro Shishikura proved in 1998 that the boundary is
        so convoluted that its Hausdorff dimension is 2, the same as a filled
        area.

        Pick a place and zoom in. Each step of the slider magnifies ten times.
        Ordinary 64-bit numbers run out of precision near a magnification of
        $10^{13}$; beyond that, deep-zoom programs use arbitrary precision.
        """
    )
    return


@app.cell
def _():
    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    return mo, np, plt


@app.cell
def _(mo):
    PLACES = {
        "Seahorse valley": complex(-0.743643887037151, 0.131825904205330),
        "Elephant valley": complex(0.285, 0.0115),
        "A small copy of the set": complex(-1.7685736562, 0.0007),
        "Spiral": complex(-0.761574, -0.0847596),
        "Whole set": complex(-0.6, 0.0),
    }
    place = mo.ui.dropdown(list(PLACES), value="Seahorse valley", label="Place")
    zoom = mo.ui.slider(0, 12, step=0.25, value=3, label="Magnification (powers of ten)")
    mo.vstack([place, zoom])
    return PLACES, place, zoom


@app.cell
def _(np):
    def escape_time(center, width, pixels=(480, 320), max_iter=500):
        """Smooth iteration count for each pixel; NaN inside the set.

        Iterates only the points that have not escaped yet, which keeps deep
        zooms fast where most of the frame is inside the set.
        """
        nx, ny = pixels
        x = center.real + (np.arange(nx) - nx / 2) * width / nx
        y = center.imag + (np.arange(ny) - ny / 2) * width / nx
        c = (x[None, :] + 1j * y[::-1, None]).ravel()
        z = np.zeros_like(c)
        result = np.full(c.shape, np.nan)
        active = np.arange(c.size)
        for n in range(max_iter):
            z = z * z + c[active]
            escaped = np.abs(z) > 256
            if escaped.any():
                done = active[escaped]
                result[done] = n + 1 - np.log2(np.log(np.abs(z[escaped])))
                active, z = active[~escaped], z[~escaped]
            if active.size == 0:
                break
        return result.reshape(ny, nx)

    return (escape_time,)


@app.cell
def _(PLACES, escape_time, mo, np, place, plt, zoom):
    width = 3.2 / 10**zoom.value
    iterations = int(300 + 150 * zoom.value)
    with mo.status.spinner(title=f"Iterating up to {iterations:,} times per pixel"):
        counts = escape_time(PLACES[place.value], width, max_iter=iterations)
    shade = np.sqrt(np.nan_to_num(counts, nan=0.0))
    fig, ax = plt.subplots(figsize=(7.2, 4.8))
    ax.imshow(np.where(np.isnan(counts), np.nan, np.sin(0.35 * shade) ** 2), cmap="twilight_shifted", interpolation="bilinear")
    ax.set_facecolor("black")
    ax.set_axis_off()
    ax.set_title(f"{place.value}, magnified {10**zoom.value:,.0f}×, frame width {width:.2e}")
    fig.tight_layout()
    fig
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        ## How much of the plane is inside?

        The area of the Mandelbrot set is not known in closed form. A Monte
        Carlo estimate samples random points in a box around the set and
        counts how many stay bounded. The error shrinks like
        $1/\sqrt{N}$, so each extra digit costs a hundred times more samples.
        Published estimates put the area near 1.5066. A point that escapes
        only after more iterations than the limit is counted as inside, so a
        short run overestimates the area; raise the limit in your fork to see
        the estimate fall.
        """
    )
    return


@app.cell
def _(mo):
    samples = mo.ui.slider(10_000, 400_000, step=10_000, value=100_000, label="Random samples")
    samples
    return (samples,)


@app.cell
def _(mo, np, samples):
    rng = np.random.default_rng(0)
    c = rng.uniform(-2.0, 0.5, samples.value) + 1j * rng.uniform(-1.25, 1.25, samples.value)
    z = np.zeros_like(c)
    inside = np.ones(c.size, dtype=bool)
    for _ in range(400):
        z[inside] = z[inside] ** 2 + c[inside]
        inside &= np.abs(z) <= 2
    box = 2.5 * 2.5
    fraction = inside.mean()
    area = box * fraction
    error = box * np.sqrt(fraction * (1 - fraction) / c.size)
    mo.md(f"Estimated area: **{area:.4f} ± {error:.4f}** from {c.size:,} samples and 400 iterations.")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(
        r"""
        ---
        **Method and limits.** Escape radius 256 with the normalised
        iteration count $n + 1 - \log_2 \log |z|$ for smooth colour; points
        that have not escaped within the iteration limit are drawn black and
        may still lie outside the set. 64-bit floating point throughout.
        **References.** B. Mandelbrot, *Ann. N.Y. Acad. Sci.* 357 (1980) 249;
        M. Shishikura, *Ann. of Math.* 147 (1998) 225.
        Fork it to colour by a different rule or to add perturbation theory
        for deeper zooms.
        """
    )
    return


if __name__ == "__main__":
    app.run()