# Серия по блокировке: D фиксирован, Ny растёт → β=D/Ny падает; экстраполяция β→0 даёт
# прямое сравнение с БЕЗГРАНИЧНОЙ литературой (Re=150: Cd≈1.33, St≈0.183, rms Cl≈0.3)
# без ручной поправки (1−β)². Запуск из этой папки:
#   python run_blockage.py
#   AMR_NY_LIST=90,128,180 python run_blockage.py    # свой список Ny
#   AMR_STEPS=20000 AMR_NY_LIST=128 python run_blockage.py   # смоук
# Стоимость прогона ≈ как у run.py: патч L1 (основная работа) от Ny не зависит (высота 4D фикс.).
import os, time
import numpy as np
import config as C
import backend as B
import solver as S
import diagnostics as Diag
import visualization as Viz

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "out"); os.makedirs(OUT, exist_ok=True)

LIT = dict(Cd=1.33, Clrms=0.3, St=0.183)   # безграничный цилиндр, Re≈150

def run_one(Ny):
    cfg = C.from_env(Ny=Ny, frame_every=0)             # серия без кадров/гифок (только ряды)
    print(f"\n[блокировка] Ny={cfg.Ny} (β={cfg.D/cfg.Ny:.3f}); домен {cfg.Nx}x{cfg.Ny}; "
          f"патч y=[{cfg.ay1},{cfg.by1}]; steps={cfg.steps}")
    t0 = time.perf_counter()
    sim = S.Solver(cfg)
    res = sim.run(Viz.make_progress(f"Ny={cfg.Ny}"))
    print("  время:", round(time.perf_counter() - t0, 1), "s")
    a = Diag.analyze(res, cfg)
    Diag.print_table(a, cfg)
    Diag.print_crosscheck(a)
    Diag.convergence_report(res, cfg)
    np.savez(os.path.join(OUT, f"probe_blockage_Ny{cfg.Ny}.npz"),
             St=a["St"], Cd=a["Cd"], Clrms=a["Clrms"], uyrms=a["uyrms"],
             Cm=a.get("Cm", float("nan")),
             Cd_st=a.get("Cd_st", float("nan")), Clrms_st=a.get("Clrms_st", float("nan")),
             Ny=cfg.Ny, beta=cfg.D/cfg.Ny, D=cfg.D, U=cfg.U, Re=cfg.Re, steps=cfg.steps,
             Cl=a["Cl_s"], uy=a["uy"], rho=B.to_cpu(res["rho"]),
             Fx=B.to_cpu(res["Fx"]), Fy=B.to_cpu(res["Fy"]), Tz=B.to_cpu(res["Tz"]))
    row = dict(Ny=cfg.Ny, beta=cfg.D/cfg.Ny, St=a["St"], Cd=a["Cd"], Clrms=a["Clrms"],
               Cm=a.get("Cm", float("nan")), Cd_st=a.get("Cd_st"))
    # освободить VRAM перед следующим прогоном (буферы, BC, захваченный CUDA-граф)
    del sim, res, a
    B.cp.get_default_memory_pool().free_all_blocks()
    return row

def save_extrapolation_figure(rows, fits, path):
    import matplotlib; matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    betas = np.array([r["beta"] for r in rows]); bb = np.linspace(0, betas.max()*1.05, 100)
    panels = [("Cd", "⟨Cd⟩", LIT["Cd"], "лит. 1.33"),
              ("Clrms", "rms Cl", LIT["Clrms"], "лит. ~0.3"),
              ("St", "St", LIT["St"], "лит. 0.183")]
    fig, axs = plt.subplots(1, 3, figsize=(13.5, 4.2), dpi=110)
    for ax, (key, lab, lit, litlab) in zip(axs, panels):
        vals = np.array([r[key] for r in rows]); ft = fits[key]
        ax.plot(betas, vals, "o", ms=7, color="#d32f2f", label="серия (raw)")
        al, sl = ft["lin"]; aq, sq = ft["quad"]
        ax.plot(bb, al + sl*bb, "-", lw=1.2, color="#1976d2", label=f"a+b·β → {al:.3f}")
        ax.plot(bb, aq + sq*bb**2, "--", lw=1.2, color="#388e3c", label=f"a+b·β² → {aq:.3f}")
        ax.axhline(lit, color="gray", lw=1.0, ls=":", label=litlab)
        if key == "St":                       # для St полезен и кинематически поправленный вид
            ax.plot(betas, vals*(1-betas), "s", ms=6, mfc="none", color="#7b1fa2", label="St·(1−β)")
        ax.set_xlabel("β = D/Ny"); ax.set_ylabel(lab); ax.set_xlim(0, bb[-1])
        ax.legend(fontsize=8); ax.set_title(f"{lab}(β): экстраполяция β→0", fontsize=10)
    fig.suptitle("Серия по блокировке: стеснение каналом — свойство ПОСТАНОВКИ, не считывания силы",
                 fontsize=11)
    fig.tight_layout(rect=(0, 0, 1, 0.93))
    fig.savefig(path, bbox_inches="tight"); plt.close(fig)
    print("saved", os.path.basename(path))

def main():
    ny_list = [int(s) for s in os.environ.get("AMR_NY_LIST", "90,128,180").split(",")]
    print(f"[серия по блокировке] backend={B.BACKEND}; Ny={ny_list}")
    rows = [run_one(Ny) for Ny in ny_list]
    if len(rows) < 2:
        print("\n(одна точка — фит/экстраполяция пропущены)")
        return
    betas = [r["beta"] for r in rows]
    fits = {key: Diag.fit_blockage(betas, [r[key] for r in rows]) for key in ("Cd", "Clrms", "St")}
    Diag.print_blockage_table(rows, fits)
    np.savez(os.path.join(OUT, "blockage_summary.npz"),
             Ny=np.array([r["Ny"] for r in rows]), beta=np.array(betas),
             St=np.array([r["St"] for r in rows]), Cd=np.array([r["Cd"] for r in rows]),
             Clrms=np.array([r["Clrms"] for r in rows]), Cm=np.array([r["Cm"] for r in rows]),
             Cd_st=np.array([r["Cd_st"] if r["Cd_st"] is not None else np.nan for r in rows]),
             Cd0_lin=fits["Cd"]["lin"][0], Cd0_quad=fits["Cd"]["quad"][0],
             Cl0_lin=fits["Clrms"]["lin"][0], Cl0_quad=fits["Clrms"]["quad"][0],
             St0_lin=fits["St"]["lin"][0], St0_quad=fits["St"]["quad"][0],
             lit_Cd=LIT["Cd"], lit_Clrms=LIT["Clrms"], lit_St=LIT["St"])
    save_extrapolation_figure(rows, fits, os.path.join(OUT, "blockage_extrapolation.png"))
    print("DONE (серия по блокировке)")

if __name__ == "__main__":
    main()