# AMR на цилиндре: полный домен 150x72, три версии (без AMR / 2x / вложенный 2x+4x).
# Встроенные зонды: сила на теле (обмен импульсом) -> Cd/Cl; u_y в следе -> Струхаль.
# Визуализация: inferno по |omega|, ТОЛЬКО рамки зон (без линий сетки), подписи осей.
# После прогона — анализ и сравнение трёх версий на реально собранных данных.
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)

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
def bb(pre,f,solid):
    for i in range(Q): f[i,solid]=pre[OPP[i],solid]
    return f
def cyl(NX,NY,ccx,ccy,rad,walls=True):
    Y,X=np.mgrid[0:NY,0:NX]; m=(X-ccx)**2+(Y-ccy)**2<=rad*rad
    if walls: m[0,:]=True; m[-1,:]=True
    return m
def force_on(fpost,solid_body):
    # обмен импульсом (Mei et al. 2002): по линкам жидкий-узел -> твёрдый сосед,
    # F = sum c_i (f_i + f_ī)^post — корректная передача импульса на стенку.
    fluid=~solid_body; Fx=0.0; Fy=0.0
    for i in range(1,Q):
        nb=np.roll(solid_body,(-int(CY[i]),-int(CX[i])),(0,1))  # nb[x]=solid[x+c_i]
        link=fluid & nb; s=(fpost[i]+fpost[OPP[i]])[link].sum()
        Fx+=CX[i]*s; Fy+=CY[i]*s
    return Fx,Fy

# --- полный домен 150x72 (вся сетка в кадре, измельчение внутри) ---
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  # зонд в следе
solid0=cyl(Nx,Ny,cx,cy,D/2)              # коарс: цилиндр + стенки
solid0_cyl=cyl(Nx,Ny,cx,cy,D/2,walls=False)
# уровень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
solid1=cyl(P1["Nfx"],P1["Nfy"],(cx-ax1)*2,(cy-ay1)*2,D,walls=False)         # D_L=32
fl1=~solid0[ay1+1:by1,ax1+1:bx1]
# уровень2 (4×, оранжевый): у тела (в коарс-координатах [26:64]x[18:54])
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
solid2=cyl(P2["Nfx"],P2["Nfy"],(cx-cx2a)*4,(cy-cy2a)*4,2*D,walls=False)      # D_L=64
fl2=np.ones((P2["by"]-P2["ay"]-1,P2["bx"]-P2["ax"]-1),bool)
onecol=np.ones((Ny,1))
print(f"домен {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")
    DL={"none":D0,"amr2x":2*D0,"nested":4*D0}[mode]
    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)
        if mode=="none": Fx[t],Fy[t]=force_on(post0,solid0_cyl)
        f0=stream(bb(pre,post0,solid0))
        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)
                if mode=="amr2x" and s1==1: Fx[t],Fy[t]=force_on(post1,solid1)
                f1=stream(bb(pre1,post1,solid1)); 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)
                        if s1==1 and s2==1: Fx[t],Fy[t]=force_on(post2,solid2)
                        f2=stream(bb(pre2,post2,solid2)); 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,DL=DL,Fx=Fx,Fy=Fy,uy=uy,frames=frames)

def save_gif(res,name,fps=24):
    nested=(res["mode"]=="nested"); frames=res["frames"]
    mid=frames[len(frames)//2][1]
    vm=np.nanpercentile(np.abs(np.where(solid0,np.nan,vort(mid))),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")
        o1=np.abs(np.where(solid1,np.nan,vort(u1)))[1:-1,1:-1]
        ax.imshow(o1,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:
            o2=np.abs(np.where(solid2,np.nan,vort(u2)))[1:-1,1:-1]
            ax.imshow(o2,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 на цилиндре: 2× (след) + 4× (у тела)" if nested else "AMR на цилиндре: одиночный блок 2×"
        ax.set_title(f"{ttl} · поверх $|\\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))
    path=os.path.join(ANIM,name)
    pil[0].save(path,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))
    fpk=freqs[1+int(np.argmax(amp[1:]))] if len(amp)>2 else 0.0
    St=fpk*D0/U
    return dict(Cd=float(Cd[h].mean()),Clrms=float(Cl[h].std()),St=float(St),
                uyrms=float(uy[h].std()),Cd_s=Cd,Cl_s=Cl,uy=uy,freqs=freqs,amp=amp,h0=n//2)

# ===== прогон трёх версий =====
print("=== none ===");   R0=run("none")
print("=== amr2x ===");  R1=run("amr2x")
print("=== nested ==="); R2=run("nested")
save_gif(R1,"amr2x_cyl.gif"); save_gif(R2,"amr_nested_cyl.gif")

A0,A1,A2=analyze(R0),analyze(R1),analyze(R2)
LABELS=["без AMR (D=16)","AMR 2× (D=32)","AMR 2×+4× (D=64)"]
AS=[A0,A1,A2]
print("\n=== СРАВНЕНИЕ НА СОБРАННЫХ ДАННЫХ (установившийся режим) ===")
print(f"{'версия':22}{'D у тела':>9}{'St':>8}{'<Cd>':>9}{'rms Cl':>9}{'rms u_y':>9}")
for lab,DL,a in zip(LABELS,[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}")
print(f"{'литература Re≈150':22}{'—':>9}{0.183:>8.3f}{1.33:>9.3f}{'~0.3':>9}")
print("(лит.: St≈0.18, Cd≈1.3 для цилиндра при Re≈150; Williamson 1996, Henderson 1995)")

# ===== фигура сравнения =====
np.savez(os.path.join(HERE,"_amr_probe_data.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]),
         Cd0=A0['Cd_s'],Cd1=A1['Cd_s'],Cd2=A2['Cd_s'],
         Cl0=A0['Cl_s'],Cl1=A1['Cl_s'],Cl2=A2['Cl_s'],
         uy0=R0['uy'],uy1=R1['uy'],uy2=R2['uy'])
cols=["#9e9e9e","#1f6feb","#d1495b"]
fig=plt.figure(figsize=(13.5,8.2))
gs=fig.add_gridspec(2,2,hspace=0.32,wspace=0.22)
# поле |omega| вложенного AMR
axA=fig.add_subplot(gs[0,0])
mid=R2["frames"][len(R2["frames"])//2]; vmf=np.nanpercentile(np.abs(np.where(solid0,np.nan,vort(mid[1]))),99)
cmf=plt.get_cmap("inferno").copy(); cmf.set_bad("white")
axA.imshow(np.abs(np.where(solid0,np.nan,vort(mid[1]))),origin="lower",cmap=cmf,vmin=0,vmax=vmf,
           extent=(0,Nx,0,Ny),interpolation="nearest")
axA.add_patch(mpatches.Rectangle((ax1,ay1),bx1-ax1,by1-ay1,fill=False,edgecolor="#aed581",lw=1.8))
axA.add_patch(mpatches.Rectangle((cx2a,cy2a),cx2b-cx2a,cy2b-cy2a,fill=False,edgecolor="#ff8a65",lw=1.8))
axA.set_title("Вложенный AMR: $|\\omega|$ + рамки зон 2×/4×"); axA.set_xlabel("x [lu]"); axA.set_ylabel("y [lu]")
# Cl(t)
axB=fig.add_subplot(gs[0,1])
for a,lab,c in zip(AS,LABELS,cols):
    axB.plot(a['Cl_s'][a['h0']:],lw=0.8,color=c,label=lab)
axB.set_title("Подъёмная сила $C_l(t)$ (зонд силы, обмен импульсом)")
axB.set_xlabel("шаг (установившийся режим)"); axB.set_ylabel("$C_l$"); axB.legend(fontsize=8); axB.grid(alpha=0.3)
# спектры u_y
axC=fig.add_subplot(gs[1,0])
for a,lab,c in zip(AS,LABELS,cols):
    sp=a['amp']/a['amp'][1:].max(); axC.plot(a['freqs']*D0/U,sp,lw=1.0,color=c,label=lab)
axC.axvline(0.183,color="k",ls="--",alpha=0.6,label="лит. St≈0.18")
axC.set_xlim(0,0.6); axC.set_title("Спектр $u_y$ в следе → число Струхаля")
axC.set_xlabel("St = f·D/U"); axC.set_ylabel("нормир. амплитуда"); axC.legend(fontsize=8); axC.grid(alpha=0.3)
# столбики St и Cd
axD=fig.add_subplot(gs[1,1]); x=np.arange(3); w=0.35
axD.bar(x-w/2,[a['St'] for a in AS],w,color="#1f6feb",label="St")
axD.bar(x+w/2,[a['Cd'] for a in AS],w,color="#fb8c00",label="<Cd>")
axD.axhline(0.183,color="#1f6feb",ls="--",alpha=0.6); axD.axhline(1.33,color="#fb8c00",ls="--",alpha=0.6)
axD.set_xticks(x); axD.set_xticklabels(["без AMR","2×","2×+4×"],fontsize=9)
axD.set_title("St и <Cd> vs разрешение у тела (пунктир — лит.)"); axD.legend(fontsize=8); axD.grid(alpha=0.3,axis="y")
for i,a in enumerate(AS):
    axD.text(i-w/2,a['St']+0.005,f"{a['St']:.3f}",ha="center",fontsize=7)
    axD.text(i+w/2,a['Cd']+0.02,f"{a['Cd']:.2f}",ha="center",fontsize=7)
fig.suptitle("AMR на цилиндре: сравнение трёх версий на реально собранных данных зондов",fontweight="bold")
fig.savefig(os.path.join(FIGS,"11e_amr_probe_comparison.png"),dpi=110,bbox_inches="tight"); plt.close(fig)
print("saved figures/11e_amr_probe_comparison.png")
print("DONE")