//! P4 option B (`docs/overset_metal_campaign.md` §5.11): the momentum
//! residual of the background's OWN staggered predictor stencil on every
//! background face, at a settled state.
//!
//! On a solved face the discrete equation the composite marched is
//! `ρ (u^{n+1} − u^n)/dt = ρ · rhs(u^n, p^{n+1})` (predictor plus the
//! correctors' `−dt ∇p'/ρ`, with `p^{n+1} = p^n + Σ p'`), so the residual
//! `r = ρ [(u^{n+1} − u^n)/dt − rhs] · dx dy` is zero to rounding there
//! — the pin that proves the diagnostic IS the solver's operator. On a
//! PRESCRIBED face the value is stamped from the patch, the equation is
//! not solved, and `r` is the momentum source the stamping injects, in the
//! solver's own metric and without the staircase curves' face-formula
//! error. Summed over the ring it is the fringe ring's momentum defect
//! (`region_force` ring outer − hole boundary, but exact).
//!
//! Validity is decided by the stencil itself: the background field is
//! copied with `NaN` on every value that is neither the solver's own nor
//! stamped from the patch (hole cells and hole–hole faces within two cells
//! of the ring get the patch's interpolated values, `OverlapMap::hole_p`
//! / `ghost_u` / `ghost_v`, from a band widened three rows into the
//! hole), the
//! operator is evaluated as is, and a `NaN` result means the face read
//! something invalid and is not counted. Under the upwind scheme every
//! value the stencil reads enters its arithmetic, so the test is exact.

use std::collections::HashSet;

use super::overlap::CellClass;
use super::{OversetField, OversetPisoSolver};
use crate::solvers::incompressible::embedded_body::FaceKind;

/// Sums over one class of faces.
#[derive(Debug, Clone, Copy, Default)]
pub struct ResidualBucket {
    /// `Σ r` on the u faces (x-momentum source, N/m).
    pub fx: f64,
    /// `Σ r` on the v faces.
    pub fy: f64,
    /// `Σ |r|` on the u faces.
    pub abs_x: f64,
    /// `Σ |r|` on the v faces.
    pub abs_y: f64,
    /// Largest `|r|` on the u faces.
    pub max_abs_x: f64,
    /// Largest `|r|` on the v faces.
    pub max_abs_y: f64,
    /// Faces whose stencil read only valid values.
    pub evaluated: usize,
    /// Of `evaluated`, the u faces.
    pub evaluated_u: usize,
    /// Of `evaluated`, the v faces.
    pub evaluated_v: usize,
    /// Faces of this class.
    pub total: usize,
}

impl ResidualBucket {
    fn add(&mut self, r: f64, is_u: bool) {
        self.total += 1;
        if !r.is_finite() {
            return;
        }
        self.evaluated += 1;
        if is_u {
            self.evaluated_u += 1;
            self.fx += r;
            self.abs_x += r.abs();
            self.max_abs_x = self.max_abs_x.max(r.abs());
        } else {
            self.evaluated_v += 1;
            self.fy += r;
            self.abs_y += r.abs();
            self.max_abs_y = self.max_abs_y.max(r.abs());
        }
    }
}

/// One prescribed face's residual.
#[derive(Debug, Clone, Copy)]
pub struct FaceResidual {
    /// A u face (x-momentum) or a v face.
    pub is_u: bool,
    /// Row.
    pub j: usize,
    /// Column.
    pub i: usize,
    /// The residual (N/m), `NaN` when not evaluable.
    pub r: f64,
    /// Between two fringe cells (else fringe–hole).
    pub fringe_fringe: bool,
    /// The pieces of `r` (N/m): unsteady, convective, diffusive, pressure
    /// (`r = time + conv − diff + pres` under upwind on interior faces;
    /// under a limited scheme the convective piece is the upwind part
    /// only and the four do not reconstruct `r`).
    pub pieces: [f64; 4],
}

/// The momentum residual by face class.
#[derive(Debug, Clone, Default)]
pub struct MomentumResidual {
    /// Solved faces whose stencil reads only solved values.
    pub solved_far: ResidualBucket,
    /// Solved faces whose stencil reads a fringe cell or a prescribed face.
    pub solved_near: ResidualBucket,
    /// Prescribed faces between two fringe cells (tangential to the ring).
    pub fringe_fringe: ResidualBucket,
    /// Prescribed faces between a fringe and a hole cell (normal to it).
    pub fringe_hole: ResidualBucket,
    /// Prescribed faces between two hole cells: not evaluated (their
    /// control volume lies in the hole).
    pub hole_hole_skipped: usize,
    /// Hole cells given a ghost pressure.
    pub hole_ghosts: usize,
    /// Hole–hole faces given a ghost velocity (beyond the solver's stamps).
    pub ghost_faces: usize,
    /// Solved u faces between an active and a fringe cell (the ring's outer
    /// boundary); their residual is the pressure LEVEL offset `δ · h`.
    pub interface_u: usize,
    /// See `interface_u`.
    pub interface_v: usize,
    /// Every prescribed fringe–fringe / fringe–hole face's residual.
    pub prescribed: Vec<FaceResidual>,
}

impl MomentumResidual {
    /// The composite's pressure level offset `δ` (Pa) between the active
    /// cells and the re-stamped fringe: `max |r| / h` over the interface.
    pub fn level_offset(&self, h: f64) -> f64 {
        self.solved_near.max_abs_x.max(self.solved_near.max_abs_y) / h
    }
}

impl OversetPisoSolver {
    /// The momentum residual of the background's own predictor stencil on
    /// every interior background face, after [`Self::advance`] (the field
    /// holds `u^{n+1}`, `u_old = u^n`, `p = p^{n+1}` with the fringe
    /// re-stamped). `dt` is the step just taken.
    pub fn momentum_residual(&self, field: &OversetField, dt: f64) -> MomentumResidual {
        let (nx, ny, dx, dy) = self.grid;
        let rho = self.background.config().density;
        let t_old = self.background.time() - dt;
        let mask = self
            .background
            .mask()
            .expect("the overset background carries a mask");
        let map = &self.overlap;
        let hole = |j: usize, i: usize| map.class(j, i) == CellClass::Hole;

        // The masked copy.
        let mut m = field.background.clone();
        for j in 0..ny {
            for i in 0..nx {
                if hole(j, i) {
                    m.p[(j, i)] = f64::NAN;
                }
            }
        }
        let ghosts = map.hole_p_values(&field.patch.p);
        map.stamp_hole_p(&mut m.p, &ghosts);
        map.stamp_ghost_faces(&mut m, &field.patch.u, &field.patch.v);
        let prescribed_u: HashSet<(usize, usize)> = map
            .fringe_u
            .iter()
            .chain(&map.ghost_u)
            .map(|e| (e.j, e.i))
            .collect();
        let prescribed_v: HashSet<(usize, usize)> = map
            .fringe_v
            .iter()
            .chain(&map.ghost_v)
            .map(|e| (e.j, e.i))
            .collect();
        for j in 0..ny {
            for i in 0..=nx {
                let valid = (i > 0 && !hole(j, i - 1))
                    || (i < nx && !hole(j, i))
                    || prescribed_u.contains(&(j, i));
                if !valid {
                    m.u[(j, i)] = f64::NAN;
                    m.u_old[(j, i)] = f64::NAN;
                }
            }
        }
        for j in 0..=ny {
            for i in 0..nx {
                let valid = (j > 0 && !hole(j - 1, i))
                    || (j < ny && !hole(j, i))
                    || prescribed_v.contains(&(j, i));
                if !valid {
                    m.v[(j, i)] = f64::NAN;
                    m.v_old[(j, i)] = f64::NAN;
                }
            }
        }

        let ghost_u = |j: usize, i: usize| mask.u_kind(j, i) == FaceKind::Ghost;
        let ghost_v = |j: usize, i: usize| mask.v_kind(j, i) == FaceKind::Ghost;
        let active = |j: usize, i: usize| map.class(j, i) == CellClass::Active;
        let mut out = MomentumResidual {
            hole_ghosts: map.hole_p.len(),
            ghost_faces: map.ghost_u.len() + map.ghost_v.len(),
            ..MomentumResidual::default()
        };
        let vol = dx * dy;
        let mut prescribed = Vec::new();
        // Under a limited scheme the far-upwind value enters through
        // `r > 0`, which a NaN fails silently (a silent upwind fallback),
        // so the two-away neighbours are checked explicitly.
        let limited = self.background.parameters().convection_scheme
            != crate::solvers::incompressible::ConvectionScheme::Upwind;
        let far_ok_u =
            |m: &crate::solvers::incompressible::flow_field::FlowField, j: usize, i: usize| {
                !limited
                    || ((i < 2 || m.u_old[(j, i - 2)].is_finite())
                        && (i + 2 > nx || m.u_old[(j, i + 2)].is_finite())
                        && (j < 2 || m.u_old[(j - 2, i)].is_finite())
                        && (j + 2 >= ny || m.u_old[(j + 2, i)].is_finite()))
            };
        let far_ok_v =
            |m: &crate::solvers::incompressible::flow_field::FlowField, j: usize, i: usize| {
                !limited
                    || ((j < 2 || m.v_old[(j - 2, i)].is_finite())
                        && (j + 2 > ny || m.v_old[(j + 2, i)].is_finite())
                        && (i < 2 || m.v_old[(j, i - 2)].is_finite())
                        && (i + 2 >= nx || m.v_old[(j, i + 2)].is_finite()))
            };
        let mu = self.background.config().viscosity;
        let upw = |f: f64, a: f64, b: f64| if f >= 0.0 { a } else { b };
        // The upwind predictor's pieces on an interior u face, × ρ·vol.
        let u_pieces = |m: &crate::solvers::incompressible::flow_field::FlowField,
                        j: usize,
                        i: usize| {
            let (u, v) = (&m.u_old, &m.v_old);
            let ue = 0.5 * (u[(j, i)] + u[(j, i + 1)]);
            let uw = 0.5 * (u[(j, i - 1)] + u[(j, i)]);
            let vn = 0.5 * (v[(j + 1, i - 1)] + v[(j + 1, i)]);
            let vs = 0.5 * (v[(j, i - 1)] + v[(j, i)]);
            let conv = (ue * upw(ue, u[(j, i)], u[(j, i + 1)])
                - uw * upw(uw, u[(j, i - 1)], u[(j, i)]))
                / dx
                + (vn * upw(vn, u[(j, i)], u[(j + 1, i)]) - vs * upw(vs, u[(j - 1, i)], u[(j, i)]))
                    / dy;
            let diff = (u[(j, i + 1)] - 2.0 * u[(j, i)] + u[(j, i - 1)]) / (dx * dx)
                + (u[(j + 1, i)] - 2.0 * u[(j, i)] + u[(j - 1, i)]) / (dy * dy);
            let pres = (m.p[(j, i)] - m.p[(j, i - 1)]) / dx;
            [
                rho * (m.u[(j, i)] - m.u_old[(j, i)]) / dt * vol,
                rho * conv * vol,
                mu * diff * vol,
                pres * vol,
            ]
        };
        let v_pieces = |m: &crate::solvers::incompressible::flow_field::FlowField,
                        j: usize,
                        i: usize| {
            let (u, v) = (&m.u_old, &m.v_old);
            let vn = 0.5 * (v[(j, i)] + v[(j + 1, i)]);
            let vs = 0.5 * (v[(j - 1, i)] + v[(j, i)]);
            let ue = 0.5 * (u[(j - 1, i + 1)] + u[(j, i + 1)]);
            let uw = 0.5 * (u[(j - 1, i)] + u[(j, i)]);
            let conv = (vn * upw(vn, v[(j, i)], v[(j + 1, i)])
                - vs * upw(vs, v[(j - 1, i)], v[(j, i)]))
                / dy
                + (ue * upw(ue, v[(j, i)], v[(j, i + 1)]) - uw * upw(uw, v[(j, i - 1)], v[(j, i)]))
                    / dx;
            let diff = (v[(j + 1, i)] - 2.0 * v[(j, i)] + v[(j - 1, i)]) / (dy * dy)
                + (v[(j, i + 1)] - 2.0 * v[(j, i)] + v[(j, i - 1)]) / (dx * dx);
            let pres = (m.p[(j, i)] - m.p[(j - 1, i)]) / dy;
            [
                rho * (m.v[(j, i)] - m.v_old[(j, i)]) / dt * vol,
                rho * conv * vol,
                mu * diff * vol,
                pres * vol,
            ]
        };

        // u faces (j, i), i = 1..nx: cells (j, i−1) | (j, i).
        for j in 0..ny {
            for i in 1..nx {
                let (w, e) = (map.class(j, i - 1), map.class(j, i));
                if (w == CellClass::Active) != (e == CellClass::Active) {
                    out.interface_u += 1;
                }
                let bucket = if !ghost_u(j, i) {
                    // Stencil: u (j, i±1), (j±1, i); v (j, i−1), (j, i), (j+1, i−1), (j+1, i); p (j, i−1), (j, i).
                    let near = !active(j, i - 1)
                        || !active(j, i)
                        || ghost_u(j, i - 1)
                        || ghost_u(j, i + 1)
                        || (j > 0 && ghost_u(j - 1, i))
                        || (j + 1 < ny && ghost_u(j + 1, i))
                        || ghost_v(j, i - 1)
                        || ghost_v(j, i)
                        || ghost_v(j + 1, i - 1)
                        || ghost_v(j + 1, i);
                    if near {
                        &mut out.solved_near
                    } else {
                        &mut out.solved_far
                    }
                } else {
                    match (w, e) {
                        (CellClass::Fringe, CellClass::Fringe) => &mut out.fringe_fringe,
                        (CellClass::Hole, CellClass::Hole) => {
                            out.hole_hole_skipped += 1;
                            continue;
                        }
                        _ => &mut out.fringe_hole,
                    }
                };
                let rhs = self.background.u_rhs(&m, j, i, t_old);
                let mut r = rho * ((m.u[(j, i)] - m.u_old[(j, i)]) / dt - rhs) * vol;
                if !far_ok_u(&m, j, i) {
                    r = f64::NAN;
                }
                bucket.add(r, true);
                if ghost_u(j, i) {
                    prescribed.push(FaceResidual {
                        is_u: true,
                        j,
                        i,
                        r,
                        fringe_fringe: w == CellClass::Fringe && e == CellClass::Fringe,
                        pieces: u_pieces(&m, j, i),
                    });
                }
            }
        }
        // v faces (j, i), j = 1..ny: cells (j−1, i) | (j, i).
        for j in 1..ny {
            for i in 0..nx {
                let (s, n) = (map.class(j - 1, i), map.class(j, i));
                if (s == CellClass::Active) != (n == CellClass::Active) {
                    out.interface_v += 1;
                }
                let bucket = if !ghost_v(j, i) {
                    let near = !active(j - 1, i)
                        || !active(j, i)
                        || ghost_v(j - 1, i)
                        || ghost_v(j + 1, i)
                        || (i > 0 && ghost_v(j, i - 1))
                        || (i + 1 < nx && ghost_v(j, i + 1))
                        || ghost_u(j - 1, i)
                        || ghost_u(j, i)
                        || ghost_u(j - 1, i + 1)
                        || ghost_u(j, i + 1);
                    if near {
                        &mut out.solved_near
                    } else {
                        &mut out.solved_far
                    }
                } else {
                    match (s, n) {
                        (CellClass::Fringe, CellClass::Fringe) => &mut out.fringe_fringe,
                        (CellClass::Hole, CellClass::Hole) => {
                            out.hole_hole_skipped += 1;
                            continue;
                        }
                        _ => &mut out.fringe_hole,
                    }
                };
                let rhs = self.background.v_rhs(&m, j, i, t_old);
                let mut r = rho * ((m.v[(j, i)] - m.v_old[(j, i)]) / dt - rhs) * vol;
                if !far_ok_v(&m, j, i) {
                    r = f64::NAN;
                }
                bucket.add(r, false);
                if ghost_v(j, i) {
                    prescribed.push(FaceResidual {
                        is_u: false,
                        j,
                        i,
                        r,
                        fringe_fringe: s == CellClass::Fringe && n == CellClass::Fringe,
                        pieces: v_pieces(&m, j, i),
                    });
                }
            }
        }
        out.prescribed = prescribed;
        out
    }
}

impl OversetPisoSolver {
    /// The force on everything inside the box `(i0, i1, j0, j1)` (cell
    /// index bounds, as `EmbeddedMask::control_volume_force`) in the
    /// SOLVER'S OWN flux form: the predictor's convective flux (upwind
    /// plus the scheme's limited correction),
    /// its diffusive flux and the cell pressure, on the momentum control
    /// volumes' faces that make up the box boundary (u volumes `i0 + 1
    /// ..= i1` × `j0 .. j1`, v volumes `i0 .. i1` × `j0 + 1 ..= j1`),
    /// minus the unsteady term over the box's evaluable volumes. On the
    /// solved faces the residual is rounding, so this is box-INDEPENDENT
    /// to rounding as long as the box stays in the active region — the
    /// gate — whereas the control-volume formula moves by ±0.5 % between
    /// boxes. Returns `(fx, fy)`, positive = drag / lift on the body.
    pub fn solver_metric_force(
        &self,
        field: &OversetField,
        dt: f64,
        (i0, i1, j0, j1): (usize, usize, usize, usize),
    ) -> (f64, f64) {
        let (nx, ny, dx, dy) = self.grid;
        assert!(
            i0 >= 1 && i1 + 1 < nx && j0 >= 1 && j1 + 1 < ny,
            "box must be interior"
        );
        let rho = self.background.config().density;
        let mu = self.background.config().viscosity;
        let bg = &field.background;
        let (u, v, p) = (&bg.u_old, &bg.v_old, &bg.p);
        // The predictor's face value: upwind plus the scheme's limited
        // correction with the same far-upwind choice (`u_rhs`), so the
        // flux form is the solver's under upwind AND under TVD.
        let scheme = self.background.parameters().convection_scheme;
        let face = |f: f64, far_up: Option<f64>, up: f64, down: f64, far_dn: Option<f64>| {
            if f >= 0.0 {
                up + scheme.face_correction(far_up, up, down)
            } else {
                down + scheme.face_correction(far_dn, down, up)
            }
        };
        // Outward x-momentum flux through the u-volume face at cell i's
        // centre (n = +x), per unit length.
        let phi_u_x = |j: usize, i: usize| {
            let ue = 0.5 * (u[(j, i)] + u[(j, i + 1)]);
            let uf = face(
                ue,
                (i >= 1).then(|| u[(j, i - 1)]),
                u[(j, i)],
                u[(j, i + 1)],
                (i + 2 <= nx).then(|| u[(j, i + 2)]),
            );
            rho * ue * uf - mu * (u[(j, i + 1)] - u[(j, i)]) / dx + p[(j, i)]
        };
        // Through the u-volume face at v-face row j (n = +y), for u face i.
        let phi_u_y = |j: usize, i: usize| {
            let vn = 0.5 * (v[(j, i - 1)] + v[(j, i)]);
            let uf = face(
                vn,
                (j >= 2).then(|| u[(j - 2, i)]),
                u[(j - 1, i)],
                u[(j, i)],
                (j + 1 < ny).then(|| u[(j + 1, i)]),
            );
            rho * vn * uf - mu * (u[(j, i)] - u[(j - 1, i)]) / dy
        };
        // y-momentum: through the v-volume face at cell j's centre (n = +y).
        let phi_v_y = |j: usize, i: usize| {
            let vn = 0.5 * (v[(j, i)] + v[(j + 1, i)]);
            let vf = face(
                vn,
                (j >= 1).then(|| v[(j - 1, i)]),
                v[(j, i)],
                v[(j + 1, i)],
                (j + 2 <= ny).then(|| v[(j + 2, i)]),
            );
            rho * vn * vf - mu * (v[(j + 1, i)] - v[(j, i)]) / dy + p[(j, i)]
        };
        // Through the v-volume face at u-face column i (n = +x), for v face j.
        let phi_v_x = |j: usize, i: usize| {
            let ue = 0.5 * (u[(j - 1, i)] + u[(j, i)]);
            let vf = face(
                ue,
                (i >= 2).then(|| v[(j, i - 2)]),
                v[(j, i - 1)],
                v[(j, i)],
                (i + 1 < nx).then(|| v[(j, i + 1)]),
            );
            rho * ue * vf - mu * (v[(j, i)] - v[(j, i - 1)]) / dx
        };
        let vol = dx * dy;
        let (mut out_x, mut out_y) = (0.0, 0.0);
        let (mut dt_x, mut dt_y) = (0.0, 0.0);
        for j in j0..j1 {
            out_x += (phi_u_x(j, i1) - phi_u_x(j, i0)) * dy;
        }
        for i in i0 + 1..=i1 {
            out_x += (phi_u_y(j1, i) - phi_u_y(j0, i)) * dx;
            for j in j0..j1 {
                let d = bg.u[(j, i)] - bg.u_old[(j, i)];
                if d.is_finite()
                    && (self.overlap.class(j, i) != CellClass::Hole
                        || self.overlap.class(j, i - 1) != CellClass::Hole)
                {
                    dt_x += rho * d / dt * vol;
                }
            }
        }
        for i in i0..i1 {
            out_y += (phi_v_y(j1, i) - phi_v_y(j0, i)) * dx;
        }
        for j in j0 + 1..=j1 {
            out_y += (phi_v_x(j, i1) - phi_v_x(j, i0)) * dy;
            for i in i0..i1 {
                let d = bg.v[(j, i)] - bg.v_old[(j, i)];
                if d.is_finite()
                    && (self.overlap.class(j, i) != CellClass::Hole
                        || self.overlap.class(j - 1, i) != CellClass::Hole)
                {
                    dt_y += rho * d / dt * vol;
                }
            }
        }
        (-out_x - dt_x, -out_y - dt_y)
    }
}