# AMR-связка L0 (грубый) ↔ L1 (тонкий, измельчение r).
#
# Идея: на тонком патче считаем r подшагов по времени на каждый шаг L0. Границы патча («ghost»)
# заполняем интерполяцией с грубого уровня (по пространству + по времени между «до» и «после» шага L0),
# а внутренность патча после подшагов проецируем обратно на грубый уровень (рестрикция).
# Неравновесная часть масштабируется множителями Rcf (грубый→тонкий) и Rfc (тонкий→грубый),
# т.к. при смене τ масштаб неравновесия меняется.
import backend as B
import lattice as L
from equilibrium import feq, macros

cp = B.cp; xp = B.xp; DTYPE = B.DTYPE

def patch(pNx, pNy, ax, bx, ay, by, r):
    """Описание патча: размеры (Nfy,Nfx), индексы/веса билинейной интерполяции с грубого уровня,
    срезы для рестрикции каждого r-го узла."""
    Wx, Wy = bx - ax, by - ay
    Nfx, Nfy = r*Wx + 1, r*Wy + 1
    FX, FY = cp.meshgrid(ax + cp.arange(Nfx)/r, ay + cp.arange(Nfy)/r)   # координаты тонких узлов в L0
    x0 = cp.floor(FX).astype(cp.int64); y0 = cp.floor(FY).astype(cp.int64)
    x1 = cp.minimum(x0 + 1, pNx - 1); y1 = cp.minimum(y0 + 1, pNy - 1)
    return dict(Nfx=int(Nfx), Nfy=int(Nfy), ax=ax, bx=bx, ay=ay, by=by, r=r,
                x0=x0, y0=y0, x1=x1, y1=y1,
                tx=(FX - x0).astype(DTYPE), ty=(FY - y0).astype(DTYPE),
                slx=slice(r, r*Wx, r), sly=slice(r, r*Wy, r))

def pint(fld, P):
    """Билинейная интерполяция поля fld (грубого) в узлы тонкого патча."""
    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):
    """Граничные значения тонкого уровня из грубого поля pf: равновесие по интерполированным ρ,u
    плюс масштабированная неравновесная часть Rcf·neq.
    u в feq передаём кортежем (без cupy.stack) — H2D ломает CUDA-graph capture."""
    r, u = macros(pf); neq = pf - feq(r, u)
    return feq(pint(r, P), (pint(u[0], P), pint(u[1], P))) + Rcf*pint(neq, P)

def fill(cf, gh):
    """Записать ghost-значения gh в рамку (1 ряд) тонкого поля cf."""
    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):
    """Спроецировать тонкое поле cf обратно на грубое pf (каждый r-й узел), масштабируя neq на 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"]] = xp.where(fluid[None], nv, cur)
    return pf