USE gnuplot
! USE rtchecks

INTEGER nx=16,ny=16,nz=16
REAL a=0.1, rad=0.1
LOOP FOR i=1 TO COMMANDLINE.HI
  IF COMMANDLINE(i,0..2)="nx=" THEN
    nx=atoi(COMMANDLINE(i,3+(0..HI-3))); ny=nx; nz=ny
  ELSE IF COMMANDLINE(i,0..2)="ny=" THEN
    ny=atoi(COMMANDLINE(i,3+(0..HI-3))); nz=ny
  ELSE IF COMMANDLINE(i,0..2)="nz=" THEN
    nz=atoi(COMMANDLINE(i,3+(0..HI-3)))
  ELSE IF COMMANDLINE(i,0..1)="a=" THEN
    a=atof(COMMANDLINE(i,2+(0..HI-2)))
  ELSE IF COMMANDLINE(i,0..1)="rad=" THEN
    rad=atof(COMMANDLINE(i,4+(0..HI-4)))
  ELSE ERROR "unknown option"  
REPEAT

baseline=1E-12+1/8
staircase=NO
nshapes=ENUM(sinriblet,gaussian,smoothedcyl,cylinder,cone,spheroid,squarecyl,ellipsoid)
shape=spheroid

BOOLEAN FUNCTION InBody(REAL x,y,z)
CASE shape OF
sinriblet: ! sinusoidal riblet
  RESULT=z<baseline+a*SIN(x*PI)^2
gaussian:
  RESULT=z<baseline+a*EXP{-160*[(x-0.5)^2+(y-0.5)^2]}
smoothedcyl: ! smoothed cylinder
  RESULT=IF (x-0.5)^2+(y-0.5)^2>4*rad^2 THEN z<baseline ELSE z<baseline+a/{1+EXP[10*(SQRT[(x-0.5)^2+(y-0.5)^2]/rad-1)]}
cylinder:
  RESULT=IF (x-0.5)^2+(y-0.5)^2<rad^2 THEN z<baseline+a ELSE z<baseline
cone:
  RESULT=IF SQRT[(x-0.5)^2+(y-0.5)^2]<rad THEN z<baseline+a*{1-SQRT[(x-0.5)^2+(y-0.5)^2]/rad} ELSE z<baseline
spheroid:
  RESULT=IF (x-0.5)^2+(y-0.5)^2<rad^2 THEN z<baseline+a*SQRT{1-[(x-0.5)^2+(y-0.5)^2]/rad^2} ELSE z<baseline
squarecyl: ! square cylinder
  RESULT=IF ABS(x-0.5)<rad AND ABS(y-0.5)<rad THEN z<baseline+a ELSE z<baseline
ellipsoid: ! 3 to 1 ellipsoid
  RESULT=IF [(x-0.5)/0.05]^2+[(y-0.5)/0.15]^2<1 THEN z<baseline+a*SQRT{1-[(x-0.5)/0.05]^2-[(y-0.5)/0.15]^2} ELSE z<baseline
END CASE
END InBody

REAL FUNCTION Bisection[FUNCTION(REAL s)->BOOLEAN f; REAL VARIABLE s1, s2]
f1 = f(s1); f2 = f(s2)
IF f1 = f2 THEN ERROR "Bisection point not included in " s1 " " s2
LOOP halve
  RESULT=(s1+s2)/2
  IF ABS(s2-s1)<1E-10 THEN EXIT Bisection
  IF f(RESULT)=f1 THEN s1=RESULT ELSE s2=RESULT
REPEAT halve
END Bisection

! laplstar=STRUCTURE(REAL c0,cn)
laplstar=STRUCTURE(REAL c0)
boundarypoint=STRUCTURED ARRAY(u,v,w) OF laplstar

boundarypoint stdpoint

USE mg3d

POINTER TO boundarypoint bcond(0..nx,0..ny,0..nz)
dx=1/nx; cn=1/dx
WITH stdpoint.u: c0=dx*dx/6 ! ; cn=1/dx
WITH stdpoint.v: c0=dx*dx/6 ! ; cn=1/dx
WITH stdpoint.w: c0=dx*dx/6 ! ; cn=1/dx
DO bcond(ix,iy,iz)=stdpoint FOR ALL ix,iy,iz
DO bcond(ix,iy,0)=NULL; bcond(ix,iy,nz)=NULL FOR ALL ix,iy
ARRAY(-1..nx+1,-1..ny+1,0..nz) OF VARS field=0,rhs=0

#define imblapl(f) [f(1,0,0)+f(-1,0,0)+f(0,1,0)+f(0,-1,0)+f(0,0,1)+f(0,0,-1)]/dx/dx
imbueq==imblapl(u)-cn*[p(1,0,0)-p(0,0,0)]
imbveq==imblapl(v)-cn*[p(0,1,0)-p(0,0,0)]
imbweq==imblapl(w)-cn*[p(0,0,1)-p(0,0,0)]

SUBROUTINE imbresids
LOOP FOR ix=1 TO field.HI1-1 AND iy=1 TO field.HI2-1 AND iz=1 TO field.HI3-1
  EXCEPT bcond(ix,iy,iz)=NULL
  REAL cont=0, dcont=0, dcont0=0
  WITH field(ix+*,iy+*,iz+*)
  IF bcond(ix MOD HI +1,iy,iz)#NULL THEN WITH bcond(ix,iy,iz).u rhs(ix,iy,iz).u=(c0*imbueq-u(0,0,0))*(6/dx/dx); cont=~+cn*u(0,0,0); dcont=~+cn*cn*c0; dcont0=~+1/6
  IF bcond(ix,iy MOD HI +1,iz)#NULL THEN WITH bcond(ix,iy,iz).v rhs(ix,iy,iz).v=(c0*imbveq-v(0,0,0))*(6/dx/dx); cont=~+cn*v(0,0,0); dcont=~+cn*cn*c0; dcont0=~+1/6
  IF bcond(ix,iy,iz+1)#NULL THEN WITH bcond(ix,iy,iz).w rhs(ix,iy,iz).w=(c0*imbweq-w(0,0,0))*(6/dx/dx); cont=~+cn*w(0,0,0); dcont=~+cn*cn*c0; dcont0=~+1/6
  IF bcond(ix-1,iy,iz)#NULL THEN WITH bcond(ix-1,iy,iz).u cont=~-cn*u(-1,0,0); dcont=~+cn*cn*c0; dcont0=~+1/6
  IF bcond(ix,iy-1,iz)#NULL THEN WITH bcond(ix,iy-1,iz).v cont=~-cn*v(0,-1,0); dcont=~+cn*cn*c0; dcont0=~+1/6
  IF bcond(ix,iy,iz-1)#NULL THEN WITH bcond(ix,iy,iz-1).w cont=~-cn*w(0,0,-1); dcont=~+cn*cn*c0; dcont0=~+1/6
  rhs(ix,iy,iz).p=cont*dcont0/dcont
REPEAT
rhs(0,*,*)=rhs(HI-1,*,*); rhs(HI,*,*)=rhs(1,*,*)
rhs(*,0,*)=rhs(*,HI-1,*); rhs(*,HI,*)=rhs(*,1,*)
END imbresids

SUBROUTINE imbsmooth
LOOP FOR parity=0 TO 1
  LOOP FOR ix=1 TO field.HI1-1 AND iy=1 TO field.HI2-1 AND iz=2-(ix+iy+parity) MOD 2 TO field.HI3-1 BY 2
    EXCEPT bcond(ix,iy,iz)=NULL
    REAL ut(-1..0),dtu(-1..0),vt(-1..0),dtv(-1..0),wt(-1..0),dtw(-1..0)
    REAL cont=0,dcont=0
    LOOP FOR dix=-1 TO 0 WITH field(ix+dix+*,iy+*,iz+*)
      IF bcond((ix+2*dix) MOD HI +1,iy,iz)=NULL THEN ut(dix)=u(0,0,0); dtu(dix)=0 ELSE
        WITH bcond(ix+dix,iy,iz).u
        dtu(dix)=cn*c0
        u(0,0,0)=imbueq*c0
        IF dix=0 THEN cont=~+cn*u(0,0,0) ELSE cont=~-cn*u(0,0,0)
        dcont=~+cn*dtu(dix)
      END IF
    REPEAT
    LOOP FOR diy=-1 TO 0 WITH field(ix+*,iy+diy+*,iz+*)
      IF bcond(ix,(iy+2*diy) MOD HI +1,iz)=NULL THEN vt(diy)=v(0,0,0); dtv(diy)=0 ELSE
        WITH bcond(ix,iy+diy,iz).v
        dtv(diy)=cn*c0
        v(0,0,0)=imbveq*c0
        IF diy=0 THEN cont=~+cn*v(0,0,0) ELSE cont=~-cn*v(0,0,0)
        dcont=~+cn*dtv(diy)
      END IF
    REPEAT
    LOOP FOR diz=-1 TO 0 WITH field(ix+*,iy+*,iz+diz+*)
      IF bcond(ix,iy,iz+2*diz+1)=NULL THEN wt(diz)=w(0,0,0); dtw(diz)=0 ELSE
        WITH bcond(ix,iy,iz+diz).w
        dtw(diz)=cn*c0
        w(0,0,0)=imbweq*c0
        IF diz=0 THEN cont=~+cn*w(0,0,0) ELSE cont=~-cn*w(0,0,0)
        dcont=~+cn*dtw(diz)
      END IF
    REPEAT
    WITH field(ix+*,iy+*,iz+*)
    cont=cont/dcont
    p(0,0,0)=~-cont
    u(0,0,0)=~-dtu(0)*cont
    u(-1,0,0)=~+dtu(-1)*cont
    v(0,0,0)=~-dtv(0)*cont
    v(0,-1,0)=~+dtv(-1)*cont
    w(0,0,0)=~-dtw(0)*cont
    w(0,0,-1)=~+dtw(-1)*cont
  REPEAT
  mirrorbc(field)
REPEAT
END imbsmooth

SUBROUTINE updatebcond(POINTER TO POINTER TO boundarypoint bp; INTEGER vcomp; REAL d)
  IF bp=^stdpoint THEN NEW bp; bp^=stdpoint
  WITH bp(vcomp)
  c0=c0*d/[c0/dx/dx*(dx-d)+d]
!  cn=~*(d+dx/2)^2/2/d/dx  !  SQRT(1/c0/6)  !  0.5*(dx+d)/dx*cn
END updatebcond

SUBROUTINE calcbcond(POINTER TO POINTER TO boundarypoint bp; INTEGER vcomp; REAL x,y,z)
IF InBody(x,y,z) THEN updatebcond(bp,vcomp,0) ELSE IF NOT staircase THEN
  IF InBody(x-dx,y,z) THEN updatebcond(bp,vcomp,x-Bisection[InBody(s,y,z),x-dx,x])
  IF InBody(x+dx,y,z) THEN updatebcond(bp,vcomp,Bisection[InBody(s,y,z),x,x+dx]-x)
  IF InBody(x,y-dx,z) THEN updatebcond(bp,vcomp,y-Bisection[InBody(x,s,z),y-dx,y])
  IF InBody(x,y+dx,z) THEN updatebcond(bp,vcomp,Bisection[InBody(x,s,z),y,y+dx]-y)
  IF InBody(x,y,z-dx) THEN updatebcond(bp,vcomp,z-Bisection[InBody(x,y,s),z-dx,z])
  IF InBody(x,y,z+dx) THEN updatebcond(bp,vcomp,Bisection[InBody(x,y,s),z,z+dx]-z)
END IF
END calcbcond

REAL mh=0,mh2=0
LOOP FOR ix=nx DOWN TO 1 AND iy=ny DOWN TO 1 AND iz=nz-1 DOWN TO 1
  IF InBody[(ix+0.5)*dx,iy*dx,iz*dx] AND InBody[(ix-0.5)*dx,iy*dx,iz*dx] AND
     InBody[ix*dx,(iy+0.5)*dx,iz*dx] AND InBody[ix*dx,(iy+0.5)*dx,iz*dx] AND
     InBody[ix*dx,iy*dx,(iz+0.5)*dx] AND InBody[ix*dx,iy*dx,(iz-0.5)*dx] THEN
    bcond(ix,iy,iz)=NULL
  ELSE
    IF InBody[ix*dx,iy*dx,(iz-1)*dx] AND NOT InBody(ix*dx,iy*dx,iz*dx) THEN
      h=Bisection[InBody(ix*dx,iy*dx,s),(iz-1)*dx,iz*dx]-baseline
      mh=~+h
      mh2=~+h^2
    END IF
    calcbcond(bcond(ix,iy,iz),0,(ix+0.5)*dx,iy*dx,iz*dx)
    calcbcond(bcond(ix,iy,iz),1,ix*dx,(iy+0.5)*dx,iz*dx)
    calcbcond(bcond(ix,iy,iz),2,ix*dx,iy*dx,(iz+0.5)*dx)
  END IF
REPEAT
bcond(0,*,*)=bcond(nx,*,*)
bcond(*,0,*)=bcond(*,ny,*)
mh=~/nx/ny
mh2=~/nx/ny
heff=mh2/mh; deff=SQRT(mh/heff*4/PI); ar=heff/deff
! deffa=SQRT(mh/a*4/PI); ara=a/deffa=(a/heff)^1.5*ar  !  va peggio

REAL r=0
LOOP FOR it=1 TO 50
  DO imbsmooth FOR 4 TIMES
  imbresids
  mg(field,rhs,bcond)
  IF it MOD 1=0 THEN
    ph=nz*dx-baseline-(SUM field(ix,iy,nz).u FOR ix=1 TO nx AND iy=1 TO ny)/nx/ny
    oldr=r; r=MAXABS(rhs)
    WRITE ar,ph/mh:1.7,ph,mh,ph/(1-2.1566*ph)/mh,a,heff,oldr/r
  END IF
  ! SPLOT rhs.p(*,ny DIV 2,*)
REPEAT
!(
uyw=CREATE("uyw.dat")
DO
  DO WRITE TO uyw (ix+0.5)*dx, iy*dx, [field(ix,iy,1).u-field(ix,iy,0).u]/dx FOR iy=0 TO ny
  WRITE TO uyw
FOR ix=0 TO nx-1
!)