"""
Photo-fractal: the photo's brightness sets the escape threshold of the
fractal iteration, so escape-time contours trace the image while the
fractal set carves into it.

Improvements over the original loop version:
  * fully vectorized with numpy  (~100x faster, so we can afford
    higher resolution + more iterations)
  * photo is resized to the render resolution and oriented correctly
  * smooth fractional escape time + histogram equalization
    -> full dynamic range instead of a washed-out n/64
  * the non-escaping interior keeps the photo instead of a flat blob
  * two render styles: colormap art + photo-colored
"""
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# ---------------- settings ----------------
W       = 1500        # render resolution
MAXIT   = 300         # iteration depth
SCALE   = 16.0        # how strongly brightness raises the escape threshold
BANDS   = 1.3         # contour-band frequency
IMG     = 'image.jpg'

# ---------------- load photo ----------------
photo = Image.open(IMG).convert('RGB').resize((W, W), Image.LANCZOS)
rgb   = np.asarray(photo, float) / 255.0
gray  = rgb.mean(2)
thresh = 0.1 + gray * SCALE          # escape threshold on |z|^2

# ---------------- complex grid ----------------
xs = np.linspace(-1, 1, W)
X, Y = np.meshgrid(xs, xs)
C = X + 1j * Y

# ---------------- vectorized iteration ----------------
Z      = np.zeros_like(C)
alive  = np.ones(C.shape, bool)
nu     = np.full(C.shape, np.nan)            # smooth escape time

for n in range(MAXIT):
    Z[alive] = Z[alive] + Z[alive]**2 + C[alive]    # same map as original
    mag2 = np.zeros(C.shape)
    mag2[alive] = Z.real[alive]**2 + Z.imag[alive]**2
    esc = alive & (mag2 > thresh)
    if esc.any():
        nu[esc] = n + np.exp(-np.sqrt(mag2[esc]))   # fractional part -> smooth
        alive &= ~esc
    if not alive.any():
        break

interior = np.isnan(nu)                       # never escaped

# ---------------- histogram-equalize escaped values ----------------
f = np.zeros(C.shape)
vals = nu[~interior]
order = vals.argsort().argsort()
f[~interior] = order / max(order.max(), 1)

band = 0.5 + 0.5 * np.sin(np.nan_to_num(nu) * BANDS)
field = np.clip(0.65 * f + 0.35 * band, 0, 1)

# ---------------- style 1: colormap art ----------------
art = plt.cm.magma(field)[..., :3]
art[interior] = rgb[interior] * 0.85          # the photo survives inside the set
plt.imsave('photofractal_art.png', np.clip(art, 0, 1))

# ---------------- style 2: photo-colored ----------------
lit = rgb * (0.30 + 0.70 * field[..., None])
glow = plt.cm.inferno(field)[..., :3] * (0.25 * field[..., None])
out = np.clip(lit + glow, 0, 1)
out[interior] = rgb[interior]
plt.imsave('photofractal_photo.png', out)

print('escaped:', (~interior).mean().round(3), ' interior:', interior.mean().round(3))