"""One shared flow, continuous particle inertia. Requires numpy and Pillow.
Run: python3 source.py. Output is written beside this file.
Palette: Flexoki by Steph Ango, MIT, https://stephango.com/flexoki
"""
from pathlib import Path
import numpy as np
from PIL import Image
SEED = 2718
SIZE = 1200
MODE = 1
PALETTE = {'black': '#100F0F', 'paper': '#FFFCF0', '50': '#F2F0E5', '100': '#E6E4D9', '150': '#DAD8CE', '200': '#CECDC3', '300': '#B7B5AC', '400': '#9F9D96', '500': '#878580', '600': '#6F6E69', '700': '#575653', '800': '#403E3C', '850': '#343331', '900': '#282726', '950': '#1C1B1A', 'red-50': '#FFE1D5', 'red-100': '#FFCABB', 'red-150': '#FDB2A2', 'red-200': '#F89A8A', 'red-300': '#E8705F', 'red-400': '#D14D41', 'red-500': '#C03E35', 'red-600': '#AF3029', 'red-700': '#942822', 'red-800': '#6C201C', 'red-850': '#551B18', 'red-900': '#3E1715', 'red-950': '#261312', 'orange-50': '#FFE7CE', 'orange-100': '#FED3AF', 'orange-150': '#FCC192', 'orange-200': '#F9AE77', 'orange-300': '#EC8B49', 'orange-400': '#DA702C', 'orange-500': '#CB6120', 'orange-600': '#BC5215', 'orange-700': '#9D4310', 'orange-800': '#71320D', 'orange-850': '#59290D', 'orange-900': '#40200D', 'orange-950': '#27180E', 'yellow-50': '#FAEEC6', 'yellow-100': '#F6E2A0', 'yellow-150': '#F1D67E', 'yellow-200': '#ECCB60', 'yellow-300': '#DFB431', 'yellow-400': '#D0A215', 'yellow-500': '#BE9207', 'yellow-600': '#AD8301', 'yellow-700': '#8E6B01', 'yellow-800': '#664D01', 'yellow-850': '#503D02', 'yellow-900': '#3A2D04', 'yellow-950': '#241E08', 'green-50': '#EDEECF', 'green-100': '#DDE2B2', 'green-150': '#CDD597', 'green-200': '#BEC97E', 'green-300': '#A0AF54', 'green-400': '#879A39', 'green-500': '#768D21', 'green-600': '#66800B', 'green-700': '#536907', 'green-800': '#3D4C07', 'green-850': '#313D07', 'green-900': '#252D09', 'green-950': '#1A1E0C', 'cyan-50': '#DDF1E4', 'cyan-100': '#BFE8D9', 'cyan-150': '#A2DECE', 'cyan-200': '#87D3C3', 'cyan-300': '#5ABDAC', 'cyan-400': '#3AA99F', 'cyan-500': '#2F968D', 'cyan-600': '#24837B', 'cyan-700': '#1C6C66', 'cyan-800': '#164F4A', 'cyan-850': '#143F3C', 'cyan-900': '#122F2C', 'cyan-950': '#101F1D', 'blue-50': '#E1ECEB', 'blue-100': '#C6DDE8', 'blue-150': '#ABCFE2', 'blue-200': '#92BFDB', 'blue-300': '#66A0C8', 'blue-400': '#4385BE', 'blue-500': '#3171B2', 'blue-600': '#205EA6', 'blue-700': '#1A4F8C', 'blue-800': '#163B66', 'blue-850': '#133051', 'blue-900': '#12253B', 'blue-950': '#101A24', 'purple-50': '#F0EAEC', 'purple-100': '#E2D9E9', 'purple-150': '#D3CAE6', 'purple-200': '#C4B9E0', 'purple-300': '#A699D0', 'purple-400': '#8B7EC8', 'purple-500': '#735EB5', 'purple-600': '#5E409D', 'purple-700': '#4F3685', 'purple-800': '#3C2A62', 'purple-850': '#31234E', 'purple-900': '#261C39', 'purple-950': '#1A1623', 'magenta-50': '#FEE4E5', 'magenta-100': '#FCCFDA', 'magenta-150': '#F9B9CF', 'magenta-200': '#F4A4C2', 'magenta-300': '#E47DA8', 'magenta-400': '#CE5D97', 'magenta-500': '#B74583', 'magenta-600': '#A02F6F', 'magenta-700': '#87285E', 'magenta-800': '#641F46', 'magenta-850': '#4F1B39', 'magenta-900': '#39172B', 'magenta-950': '#24131D'}

def render():
    rng=np.random.default_rng(SEED)
    n=75000
    y=rng.uniform(-0.30,1.35,n)
    x=np.full(n,-0.22)
    response=rng.uniform(0,1,n)
    tau=0.008+0.22*response**1.7
    vx=np.full(n,.57); vy=np.full(n,-.18)
    maps=np.zeros((6,SIZE,SIZE),dtype=np.float32)
    band=np.minimum((response*6).astype(int),5)
    # One family of unequal vortices. MODE modifies only their geometry.
    eddies=[(.24,.75, .082,.21),(.61,.34,-.094,.23),(.94,.63,.056,.22)]
    if MODE == 1: eddies=[(.20,.82,.086,.23),(.60,.38,-.118,.26),(.96,.69,.065,.20)]
    if MODE == 2: eddies=[(.20,.82,.086,.23),(.57,.35,-.118,.26),(.96,.69,.065,.20)]
    if MODE == 3: eddies=[(.20,.82,.086,.23),(.60,.38,-.118,.29),(.96,.69,.065,.20)]
    dt=.002
    for step in range(1820):
        u=np.full(n,.57); v=np.full(n,-.18)
        for cx,cy,strength,width in eddies:
            dx=x-cx;dy=y-cy;rr=dx*dx+dy*dy
            factor=strength*np.exp(-rr/(2*width*width))/(rr+.012)
            u-=dy*factor;v+=dx*factor
        relax=1-np.exp(-dt/tau)
        vx+=(u-vx)*relax;vy+=(v-vy)*relax
        x+=vx*dt;y+=vy*dt
        if True:
            xx=(x*SIZE).astype(int);yy=(y*SIZE).astype(int)
            good=(xx>=0)&(xx<SIZE)&(yy>=0)&(yy<SIZE)
            np.add.at(maps,(band[good],yy[good],xx[good]),1)
    # Separable binomial convolution works in float without introducing display colors.
    for axis in (1,2): maps=(np.roll(maps,1,axis)+2*maps+np.roll(maps,-1,axis))/4
    total=maps.sum(axis=0)
    mean=(maps*np.arange(6)[:,None,None]).sum(axis=0)/(total+1e-8)
    mix=maps.max(axis=0)/(total+1e-8)
    colors=np.empty((SIZE,SIZE,3),dtype=np.uint8)
    colors[:]=tuple(bytes.fromhex(PALETTE['black'][1:]))
    families=[['cyan-950','cyan-900','cyan-700','cyan-600','cyan-400','cyan-200','cyan-100'],
              ['blue-950','blue-900','blue-700','blue-600','blue-400','blue-200','blue-100'],
              ['orange-950','orange-900','orange-700','orange-600','orange-400','orange-200','orange-100']]
    lo,hi=np.quantile(mean[total>2],[.12,.88])
    groups=np.where(mean<lo,0,np.where(mean<hi,1,2))
    # Quantile bands preserve low-density folded lanes, not a painted overlay.
    q=np.quantile(total[total>0],[.12,.35,.58,.78,.91,.975])
    level=np.searchsorted(q,total)
    for g,family in enumerate(families):
        for k,key in enumerate(family): colors[(groups==g)&(level==k)&(total>0)]=tuple(bytes.fromhex(PALETTE[key][1:]))
    # Near-white is reserved for genuinely crowded crossings.
    colors[total>np.quantile(total[total>0],.994)]=tuple(bytes.fromhex(PALETTE['paper'][1:]))
    Image.fromarray(colors).save(Path(__file__).with_name('render-1.png'))

if __name__=='__main__': render()