"""
Benchmark 3 -- square cylinder (sharp-edged bluff body).

Why this case, given the circular cylinder is already covered
------------------------------------------------------------
On a circle the separation point is set by the boundary layer, so a solver
that under-resolves the wall gets separation -- and therefore drag -- wrong in
a way that depends on resolution. On a square the separation point is pinned
by the sharp leading corners regardless of resolution. That makes it a much
fairer test of an immersed-boundary method, and it is the closer analogue of
a road vehicle, where separation is likewise fixed by body-line breaks rather
than by a delicately balanced boundary layer.

It also exposes the specific weakness of a Cartesian immersed boundary: the
square's faces align with the grid, so this is the best case the method will
ever see. Agreement here and disagreement on the circle would point straight
at the staircase representation of curved surfaces.

Reference values:
  Re=100  Cd 1.44-1.51, St 0.144-0.150
  Re=200  Cd 1.37-1.48, St 0.148-0.156
"""

import sys, os, json
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np

from waterfall_core import Grid, Solver, external_flow
from waterfall_core import shapes
from cylinder import strouhal

Dh = 1.0
U = 1.0


def build(re, cells_per_d=24, lx_d=30.0, ly_d=20.0, x_c_d=8.0):
    nx, ny = int(lx_d * cells_per_d), int(ly_d * cells_per_d)
    grid = Grid(nx, ny, 1, lx_d * Dh, ly_d * Dh, 1.0)
    bc = external_flow(grid, u_inf=(U, 0.0, 0.0), sides="farfield")
    center = (x_c_d * Dh, 0.5 * ly_d * Dh)
    eu, ev, ew, ec = shapes.square_cylinder(grid, center, Dh, samples=4)
    solver = Solver(
        grid, bc, nu=U * Dh / re, rho=1.0,
        eta_u=eu, eta_v=ev, eta_w=ew, eta_c=ec, cfl=0.4,
    )
    solver.center = center
    return solver


def run_case(re, t_end, cells_per_d=24, verbose=True, **kw):
    s = build(re, cells_per_d=cells_per_d, **kw)
    if verbose:
        print(f"\n  Re={re}  {s.grid}  D/dx={Dh/s.grid.dx:.0f}")
    g = s.grid
    rng = np.random.default_rng(re)
    s.v[1:-1, 2:g.ny + 1, 1:-1] += 1e-3 * U * rng.standard_normal(
        (g.nx, g.ny - 1, g.nz)
    )
    elapsed = s.run(t_end=t_end, log_every=3000 if verbose else 0)

    area = Dh * g.dz
    coef = s.coefficients(area, u_ref=U, window=0.5)
    q = 0.5 * s.rho * U**2 * area
    f = s.filtered_force_history()
    cl = f[:, 1] / q
    st, power = strouhal(np.asarray(s.time_history)[1:], cl, d=Dh, u=U)
    half = len(cl) // 2
    cl_rms = float(np.sqrt(np.mean((cl[half:] - cl[half:].mean()) ** 2)))

    out = {
        "Re": re, "cells_per_D": cells_per_d,
        "cd": coef["cd"], "cd_std": coef["cd_std"],
        "cl_rms": cl_rms, "strouhal": st,
        "spectral_peak_fraction": power,
        "steps": s.step_count, "elapsed_s": elapsed,
        "divergence": s.divergence_norm(),
        "blockage_pct": 100.0 * Dh / g.ly,
    }
    if verbose:
        print(f"    Cd={out['cd']:.3f} +/- {out['cd_std']:.3f}   "
              f"Cl_rms={cl_rms:.3f}   St={st:.4f}   {elapsed:.0f}s")
    return s, out


REFERENCE = {
    100: {"cd": (1.44, 1.51), "st": (0.144, 0.150)},
    200: {"cd": (1.37, 1.48), "st": (0.148, 0.156)},
}


def main():
    print("Benchmark: square cylinder")
    results = []
    for re, t in [(100, 160.0), (200, 160.0)]:
        _s, out = run_case(re, t)
        ref = REFERENCE[re]
        for key, name in (("cd", "cd"), ("strouhal", "st")):
            lo, hi = ref[name]
            v = out[key]
            mid = 0.5 * (lo + hi)
            out[f"{name}_verdict"] = (
                "in range" if lo <= v <= hi else f"{100*(v-mid)/mid:+.1f}%"
            )
            out[f"{name}_err_pct"] = float(100 * (v - mid) / mid)
        out["reference"] = ref
        results.append(out)

    print()
    print(f"  {'Re':>5} {'Cd':>7} {'ref Cd':>12} {'verdict':>10} "
          f"{'St':>8} {'ref St':>14} {'verdict':>10}")
    for r in results:
        ref = REFERENCE[r["Re"]]
        # Ranges are formatted outside the f-string: nesting the same quote
        # character inside an f-string expression is a syntax error before
        # Python 3.12.
        cds = "{}-{}".format(*ref["cd"])
        sts = "{}-{}".format(*ref["st"])
        print(f"  {r['Re']:5d} {r['cd']:7.3f} "
              f"{cds:>12} {r['cd_verdict']:>10} "
              f"{r['strouhal']:8.4f} "
              f"{sts:>14} {r['st_verdict']:>10}")

    os.makedirs("validation", exist_ok=True)
    with open("validation/bluff_body.json", "w") as fh:
        json.dump({"case": "square_cylinder", "results": results}, fh, indent=2)
    print("\n  wrote validation/bluff_body.json")
    return results


if __name__ == "__main__":
    main()