# СЛЕДУЮЩАЯ ИТЕРАЦИЯ решателя: тот же AMR (без / 2× / вложенный 2×+4×), но граница тела
# задана ПОЛНОЦЕННЫМ SDF (знаковое поле расстояний) + интерполированное отражение Bouzidi
# (Bouzidi-Firdaouss-Lallemand 2001) — сглаживает ступенчатость криволинейной границы.
# Геометрия/Re/шаги ИДЕНТИЧНЫ _amr_anims.py — отличие ТОЛЬКО в граничном условии, поэтому
# результаты прямо сравнимы. В конце — сводное сравнение с версиями без SDF.
import sys, os, io as _io, numpy as np, matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt, matplotlib.patches as mpatches
from PIL import Image
try: sys.stdout.reconfigure(encoding="utf-8")
except Exception: pass
HERE = os.path.dirname(os.path.abspath(__file__))
ANIM = os.path.join(HERE, "anim"); os.makedirs(ANIM, exist_ok=True)
FIGS = os.path.join(HERE, "figures"); os.makedirs(FIGS, exist_ok=True)

# ===== проверенное ядро D2Q9 KBC (идентично _amr_anims.py) =====
C=np.array([[0,0],[1,0],[0,1],[-1,0],[0,-1],[1,1],[-1,1],[-1,-1],[1,-1]],float)
W=np.array([4/9,1/9,1/9,1/9,1/9,1/36,1/36,1/36,1/36]); CS2=1/3; Q=9
CX,CY=C[:,0],C[:,1]; OPP=np.array([0,3,4,1,2,7,8,5,6])
M=np.zeros((Q,Q)); M[0]=1;M[1]=CX;M[2]=CY;M[3]=3*(CX**2+CY**2)-2
M[4]=CX**2-CY**2;M[5]=CX*CY;M[6]=CX**2*CY;M[7]=CX*CY**2;M[8]=CX**2*CY**2
Dm=np.zeros((Q,Q)); Dm[4,4]=Dm[5,5]=1.0; Ps=np.linalg.inv(M)@Dm@M
def feq(rho,u):
    ux=np.clip(u[0],-0.95,0.95); uy=np.clip(u[1],-0.95,0.95)
    sx=np.sqrt(1+3*ux**2); sy=np.sqrt(1+3*uy**2); base=rho*(2-sx)*(2-sy)
    Bx=((2*ux+sx)/(1-ux))[None]**CX[:,None,None]; By=((2*uy+sy)/(1-uy))[None]**CY[:,None,None]
    return W[:,None,None]*base[None]*Bx*By
def macros(f):
    r=f.sum(0); return r,np.stack([(CX[:,None,None]*f).sum(0)/r,(CY[:,None,None]*f).sum(0)/r])
def collide(f,fe,beta):
    dfn=f-fe; ds=np.einsum("ij,jab->iab",Ps,dfn); dh=dfn-ds; inv=1.0/fe
    num=(ds*dh*inv).sum(0); den=(dh*dh*inv).sum(0)
    with np.errstate(divide="ignore",invalid="ignore"):
        g=np.where(den>1e-14,1/beta-(2-1/beta)*num/den,2.0)
    return f-beta*(2*ds+g[None]*dh)
def stream(f):
    fs=np.empty_like(f)
    for i in range(Q): fs[i]=np.roll(f[i],(int(CY[i]),int(CX[i])),(0,1))
    return fs
def smoothstep(x): x=min(max(x,0.0),1.0); return x*x*(3-2*x)
def tauc(tp,r): return r*tp-(r-1)/2
def vort(u):
    return (np.roll(u[1],-1,1)-np.roll(u[1],1,1))*0.5-(np.roll(u[0],-1,0)-np.roll(u[0],1,0))*0.5
def patch(pNx,pNy,ax,bx,ay,by,r):
    Wx,Wy=bx-ax,by-ay; Nfx,Nfy=r*Wx+1,r*Wy+1
    FX,FY=np.meshgrid(ax+np.arange(Nfx)/r, ay+np.arange(Nfy)/r)
    x0=np.floor(FX).astype(int); y0=np.floor(FY).astype(int)
    x1=np.minimum(x0+1,pNx-1); y1=np.minimum(y0+1,pNy-1)
    return dict(Nfx=Nfx,Nfy=Nfy,ax=ax,bx=bx,ay=ay,by=by,r=r,x0=x0,y0=y0,x1=x1,y1=y1,
                tx=FX-x0,ty=FY-y0,slx=slice(r,r*Wx,r),sly=slice(r,r*Wy,r))
def pint(fld,P):
    return (fld[...,P["y0"],P["x0"]]*(1-P["tx"])*(1-P["ty"])+fld[...,P["y0"],P["x1"]]*P["tx"]*(1-P["ty"])
            +fld[...,P["y1"],P["x0"]]*(1-P["tx"])*P["ty"]+fld[...,P["y1"],P["x1"]]*P["tx"]*P["ty"])
def ghost(pf,P,Rcf):
    r,u=macros(pf); neq=pf-feq(r,u)
    return feq(pint(r,P),np.stack([pint(u[0],P),pint(u[1],P)]))+Rcf*pint(neq,P)
def fill(cf,gh):
    cf[:,0,:]=gh[:,0,:];cf[:,-1,:]=gh[:,-1,:];cf[:,:,0]=gh[:,:,0];cf[:,:,-1]=gh[:,:,-1]; return cf
def restrict(cf,pf,P,Rfc,fluid):
    r,u=macros(cf); neq=cf-feq(r,u); slx,sly=P["slx"],P["sly"]
    nv=feq(r[sly,slx],u[:,sly,slx])+Rfc*neq[:,sly,slx]
    cur=pf[:,P["ay"]+1:P["by"],P["ax"]+1:P["bx"]]
    pf[:,P["ay"]+1:P["by"],P["ax"]+1:P["bx"]]=np.where(fluid[None],nv,cur); return pf

# ===== ПОЛНОЦЕННЫЙ SDF + интерполированное отражение Bouzidi =====
def sdf_circle(NX,NY,ccx,ccy,R):
    # знаковое расстояние до окружности: >0 жидкость, <0 внутри тела
    Y,X=np.mgrid[0:NY,0:NX]; return np.sqrt((X-ccx)**2+(Y-ccy)**2)-R
def sdf_from_mask(mask):
    # общий SDF из произвольной бинарной маски (для не-круглых тел), через EDT
    from scipy import ndimage
    din=ndimage.distance_transform_edt(mask); dout=ndimage.distance_transform_edt(~mask)
    return dout-din-0.5   # >0 жидкость, <0 тело; -0.5 — поверхность между узлами
def add_walls(phi):
    NY=phi.shape[0]; Y=np.arange(NY)[:,None]*np.ones((1,phi.shape[1]))
    return np.minimum(phi, np.minimum(Y-0.5, (NY-1.5)-Y))  # стенки: поверхность на y=0.5 и Ny-1.5
def build_bc(solid, phi, cyl_mask):
    # для каждого направления: линки жидкий->твёрдый, доля стенки q (по SDF),
    # есть ли второй жидкий узел (для интерполяции), и принадлежит ли сосед телу (для силы)
    fluid=~solid; bc=[]
    for i in range(1,Q):
        cy_,cx_=int(CY[i]),int(CX[i])
        nb_solid=np.roll(solid,(-cy_,-cx_),(0,1))            # solid[x+c_i]
        mask=fluid & nb_solid
        cl=mask & np.roll(cyl_mask,(-cy_,-cx_),(0,1))        # сосед — тело (не стенка)
        phinb=np.roll(phi,(-cy_,-cx_),(0,1))
        with np.errstate(divide="ignore",invalid="ignore"):
            qq=phi/(phi-phinb)                               # доля линка до стенки (phi=0)
        q=np.where(mask, np.clip(qq,0.02,0.98), 0.5)
        xff=mask & (~np.roll(solid,(cy_,cx_),(0,1)))         # x_f - c_i — жидкий?
        bc.append((i,int(OPP[i]),mask,q,xff,cl))
    return bc
def apply_bc(f, fpost, bc):
    # Bouzidi (линейный): пристеночный жидкий узел получает корректно отражённую популяцию
    for (i,ib,mask,q,xff,cl) in bc:
        cy_,cx_=int(CY[i]),int(CX[i])
        fi=fpost[i]; fib=fpost[ib]; fiback=np.roll(fpost[i],(cy_,cx_),(0,1))  # f_i в x_f-c_i
        near=mask&(q<0.5)&xff; bad=mask&(q<0.5)&(~xff); far=mask&(q>=0.5)
        qn=q[near]; f[ib][near]=2*qn*fi[near]+(1-2*qn)*fiback[near]
        f[ib][bad]=fi[bad]                                   # запас: нет второго узла -> halfway
        qf=q[far];  f[ib][far]=(1/(2*qf))*fi[far]+(1-1/(2*qf))*fib[far]
    return f
def force_bc(fpost, f, bc):
    # обмен импульсом (Mei et al. 2002) для интерполированной границы: F=sum c_i (f_i^post + f_ī^bc)
    Fx=0.0; Fy=0.0
    for (i,ib,mask,q,xff,cl) in bc:
        s=(fpost[i]+f[ib])[cl].sum(); Fx+=CX[i]*s; Fy+=CY[i]*s
    return Fx,Fy

# ===== геометрия (идентична версии без SDF) =====
Nx,Ny=150,72; D=16; cx,cy=40,36; U=0.07; Re=150.0; STEPS=8000; ramp=1000; FE=34
D0=D; nu=U*D/Re; tau0=nu/CS2+0.5; b0=1/(2*tau0); px,py=cx+3*D,cy
# коарс: цилиндр (SDF) + стенки
phi0=add_walls(sdf_circle(Nx,Ny,cx,cy,D/2)); solid0=phi0<0
cyl0=sdf_circle(Nx,Ny,cx,cy,D/2)<0
bc0=build_bc(solid0,phi0,cyl0)
# уровень1 (2×): тело+след
ax1,bx1,ay1,by1=16,118,8,64; P1=patch(Nx,Ny,ax1,bx1,ay1,by1,2)
t1=tauc(tau0,2); b1=1/(2*t1); R01=t1/(2*tau0); Rf01=1/R01
phi1=sdf_circle(P1["Nfx"],P1["Nfy"],(cx-ax1)*2,(cy-ay1)*2,D); solid1=phi1<0
bc1=build_bc(solid1,phi1,solid1); fl1=~solid0[ay1+1:by1,ax1+1:bx1]
# уровень2 (4×): у тела
cx2a,cx2b,cy2a,cy2b=26,64,18,54
P2=patch(P1["Nfx"],P1["Nfy"],(cx2a-ax1)*2,(cx2b-ax1)*2,(cy2a-ay1)*2,(cy2b-ay1)*2,2)
t2=tauc(t1,2); b2=1/(2*t2); R12=t2/(2*t1); Rf12=1/R12
phi2=sdf_circle(P2["Nfx"],P2["Nfy"],(cx-cx2a)*4,(cy-cy2a)*4,2*D); solid2=phi2<0
bc2=build_bc(solid2,phi2,solid2); fl2=np.ones((P2["by"]-P2["ay"]-1,P2["bx"]-P2["ax"]-1),bool)
onecol=np.ones((Ny,1))
print(f"[SDF] домен {Nx}x{Ny}; L1(2×) {P1['Nfx']}x{P1['Nfy']}; L2(4×) {P2['Nfx']}x{P2['Nfy']}; "
      f"Re={Re} D0={D0} steps={STEPS}")

def run(mode):  # "none" | "amr2x" | "nested"
    use1=mode in ("amr2x","nested"); use2=(mode=="nested")
    f0=feq(np.ones((Ny,Nx)),np.zeros((2,Ny,Nx)))
    f1=feq(np.ones((P1["Nfy"],P1["Nfx"])),np.zeros((2,P1["Nfy"],P1["Nfx"]))) if use1 else None
    f2=feq(np.ones((P2["Nfy"],P2["Nfx"])),np.zeros((2,P2["Nfy"],P2["Nfx"]))) if use2 else None
    Fx=np.zeros(STEPS+1); Fy=np.zeros(STEPS+1); uy=np.zeros(STEPS+1); frames=[]
    for t in range(STEPS+1):
        rho,u0=macros(f0)
        if not np.all(np.isfinite(u0)): print("BLEW UP",mode,t); Fx=Fx[:t];Fy=Fy[:t];uy=uy[:t]; break
        pre=f0.copy(); post0=collide(f0,feq(rho,u0),b0); f0=stream(post0); f0=apply_bc(f0,post0,bc0)
        if mode=="none": Fx[t],Fy[t]=force_bc(post0,f0,bc0)
        uin=np.zeros((2,Ny,1)); uin[0,:,0]=U*smoothstep(t/ramp)
        f0[:,1:-1,0]=feq(onecol,uin)[:,1:-1,0]; f0[:,1:-1,-1]=f0[:,1:-1,-2]
        if use1:
            g1o=ghost(pre,P1,R01); g1n=ghost(f0,P1,R01)
            for s1 in range(2):
                pre1=f1.copy(); r1,u1=macros(f1); post1=collide(f1,feq(r1,u1),b1)
                f1=stream(post1); f1=apply_bc(f1,post1,bc1)
                if mode=="amr2x" and s1==1: Fx[t],Fy[t]=force_bc(post1,f1,bc1)
                f1=fill(f1,(1-(s1+1)/2)*g1o+((s1+1)/2)*g1n)
                if use2:
                    g2o=ghost(pre1,P2,R12); g2n=ghost(f1,P2,R12)
                    for s2 in range(2):
                        pre2=f2.copy(); r2,u2=macros(f2); post2=collide(f2,feq(r2,u2),b2)
                        f2=stream(post2); f2=apply_bc(f2,post2,bc2)
                        if s1==1 and s2==1: Fx[t],Fy[t]=force_bc(post2,f2,bc2)
                        f2=fill(f2,(1-(s2+1)/2)*g2o+((s2+1)/2)*g2n)
                    f1=restrict(f2,f1,P2,Rf12,fl2)
            f0=restrict(f1,f0,P1,Rf01,fl1)
        uc=macros(f0)[1]; uy[t]=uc[1,py,px]
        if use1 and t%FE==0:
            frames.append((t,uc.copy(),macros(f1)[1].copy(),(macros(f2)[1].copy() if use2 else None)))
    return dict(mode=mode,Fx=Fx,Fy=Fy,uy=uy,frames=frames,DL={"none":D0,"amr2x":2*D0,"nested":4*D0}[mode])

def save_gif(res,name,fps=24):
    nested=(res["mode"]=="nested"); frames=res["frames"]
    vm=np.nanpercentile(np.abs(np.where(solid0,np.nan,vort(frames[len(frames)//2][1]))),99.0)
    cm=plt.get_cmap("inferno").copy(); cm.set_bad("white")
    pil=[]
    for (t,u0,u1,u2) in frames:
        fig=plt.figure(figsize=(11.2,5.7),dpi=96); ax=fig.add_subplot(111)
        ax.imshow(np.abs(np.where(solid0,np.nan,vort(u0))),origin="lower",cmap=cm,vmin=0,vmax=vm,
                  extent=(0,Nx,0,Ny),interpolation="nearest")
        ax.imshow(np.abs(np.where(solid1,np.nan,vort(u1)))[1:-1,1:-1],origin="lower",cmap=cm,vmin=0,vmax=vm,
                  extent=(ax1+0.5,bx1-0.5,ay1+0.5,by1-0.5),interpolation="nearest")
        if nested and u2 is not None:
            ax.imshow(np.abs(np.where(solid2,np.nan,vort(u2)))[1:-1,1:-1],origin="lower",cmap=cm,vmin=0,vmax=vm,
                      extent=(cx2a+0.25,cx2b-0.25,cy2a+0.25,cy2b-0.25),interpolation="nearest")
        ax.add_patch(mpatches.Rectangle((0.2,0.2),Nx-0.4,Ny-0.4,fill=False,edgecolor="#4fc3f7",lw=1.6))
        ax.text(1.5,Ny-4,"уровень 0: Δx (крупный)",color="#4fc3f7",fontsize=9,weight="bold")
        ax.add_patch(mpatches.Rectangle((ax1,ay1),bx1-ax1,by1-ay1,fill=False,edgecolor="#aed581",lw=2.2))
        ax.text(ax1+1,by1-3.5,"уровень 1: Δx/2 (тело+след)",color="#aed581",fontsize=9,weight="bold")
        if nested:
            ax.add_patch(mpatches.Rectangle((cx2a,cy2a),cx2b-cx2a,cy2b-cy2a,fill=False,edgecolor="#ff8a65",lw=2.2))
            ax.text(cx2a+1,cy2b-3.5,"уровень 2: Δx/4 (у тела)",color="#ff8a65",fontsize=9,weight="bold")
        ttl="Вложенный AMR + SDF: 2×(след)+4×(тело)" if nested else "AMR + SDF: одиночный блок 2×"
        ax.set_title(f"{ttl} · граница тела по SDF (Bouzidi) · поверх $|\\omega|$ · t={t}",fontsize=11)
        ax.set_xlabel("x [lu]"); ax.set_ylabel("y [lu]"); ax.set_xlim(0,Nx); ax.set_ylim(0,Ny)
        buf=_io.BytesIO(); fig.savefig(buf,format="png",bbox_inches="tight"); buf.seek(0); plt.close(fig)
        pil.append(Image.open(buf).convert("RGB").convert("P",palette=Image.ADAPTIVE))
    pil[0].save(os.path.join(ANIM,name),save_all=True,append_images=pil[1:],duration=int(1000/fps),loop=0,optimize=True)
    print("saved",name,len(frames),"кадров")

def analyze(res):
    Fx,Fy,uy,DL=res["Fx"],res["Fy"],res["uy"],res["DL"]
    n=len(uy); h=slice(n//2,n); Cd=2*Fx/(U**2*DL); Cl=2*Fy/(U**2*DL)
    sig=uy[h]-uy[h].mean(); npad=8192
    freqs=np.fft.rfftfreq(npad,1.0); amp=np.abs(np.fft.rfft(sig,n=npad))
    St=(freqs[1+int(np.argmax(amp[1:]))] if len(amp)>2 else 0.0)*D0/U
    return dict(Cd=float(Cd[h].mean()),Clrms=float(Cl[h].std()),St=float(St),uyrms=float(uy[h].std()),
                Cl_s=Cl,uy=uy,freqs=freqs,amp=amp,h0=n//2)

print("=== SDF none ===");   R0=run("none")
print("=== SDF amr2x ===");  R1=run("amr2x")
print("=== SDF nested ==="); R2=run("nested")
save_gif(R1,"amr2x_cyl_sdf.gif"); save_gif(R2,"amr_nested_cyl_sdf.gif")
AS=[analyze(R0),analyze(R1),analyze(R2)]; LAB=["без AMR (D=16)","AMR 2× (D=32)","AMR 2×+4× (D=64)"]
print("\n=== SDF: СРАВНЕНИЕ НА СОБРАННЫХ ДАННЫХ ===")
print(f"{'версия':22}{'D у тела':>9}{'St':>8}{'<Cd>':>9}{'rms Cl':>9}{'rms u_y':>9}")
for lab,DL,a in zip(LAB,[16,32,64],AS):
    print(f"{lab:22}{DL:>9}{a['St']:>8.3f}{a['Cd']:>9.3f}{a['Clrms']:>9.4f}{a['uyrms']:>9.4f}")
np.savez(os.path.join(HERE,"_amr_probe_data_sdf.npz"),
         St=np.array([a['St'] for a in AS]),Cd=np.array([a['Cd'] for a in AS]),
         Clrms=np.array([a['Clrms'] for a in AS]),uyrms=np.array([a['uyrms'] for a in AS]))

# ===== сводное сравнение SDF vs без SDF (если есть данные без SDF) =====
nosdf_path=os.path.join(HERE,"_amr_probe_data.npz")
if os.path.exists(nosdf_path) and ("St" in np.load(nosdf_path).files):
    N=np.load(nosdf_path); St_n,Cd_n,Cl_n=N["St"],N["Cd"],N["Clrms"]
    St_s=np.array([a['St'] for a in AS]); Cd_s=np.array([a['Cd'] for a in AS]); Cl_s=np.array([a['Clrms'] for a in AS])
    print("\n=== СВОДНО: SDF vs без SDF ===")
    print(f"{'версия':14}{'St(noSDF)':>11}{'St(SDF)':>9}{'Cd(noSDF)':>11}{'Cd(SDF)':>9}")
    for k,lab in enumerate(["без AMR","2×","2×+4×"]):
        print(f"{lab:14}{St_n[k]:>11.3f}{St_s[k]:>9.3f}{Cd_n[k]:>11.3f}{Cd_s[k]:>9.3f}")
    print(f"{'лит. Re≈150':14}{0.183:>11.3f}{0.183:>9.3f}{1.330:>11.3f}{1.330:>9.3f}")
    x=np.arange(3); w=0.38
    fig,axs=plt.subplots(1,2,figsize=(13.5,4.6))
    axs[0].bar(x-w/2,Cd_n,w,color="#9e9e9e",label="без SDF (ступенч.)")
    axs[0].bar(x+w/2,Cd_s,w,color="#1f6feb",label="с SDF (Bouzidi)")
    axs[0].axhline(1.33,color="k",ls="--",alpha=0.6,label="лит. Cd≈1.33")
    axs[0].set_xticks(x); axs[0].set_xticklabels(["без AMR","2×","2×+4×"]); axs[0].set_ylabel("<Cd>")
    axs[0].set_title("Сопротивление: SDF убирает завышение от ступенчатой границы"); axs[0].legend(fontsize=8); axs[0].grid(alpha=0.3,axis="y")
    axs[1].bar(x-w/2,St_n,w,color="#9e9e9e",label="без SDF")
    axs[1].bar(x+w/2,St_s,w,color="#1f6feb",label="с SDF")
    axs[1].axhline(0.183,color="k",ls="--",alpha=0.6,label="лит. St≈0.18")
    axs[1].set_xticks(x); axs[1].set_xticklabels(["без AMR","2×","2×+4×"]); axs[1].set_ylabel("St")
    axs[1].set_title("Число Струхаля"); axs[1].legend(fontsize=8); axs[1].grid(alpha=0.3,axis="y")
    fig.suptitle("Цилиндр: SDF (Bouzidi) против ступенчатого bounce-back, на реально собранных данных",fontweight="bold")
    fig.savefig(os.path.join(FIGS,"11f_sdf_vs_nosdf.png"),dpi=110,bbox_inches="tight"); plt.close(fig)
    print("saved figures/11f_sdf_vs_nosdf.png")
else:
    print("\n(нет свежего _amr_probe_data.npz со скалярами — сначала запусти _amr_anims.py)")
print("DONE")