//! A-P2, step S3 (`docs/overset_metal_campaign.md` §5.9): each half of the
//! overset coupling against the EXACT manufactured field, before Schwarz
//! joins them.
//!
//! (a) The patch with its outer row turned into ACCEPTORS stamped from the
//!     exact field every step (`u, v, p` at the acceptor centres, `p' = 0`)
//!     must keep the P0 annulus orders: Stokes ≥ 1.8, upwind ≈ 1.
//! (b) The background with the P2 hole and its FRINGE stamped from the exact
//!     field every step (velocity on the prescribed faces, pressure on the
//!     fringe cells, `p' = 0`) must land at the embedded-circle MMS level
//!     (L2 u 8.489e-3 / 4.341e-3 at n = 32 / 64, orders 0.92 / 0.97) and
//!     keep first order to n = 128.

use rtx_cfd::mesh::PatchMesh;
use rtx_cfd::mesh::patch_gen::annulus_skewed;
use rtx_cfd::solvers::incompressible::overset::overlap::DEFAULT_OVERLAP_ROWS;
use rtx_cfd::solvers::incompressible::{
    CurvilinearParameters, CurvilinearPisoSolver, EmbeddedParameters, EmbeddedPisoSolver, FaceKind,
    FlowField, NormalDiffusion, OverlapMap, PatchConvection, PatchField, PoissonSolverKind,
};
use rtx_cfd::{CfdConfig, CfdResult};
use std::f64::consts::PI;

const RHO: f64 = 1.0;
const MU: f64 = 0.05;

fn u_exact(x: f64, y: f64) -> f64 {
    (PI * x).sin() * (PI * y).cos()
}
fn v_exact(x: f64, y: f64) -> f64 {
    -(PI * x).cos() * (PI * y).sin()
}
fn p_exact(x: f64, y: f64) -> f64 {
    (PI * x).sin() * (PI * y).sin()
}
fn source(x: f64, y: f64, convecting: bool) -> (f64, f64) {
    let conv = if convecting { RHO * 0.5 * PI } else { 0.0 };
    (
        conv * (2.0 * PI * x).sin()
            + 2.0 * PI * PI * MU * u_exact(x, y)
            + PI * (PI * x).cos() * (PI * y).sin(),
        conv * (2.0 * PI * y).sin()
            + 2.0 * PI * PI * MU * v_exact(x, y)
            + PI * (PI * x).sin() * (PI * y).cos(),
    )
}

fn min_spacing(mesh: &PatchMesh) -> f64 {
    let mut h = f64::INFINITY;
    for c in 0..mesh.cell_count() {
        for (f, _) in mesh.cell_faces(c) {
            let d = mesh.faces()[f].d;
            h = h.min((d[0] * d[0] + d[1] * d[1]).sqrt());
        }
    }
    h
}

fn orders(errs: &[f64]) -> Vec<f64> {
    errs.windows(2).map(|p| (p[0] / p[1]).log2()).collect()
}

// ------------------------------------------------------------ (a) patch

/// March the P0 annulus with an acceptor ring stamped from the exact field
/// to steady state; L2 velocity error on the interior cells.
async fn patch_with_exact_acceptors(
    ns: usize,
    convection: PatchConvection,
    diffusion: NormalDiffusion,
) -> CfdResult<(f64, usize)> {
    let mesh = annulus_skewed([0.0, 0.0], 0.5, 1.5, ns, ns / 4, 0.3, 3.0)?;
    let nu = MU / RHO;
    let h = min_spacing(&mesh);
    let dt = 0.4 * (h * h / (4.0 * nu)).min(h);
    let config = CfdConfig::new()
        .with_density(RHO)
        .with_viscosity(MU)
        .with_reference_velocity(1.0)
        .with_reference_length(1.0);
    let params = CurvilinearParameters {
        tolerance: 1e-5,
        convection,
        normal_diffusion: diffusion,
        ..CurvilinearParameters::default()
    };
    let convecting = convection != PatchConvection::None;
    let mut solver = CurvilinearPisoSolver::new(config, params, mesh)?;
    solver.set_boundary_velocity(|x, y, _| (u_exact(x, y), v_exact(x, y)));
    solver.set_momentum_source(move |x, y, _| source(x, y, convecting));
    solver.set_acceptor_ring(true);
    let (ns, nn) = (solver.mesh().ns(), solver.mesh().nn());
    let acceptor_values: Vec<(f64, f64, f64)> = (0..ns)
        .map(|i| {
            let xy = solver.mesh().centre(solver.mesh().cell(nn - 1, i));
            (
                u_exact(xy[0], xy[1]),
                v_exact(xy[0], xy[1]),
                p_exact(xy[0], xy[1]),
            )
        })
        .collect();
    let zeros = vec![0.0; ns];
    let mut field = PatchField::new(solver.mesh());
    solver.initialize(&mut field, |_, _| (0.0, 0.0));
    solver.stamp_acceptors(&mut field, &acceptor_values);
    solver.set_acceptor_correction(&zeros);

    let steady_tol = if convecting { 1e-6 } else { 1e-7 };
    let mut steady = f64::INFINITY;
    let mut steps = 0;
    for step in 0..400_000 {
        let before = (field.u.clone(), field.v.clone());
        let r = solver.advance(&mut field, dt).await?;
        assert!(
            r.poisson_converged,
            "pressure solve did not converge at step {step}: {r:?}"
        );
        solver.stamp_acceptors(&mut field, &acceptor_values);
        steps = step + 1;
        let change = field
            .u
            .iter()
            .zip(&before.0)
            .chain(field.v.iter().zip(&before.1))
            .map(|(a, b)| (a - b).abs())
            .fold(0.0, f64::max);
        steady = change / dt;
        if steady < steady_tol {
            break;
        }
    }
    assert!(
        steady < steady_tol,
        "no steady state: |du/dt| = {steady:.3e}"
    );
    let mesh = solver.mesh();
    let (mut sq, mut vol) = (0.0, 0.0);
    for c in 0..mesh.cell_count() {
        if solver.is_acceptor(c) {
            continue;
        }
        let xy = mesh.centre(c);
        let eu = field.u[c] - u_exact(xy[0], xy[1]);
        let ev = field.v[c] - v_exact(xy[0], xy[1]);
        sq += (eu * eu + ev * ev) * mesh.area(c);
        vol += mesh.area(c);
    }
    Ok(((sq / vol).sqrt(), steps))
}

#[tokio::test]
async fn patch_with_exact_acceptors_keeps_the_p0_orders() -> CfdResult<()> {
    let only_upwind = std::env::var("RTX_OVERSET_UPWIND_ONLY").is_ok();
    for (convection, diffusion, gate) in [
        (PatchConvection::None, NormalDiffusion::Explicit, 1.8..2.6),
        (
            PatchConvection::None,
            NormalDiffusion::LineImplicit,
            1.8..2.6,
        ),
        (PatchConvection::Upwind, NormalDiffusion::Explicit, 0.7..1.6),
    ] {
        if only_upwind && convection != PatchConvection::Upwind {
            continue;
        }
        let mut errs = Vec::new();
        for ns in [32usize, 64, 128] {
            let (l2, steps) = patch_with_exact_acceptors(ns, convection, diffusion).await?;
            println!(
                "  acceptor patch {convection:?} {diffusion:?} ns={ns}: L2 {l2:.6e}, {steps} steps"
            );
            errs.push(l2);
        }
        let o = orders(&errs);
        println!("  acceptor patch {convection:?} {diffusion:?} orders {o:?}");
        assert!(
            o.iter().all(|x| gate.contains(x)),
            "{convection:?} orders {o:?} outside {gate:?}"
        );
    }
    Ok(())
}

// ------------------------------------------------------- (b) background

const CX: f64 = 0.6;
const CY: f64 = 0.45;
const R0: f64 = 0.2;
const R1: f64 = 0.354;

fn p2_patch(n: usize) -> CfdResult<PatchMesh> {
    annulus_skewed([CX, CY], R0, R1, 9 * n / 4, n / 4, 0.3, 3.0)
}

fn boundary_exact(x: f64, y: f64) -> (f64, f64) {
    let u = if x <= 0.0 || x >= 1.0 {
        0.0
    } else {
        u_exact(x, y)
    };
    let v = if y <= 0.0 || y >= 1.0 {
        0.0
    } else {
        v_exact(x, y)
    };
    (u, v)
}

/// Stamp the exact field onto the fringe faces and cells.
fn stamp_exact_fringe(map: &OverlapMap, field: &mut FlowField, h: f64) {
    for e in &map.fringe_u {
        field.u[(e.j, e.i)] = u_exact(e.i as f64 * h, (e.j as f64 + 0.5) * h);
    }
    for e in &map.fringe_v {
        field.v[(e.j, e.i)] = v_exact((e.i as f64 + 0.5) * h, e.j as f64 * h);
    }
    for e in &map.fringe_cells {
        field.p[(e.j, e.i)] = p_exact((e.i as f64 + 0.5) * h, (e.j as f64 + 0.5) * h);
    }
}

/// The embedded solver's stationarity floor grows with the grid (its
/// inner stop has an absolute part — the P0 finding on the patch): 1e-6
/// resolves n ≤ 64; at n = 128 |du/dt| floored at 2.2e-6 after 400k steps.
fn steady_tolerance(n: usize) -> f64 {
    1e-6 * (n as f64 / 64.0).powi(2).max(1.0)
}

async fn background_with_exact_fringe(n: usize) -> CfdResult<(f64, f64, usize)> {
    let steady_tol = steady_tolerance(n);
    let h = 1.0 / n as f64;
    let patch = p2_patch(n)?;
    let map = OverlapMap::build(&patch, n, n, h, h, DEFAULT_OVERLAP_ROWS)?;
    let nu = MU / RHO;
    let dt = 0.4 * (h * h / (4.0 * nu)).min(h);
    let config = CfdConfig::new()
        .with_density(RHO)
        .with_viscosity(MU)
        .with_reference_velocity(1.0)
        .with_reference_length(1.0);
    let mut solver = EmbeddedPisoSolver::new(
        config,
        EmbeddedParameters {
            corrector_steps: 2,
            tolerance: 1e-8,
            poisson_solver: PoissonSolverKind::Multigrid,
            ..EmbeddedParameters::default()
        },
    )?;
    solver.set_momentum_source(|x, y, _| source(x, y, true));
    solver.set_boundary_velocity(|x, y, _| boundary_exact(x, y));
    solver.set_overlap(map.background_mask(), map.fringe_flags());
    let mut field = FlowField::new(n, n, h, h)?;
    for j in 0..n {
        let y = (j as f64 + 0.5) * h;
        field.u[(j, 0)] = boundary_exact(0.0, y).0;
        field.u[(j, n)] = boundary_exact(1.0, y).0;
    }
    for i in 0..n {
        let x = (i as f64 + 0.5) * h;
        field.v[(0, i)] = boundary_exact(x, 0.0).1;
        field.v[(n, i)] = boundary_exact(x, 1.0).1;
    }
    stamp_exact_fringe(&map, &mut field, h);
    solver.initialize(&mut field)?;

    let mut steady = f64::INFINITY;
    let mut steps = 0;
    let mut max_div = 0.0_f64;
    for step in 0..400_000 {
        let before = (field.u.clone(), field.v.clone());
        let r = solver.advance(&mut field, dt).await?;
        // The embedded solver's `converged` flag asks for the normalised
        // residual below `tolerance` within its correctors, which the
        // reference harness (embedded_mms.rs) never asserts either; the
        // worst residual is reported instead.
        max_div = max_div.max(r.solver_result.final_residual);
        stamp_exact_fringe(&map, &mut field, h);
        steps = step + 1;
        let change = (&field.u - &before.0)
            .abs()
            .max()
            .max((&field.v - &before.1).abs().max());
        steady = change / dt;
        if steady < steady_tol {
            break;
        }
    }
    assert!(
        steady < steady_tol,
        "no steady state: |du/dt| = {steady:.3e}"
    );

    // L2 velocity on the fluid faces (the embedded_mms measure).
    let mask = solver.mask().expect("mask");
    let (mut sq, mut area) = (0.0, 0.0);
    for j in 0..n {
        for i in 1..n {
            if mask.u_kind(j, i) == FaceKind::Fluid {
                let e = field.u[(j, i)] - u_exact(i as f64 * h, (j as f64 + 0.5) * h);
                sq += e * e * h * h;
                area += h * h;
            }
        }
    }
    for j in 1..n {
        for i in 0..n {
            if mask.v_kind(j, i) == FaceKind::Fluid {
                let e = field.v[(j, i)] - v_exact((i as f64 + 0.5) * h, j as f64 * h);
                sq += e * e * h * h;
                area += h * h;
            }
        }
    }
    Ok(((sq / area).sqrt(), max_div, steps))
}

#[tokio::test]
async fn background_with_exact_fringe_lands_at_the_embedded_circle_level() -> CfdResult<()> {
    // embedded_mms.rs circle (0.6, 0.45) r = 0.2, upwind: 8.489e-3 / 4.341e-3
    // at n = 32 / 64. The hole here is larger (the patch's inner rows), so
    // the numbers are a level, not a pin.
    let reference = [8.489e-3, 4.341e-3];
    let mut errs = Vec::new();
    for (idx, n) in [32usize, 64, 128].into_iter().enumerate() {
        let (l2, max_div, steps) = background_with_exact_fringe(n).await?;
        println!(
            "  fringe background n={n}: L2 u {l2:.6e}, worst mass residual {max_div:.2e}, {steps} steps"
        );
        if let Some(r) = reference.get(idx) {
            assert!(
                l2 < 1.5 * r,
                "n = {n}: L2 {l2:.3e} > 1.5x the embedded circle's {r:.3e}"
            );
        }
        errs.push(l2);
    }
    let o = orders(&errs);
    println!("  fringe background orders {o:?}");
    // Measured 1.44 / 1.41: with the exact field on the fringe the hole
    // removes the body's near-wall layer, where upwind's error is largest,
    // and the remaining domain converges faster than the embedded circle's
    // 0.92 / 0.97 — an upper band of 1.6 (the P0 upwind band).
    assert!(o.iter().all(|&x| (0.75..1.6).contains(&x)), "orders {o:?}");
    Ok(())
}