;Smoke Demo
;By Jim Shaw 19/4/2005
;Uses Navier-Stokes density equations To produce a smoke Field.
;Based on (not-very real-time and slightly broken) source from 
;Real-Time Fluid Dynamics For Games
;available here
;<http://www.dgp.toronto.edu/people/stam/reality/index.html>

Graphics 320,240,0,2

Text 10,10,"Press SPACE for full screen"
Text 10,25,"Any other key for windowed"

key=WaitKey()
If key=32 Then
	mode=0
	bpp=32
Else
	mode=2
	bpp=0
End If

Graphics 320,240,bpp,mode

SetBuffer BackBuffer()

Const N%=200
Const ITER%=3
Const Nplus2%=N+2
Const size%=(N+2)*(N+2)
Global u#[size]
Global v#[size]
Global u_prev#[size]
Global v_prev#[size]
Global dens#[size]
Global dens_prev#[size]
Global difftmp#[size]

Global xo%=(GraphicsWidth()-N)/2
Global yo%=(GraphicsHeight()-N)/2

Global wind_dir#
Global wind_mag#
autowind%=True
wind_mag=0.005
wind_dir=0

Global dens_mag#
dens_mag=1

Const dt#=1
Const diff#=0.00001
Const visc#=0.00001

Global mx%,my%
Global inwindow%=0

Const palsize%=255
Dim palette(palsize)

SeedRnd(MilliSecs())

Global b%=CreateBank(4)
Global log2#=Log(2)

Const mbits%=8
Const mcount%=1 Shl mbits
Const mmask%=mcount-1

Dim logtable#(mcount)
For x#=0 To mcount
	logtable(x)=Log(1.0+(x/Float(mcount)))
Next

pal%=1
createpalette(pal)
init(dens_prev, u_prev, v_prev)

sx%=N/2
sy%=N-25
Repeat
	t%=MilliSecs()
	mx = MouseX()
	my = MouseY()

	If autowind Then
		wind_dir=180+20*Sin(q)+Sin(q/3.0)
		q=q+10
	End If
	wind(wind_dir, wind_mag, dens_prev, u_prev, v_prev, dt)
	get_from_UI(dens_prev, u_prev, v_prev)

	dens_prev[IX(sx,sy)]=dens_prev[IX(sx,sy)]+dens_mag
	u_prev[IX(sx,sy)]=u_prev[IX(sx,sy)]+Rnd(2)-1
	v_prev[IX(sx,sy)]=v_prev[IX(sx,sy)]-Rnd(2)
	vel_step(N, u, v, u_prev, v_prev, visc, dt)

	dens_step(N, dens, dens_prev, u, v, diff, dt)

;	Cls
;	draw_vel(N, u, v)
	
	LockBuffer BackBuffer()
	draw_dens(N, dens)
	showpal()
	UnlockBuffer BackBuffer()
	
	s%=MilliSecs()
	fps#=1000.0/(s-t)
	Color 0,0,0
	Rect 0,N+10,xo-3,20,1
	Color 255,255,255
	Text 0,N+10,fps
	If inwindow WritePixel mx,my,$ff0000
	Rect xo-2,yo-2,N+4,N+4,0	
	
	Flip False
	
	If KeyDown(57) Then
		createpalette(pal)
		pal=pal+1
		pal=pal Mod 3
	End If
	If KeyDown(51) Then
		wind_dir=wind_dir+5
		autowind=False
	End If
	If KeyDown(52) Then
		wind_dir=wind_dir-5
		autowind=False
	End If
	If KeyDown(203) Then
		sx=sx-1
		If sx<0 sx=0
	Else If KeyDown(205) Then
		sx=sx+1
		If sx>N sx=N
	End If
	If KeyDown(200) Then 
		sy=sy-1
		If sy<0 sy=0
	Else If KeyDown(208) Then
		sy=sy+1
		If sy>N sy=N
	End If
	If KeyDown(44) Then
		wind_mag=wind_mag-0.001
		If wind_mag<0 wind_mag=0
	Else If KeyDown(30) Then
		wind_mag=wind_mag+0.001
		If wind_mag>0.01 wind_mag=0.01
	End If
	If KeyDown(31) Then
		dens_mag=dens_mag+0.01
	Else If KeyDown(45) Then
		dens_mag=dens_mag-0.01
		If dens_mag<0 dens_mag=0
	End If
	If KeyDown(46) Then
		For i=0 To size
			dens_prev[i]=0
			dens[i]=0
		Next
	End If
Until KeyDown(1)
End

Function init(d#[size], u#[size], v#[size])
	Local i
	For i=0 To size
		d[i]=0
		u[i]=0
		v[i]=0
	Next
End Function

Function get_from_UI(d#[size], u#[size], v#[size])
	Local dmx%,dmy%
	inwindow=0
	If mx>=xo And mx<=xo+N Then
		If my>=yo And my<=yo+N Then
			inwindow=1
			If MouseDown(1) Then
				dmx=mx-xo
				dmy=my-yo
				d[dmx+dmy*Nplus2]=d[dmx+dmy*Nplus2]+1
				u[dmx+dmy*Nplus2]=u[dmx+dmy*Nplus2]+Rnd(4)-2
				v[dmx+dmy*Nplus2]=v[dmx+dmy*Nplus2]+Rnd(4)-2
			End If
		End If
	End If
End Function

Function wind(angle#, mag#, d#[size], u#[size], v#[size], dt#)
	Local sn#,cz#
	Local i%
	Local x%,y%
	sn = Sin(angle)*mag*dt
	cz = Cos(angle)*mag*dt
	For i=0 To size
		u[i] = sn
		v[i] = cz
		d[i] = 0
	Next
	For i=0 To N+1
		u[i*Nplus2]=0
		v[i*Nplus2]=0
		u[i]=0
		v[i]=0
		u[i*Nplus2+Nplus2-1]=0
		v[i*Nplus2+Nplus2-1]=0
		u[i+(N+1)*Nplus2]=0
		v[i+(N+1)*Nplus2]=0
	Next
End Function

Function draw_dens(N%, d#[size])
	Local maxd#
	Local x%,y%
	Local idx%
	
	maxd=110
	maxd=-255.0/maxd
	LockBuffer BackBuffer()
	idx=0
	For y=yo To yo+N-1
		For x=xo To xo+N-1
			If d[idx]>1 Then
				WritePixelFast x,y,$ffffff
			Else If d[idx]>0 Then
				;WritePixelFast x,y,palette(255-Int(maxd*Log(d[idx])))
				PokeFloat(b,0,d[idx])
				k%=PeekInt(b,0)
				m#=logtable((k Shr 15) And 255)+log2*(((k Shr 23) And $ff)-127)
				WritePixelFast x,y,palette(255-Int(maxd*m))
			Else
				WritePixelFast x,y,$000000
			End If
			idx=idx+1
		Next
		idx=idx+2
	Next
	UnlockBuffer BackBuffer()
End Function

Function draw_vel(N%, u#[size], v#[size])
	Local x%,y%
	Local uu#,vv#,mag#
	Color 255,0,0
	For y=0 To N-1
		For x=0 To N-1
			uu=u[IX(x,y)]
			vv=v[IX(x,y)]
			mag=uu*uu+vv*vv
			If mag <> 0 Then
				mag = 1.0/Sqr(mag)
				uu=uu*mag*3
				vv=vv*mag*3
				Line x*5,y*5,x*5+uu,y*5+vv
			End If
		Next
	Next
End Function

Function IX%(i%,j%)
	Return i+Nplus2*j
End Function

Function add_source(x#[size], s#[size], dt#)
	Local i%
	For i=0 To size-1
		x[i] = x[i] + dt*s[i]
	Next
End Function

Function diffuse(N%, b%, x#[size], x0#[size], diff#, dt#)
	Local i%,j%,k%
	Local a#,a2#
	Local idx%
	
	a=dt*diff*N*N
	a2#=1.0/(1.0+4.0*a)
	For k=0 To ITER
		idx=Nplus2+1
		For i=1 To N
			For j=1 To N
				difftmp[idx] = (x0[idx] + a*(x[idx-1]+x[idx+1]+x[idx-Nplus2]+x[idx+Nplus2]))*a2
				idx=idx+1
			Next
			idx=idx+2
		Next
		idx=Nplus2
		For i=1 To N
			For j=1 To N
				x[idx]=difftmp[idx]
				idx=idx+1
			Next
			idx=idx+2
		Next
		set_bnd(N, b, x)
	Next
End Function

Function advect(N%, b%, d#[size], d0#[size], u#[size], v#[size], dt#)
	Local i%, j%, i0%, j0%, i1%, j1%
	Local x#, y#, s0#, t0#, s1#, t1#, dt0#
	Local idx%,idx2%
	
	dt0 = dt*N
	
	idx=Nplus2+1
	For i=1 To N
		For j=1 To N
			x = Float(j)-dt0*u[idx]
			y = Float(i)-dt0*v[idx]
			
			
			If x<0.5 Then : x=0.5
			Else If x>N+0.5 Then : x=N+0.5
			End If
			i0=x-0.5
			
			If y<0.5 Then : y=0.5
			Else If y>N+0.5 Then : y=N+0.5
			End If
			j0=y-0.5
			
			s1 = x-Float(i0)
			s0 = 1.0-s1
			t1 = y-Float(j0)
			t0 = 1.0-t1

			idx2=i0+j0*Nplus2
			d[idx] = s0*(t0*d0[idx2] + t1*d0[idx2+Nplus2]) + s1*(t0*d0[idx2+1] + t1*d0[idx2+Nplus2+1])
			idx=idx+1
		Next
		idx=idx+2
	Next
	set_bnd(N, b, d)
End Function

Function dens_step(N%, x#[size], x0#[size], u#[size], v#[size], diff#, dt#)
	add_source(x, x0, dt)
	diffuse(N, 0, x0, x, diff, dt)
	advect(N, 0, x, x0, u, v, dt)
End Function

Function vel_step(N%, u#[size], v#[size], u0#[size], v0#[size], visc#, dt#)
	add_source(u, u0, dt)
	add_source(v, v0, dt)
	
	diffuse(N, 1, u0, u, visc, dt)
	diffuse(N, 2, v0, v, visc, dt)
	project(N, u0, v0, u, v)
	
	advect(N, 1, u, u0, u0, v0, dt)
	advect(N, 2, v, v0, u0, v0, dt)
	project(N, u, v, u0, v0)
End Function

Function project(N%, u#[size], v#[size], p#[size], div#[size])
	Local i%, j%, k%
	Local h#,h2#
	Local idx%
	
	h = 1.0/Float(N)
	idx=Nplus2+1
	For i=1 To N
		For j=1 To N
			div[idx] = -0.5*h*(u[idx+1]-u[idx-1]+v[idx+Nplus2]-v[idx-Nplus2])
			p[idx] = 0
			idx=idx+1
		Next
		idx=idx+2
	Next
	set_bnd(N, 0, div)
	set_bnd(N, 0, p)
	For k=0 To ITER
		idx=Nplus2+1
		For i=1 To N
			For j=1 To N
				difftmp[idx] = (div[idx]+p[idx-1]+p[idx+1]+p[idx-Nplus2]+p[idx+Nplus2])*0.25		
				idx=idx+1
			Next
			idx=idx+2
		Next
		idx=Nplus2+1
		For i=1 To N
			For j=1 To N
				p[idx]=difftmp[idx]
				idx=idx+1
			Next
			idx=idx+2
		Next
		set_bnd(N, 0, p)
	Next
	
	idx=Nplus2+1
	h2#=0.5/h
	For i=1 To N
		For j=1 To N
			u[idx] = u[idx] - ((p[idx+1]-p[idx-1])*h2)
			v[idx] = v[idx] - ((p[idx+Nplus2]-p[idx-Nplus2])*h2)
			idx=idx+1
		Next
		idx=idx+2
	Next
	set_bnd(N, 1, u)
	set_bnd(N, 2, v)
End Function

Function set_bnd(N%, b%, x#[size])
	Local i%
	For i=1 To N
		If b=1 Then
			x[IX(0 ,i)]  = -x[IX(1,i)]
			x[IX(N+1,i)] = -x[IX(N,i)]
		Else
			x[IX(0 ,i)]  =  x[IX(1,i)]
			x[IX(N+1,i)] =  x[IX(N,i)]
		End If
		If b=2 Then
			x[IX(i,0 )]  = -x[IX(i,1)]
			x[IX(i,N+1)] = -x[IX(i,N)]
		Else
			x[IX(i,0 )]  =  x[IX(i,1)]
			x[IX(i,N+1)] =  x[IX(i,N)]
		End If
	Next
	x[IX(0 ,0 )]   = 0.5*(x[IX(1,0 )]+x[IX(0 ,1)])
	x[IX(0 ,N+1)]  = 0.5*(x[IX(1,N+1)]+x[IX(0 ,N )])
	x[IX(N+1,0 )]  = 0.5*(x[IX(N,0 )]+x[IX(N+1,1)])
	x[IX(N+1,N+1)] = 0.5*(x[IX(N,N+1)]+x[IX(N+1,N )])
End Function


Function spread(r0#,g0#,b0#, r1#,g1#,b1#, start#, length#)
	Local dr#,dg#,db#
	dr = (r1-r0)/length
	dg = (g1-g0)/length
	db = (b1-b0)/length
	For i=start To start+length-1
		palette(i)=r0 Shl 16 + g0 Shl 8 + b0
		r0 = r0 + dr
		g0 = g0 + dg
		b0 = b0 + db
	Next
End Function

Function createpalette(p%)
	Local s%
	palette(0)=0
	s=1
	If p=0 Then
		spread(0,0,0,     255,255,255,0,palsize+1)
	Else If p=1 Then
		spread(0,0,0,     0,0,0,      s,208) : s=s+208
		spread(0,0,0,   255,255,255,s,42) : s=s+42
		spread(255,255,255,255,255,255,s,palsize-s+1)
	Else If p=2 Then
		spread(255,0,255, 0,0,255,    s,42) : s=s+42
		spread(0,0,255,   0,255,0,    s,42) : s=s+42
		spread(0,255,0,   255,255,0,  s,42) : s=s+42
		spread(255,255,0, 255,128,0,  s,42) : s=s+42
		spread(255,128,0, 255,0,0,    s,42) : s=s+42
		spread(255,0,0,   255,255,255,s,42) : s=s+42
		spread(255,255,255,255,255,255,s,palsize-s+1)
	End If
End Function

Function showpal()
	For x=0 To palsize
	For y=0 To 3
	WritePixelFast x,y,palette(x)
	Next
	Next
End Function