"""First-gradient POST-G initialization counterfactual, not a qualification preset.

PRE F = feq - tau*A; POST G = feq - (tau - 1)*A. This low-Mach,
first-order Chapman-Enskog approximation is not an exact discrete eigenstate.
The physical viscous-stress factor tau - 1/2 is not the raw POST-G factor.
"""

import numpy as np

import shear_nested as nested
import shear_wave as sw
from oracle import CX, CY, W, U, centers

REFINEMENTS = nested.REFINEMENTS
CASES = ((100, 128, 1, "constant"), (100, 512, 1, "constant"),
         (100, 128, 3, "projected_linear"))
INITIALIZATION_ID = "postg_first_gradient_v1"
INITIALIZATION = (
    "physical FP64 POST G=feq-(tau-1)A at actual cell centers; "
    "A=3*w*U*cos(theta)*(c.e)*(c.k_level); rho=1; first-gradient low-Mach CE"
)


def initial(shape, level, length, viscosity):
    """Return physical FP64 G; production initialize owns shifted FP32 encoding."""
    x, y = centers(shape, level)
    k = sw.wavevector(length)
    theta = k[0] * x + k[1] * y + sw.PHI
    k_level = k * 2.0**-level
    tau = 0.5 + 3 * viscosity * 2**level
    ce = CX * sw.E[0] + CY * sw.E[1]
    ck = CX * k_level[0] + CY * k_level[1]
    correction = (3 * W * ce * ck)[:, None, None] * U * np.cos(theta)
    return sw.initial(shape, level, length, viscosity) - (tau - 1) * correction