//! The wall-exchange part of the operator load route (S2-1 remedy, second
//! form). A fluid face control volume next to a prescribed (ghost or
//! solid) face still exchanges momentum with it: the diffusive flux
//! `μ g A (U_b − u_f)/h` through the half-aperture face at the cell
//! centre and the convective flux `ρ m (u_face − u_f)` on the same face.
//! Both are forces the wall exerts on the fluid that the closure polygon's
//! shear `μ A_w (u_f − U_b)/d_f` does not carry, so the operator route
//! read short of the box route by exactly this exchange (a conservation
//! gap that did not shrink with h). Summed here with the predictor's own
//! flux formulas, the operator route closes the discrete momentum balance.
use super::body::Body;
use super::field::Field;
use super::wall::{FaceKind, Mask};
use crate::solvers::incompressible::ConvectionScheme;

/// DIAGNOSTIC: an x window on every load route (`set_load_window`): cells and
/// faces outside `[x0, x1)` are skipped — the cylinder and the flag read
/// apart. Process-wide; `None` (the default) reads the whole body.
static LOAD_WINDOW: std::sync::Mutex<Option<(f64, f64)>> = std::sync::Mutex::new(None);

/// Set or clear the diagnostic x window of the load routes.
pub fn set_load_window(window: Option<(f64, f64)>) {
    *LOAD_WINDOW.lock().expect("load window") = window;
}

pub(super) fn in_load_window(x: f64) -> bool {
    LOAD_WINDOW
        .lock()
        .expect("load window")
        .is_none_or(|(x0, x1)| x >= x0 && x < x1)
}

impl Mask {
    /// The cut-cell load route: the force on the body from the operators
    /// themselves — `Σ_c p_c W_c` over the cells plus the implicit wall
    /// shear `Σ_f μ A_w (u_f − U_b)/d_f` over the unknown faces. `None`
    /// without a cut geometry.
    pub fn cut_wall_force(&self, body: &Body, f: &Field, mu: f64, t: f64) -> Option<[f64; 3]> {
        let (p, s) = self.cut_wall_force_parts(body, f, mu, t)?;
        let x = self.cut_wall_exchange_force(body, f, mu, self.density, t, None)?;
        Some([p[0] + s[0] + x[0], p[1] + s[1] + x[1], p[2] + s[2] + x[2]])
    }

    /// The cut-cell load route restricted to the cells (and faces) of the
    /// planes `k0..k1`, divided by the slab's thickness: the load per unit
    /// span on a body's mid-section.
    pub fn cut_wall_force_per_span(
        &self,
        body: &Body,
        f: &Field,
        mu: f64,
        t: f64,
        (k0, k1): (usize, usize),
    ) -> Option<[f64; 3]> {
        let (p, s) = self.cut_wall_force_parts_in(body, f, mu, t, Some((k0, k1)))?;
        let x = self.cut_wall_exchange_force(body, f, mu, self.density, t, Some((k0, k1)))?;
        let lz = (k1 - k0) as f64 * self.grid.dz;
        Some([
            (p[0] + s[0] + x[0]) / lz,
            (p[1] + s[1] + x[1]) / lz,
            (p[2] + s[2] + x[2]) / lz,
        ])
    }

    /// The momentum the fluid's face control volumes exchange with the
    /// prescribed faces beside them, as a force on the body (the negative
    /// of the force on the fluid), over the z planes `planes` (all when
    /// `None`). `None` without a cut geometry.
    pub fn cut_wall_exchange_force(
        &self,
        body: &Body,
        f: &Field,
        mu: f64,
        rho: f64,
        t: f64,
        planes: Option<(usize, usize)>,
    ) -> Option<[f64; 3]> {
        let (d, c) = self.cut_wall_exchange_parts(body, f, mu, rho, t, planes)?;
        Some([d[0] + c[0], d[1] + c[1], d[2] + c[2]])
    }

    /// The exchange split into its DIFFUSIVE and CONVECTIVE parts (forces on
    /// the body). On a wall at rest the convective part is the scheme's
    /// flux correction only; on a moving wall it carries `ρ m (u_face − u_f)`
    /// with `m` the wall's own swept flux — O(v_wall h / ν) times the shear.
    pub fn cut_wall_exchange_parts(
        &self,
        body: &Body,
        f: &Field,
        mu: f64,
        rho: f64,
        t: f64,
        planes: Option<(usize, usize)>,
    ) -> Option<([f64; 3], [f64; 3])> {
        let _ = body;
        let _ = t;
        self.cut.as_ref()?;
        let g = self.grid;
        let lat = self.lattice();
        let h = [g.dx, g.dy, g.dz];
        let area = [g.dy * g.dz, g.dx * g.dz, g.dx * g.dy];
        let (nx, ny, nz) = (g.nx, g.ny, g.nz);
        let (k0, k1) = planes.unwrap_or((0, nz));
        let w_range = if self.periodic_z { 0..nz } else { 1..nz };
        let vals: [&[f64]; 3] = [&f.u, &f.v, &f.w];
        let scheme = self.scheme;
        let kind = |cc: usize, idx: usize| match cc {
            0 => self.u_kind[idx],
            1 => self.v_kind[idx],
            _ => self.w_kind[idx],
        };
        let val = |cc: usize, q: [i64; 3]| lat.face(cc, q).map(|i| vals[cc][i]);
        let ap = |cc: usize, q: [i64; 3]| self.aperture(cc, q);
        let e = |d: usize| {
            let mut v = [0i64; 3];
            v[d] = 1;
            v
        };
        let add =
            |a: [i64; 3], b: [i64; 3], s: i64| [a[0] + s * b[0], a[1] + s * b[1], a[2] + s * b[2]];
        let upwind = |m: f64, up: f64, dn: f64| if m >= 0.0 { up } else { dn };
        let mut force = [0.0; 3];
        let mut convective = [0.0; 3];
        for c in 0..3 {
            let (ir, jr, kr) = match c {
                0 => (1..nx, 0..ny, k0..k1),
                1 => (0..nx, 1..ny, k0..k1),
                _ => (0..nx, 0..ny, w_range.start.max(k0)..w_range.end.min(k1)),
            };
            let ec = e(c);
            for k in kr {
                for j in jr.clone() {
                    for i in ir.clone() {
                        let p = [i as i64, j as i64, k as i64];
                        let idx = lat.face(c, p).expect("face");
                        if kind(c, idx) != FaceKind::Fluid
                            || !in_load_window(lat.face_position(c, p)[0])
                        {
                            continue;
                        }
                        let cv = self.cv_geometry(c, p);
                        let u0 = vals[c][idx];
                        let shift0 = self.face_shift(c, p);
                        let cell_minus = add(p, ec, -1);
                        let cell_plus = p;
                        for d in 0..3 {
                            let ed = e(d);
                            let a_d = area[d];
                            let up1 = val(c, add(p, ed, 1));
                            let up2 = val(c, add(p, ed, 2));
                            let dn1 = val(c, add(p, ed, -1));
                            let dn2 = val(c, add(p, ed, -2));
                            let (m_plus, m_minus) = if d == c {
                                let f_up =
                                    ap(c, add(p, ec, 1)).unwrap_or(cv.alpha) * up1.unwrap_or(u0);
                                let f_dn =
                                    ap(c, add(p, ec, -1)).unwrap_or(cv.alpha) * dn1.unwrap_or(u0);
                                let f0 = cv.alpha * u0;
                                (0.5 * (f0 + f_up) * a_d, 0.5 * (f_dn + f0) * a_d)
                            } else {
                                let flux = |q: [i64; 3]| {
                                    ap(d, q).unwrap_or(1.0) * val(d, q).unwrap_or(0.0)
                                };
                                (
                                    0.5 * (flux(add(cell_minus, ed, 1))
                                        + flux(add(cell_plus, ed, 1)))
                                        * a_d,
                                    0.5 * (flux(cell_minus) + flux(cell_plus)) * a_d,
                                )
                            };
                            // The spacing the predictor uses toward a solid
                            // neighbour (its shift is zero): the centroid
                            // spacing in a cross direction (S2-5), else the
                            // exchange distance.
                            let solid_spacing = |sign: f64| -> f64 {
                                if self.diffusion_centroid && d != c && !self.wall_exchange_axis {
                                    (h[d] - sign * shift0[d]).clamp(0.25 * h[d], 2.0 * h[d])
                                } else {
                                    self.exchange_delta(&cv, d)
                                }
                            };
                            // Plus side: a prescribed neighbour face.
                            if let Some(fp) = lat.face(c, add(p, ed, 1)) {
                                if kind(c, fp) != FaceKind::Fluid {
                                    let un = vals[c][fp];
                                    let delta = if scheme == ConvectionScheme::Upwind {
                                        0.0
                                    } else if m_plus >= 0.0 {
                                        scheme.face_correction(dn1, u0, un)
                                    } else {
                                        scheme.face_correction(up2, un, u0)
                                    };
                                    let u_face = upwind(m_plus, u0, un) + delta;
                                    if !self.exchange_convection_off {
                                        convective[c] -= -rho * m_plus * (u_face - u0);
                                    }
                                    force[c] -=
                                        mu * cv.ap[d][1] * a_d * (un - u0) / solid_spacing(1.0);
                                }
                            }
                            // Minus side.
                            if let Some(fm) = lat.face(c, add(p, ed, -1)) {
                                if kind(c, fm) != FaceKind::Fluid {
                                    let ud = vals[c][fm];
                                    let delta = if scheme == ConvectionScheme::Upwind {
                                        0.0
                                    } else if m_minus >= 0.0 {
                                        scheme.face_correction(dn2, ud, u0)
                                    } else {
                                        scheme.face_correction(up1, u0, ud)
                                    };
                                    let u_face = upwind(m_minus, ud, u0) + delta;
                                    if !self.exchange_convection_off {
                                        convective[c] -= rho * m_minus * (u_face - u0);
                                    }
                                    force[c] -=
                                        mu * cv.ap[d][0] * a_d * (ud - u0) / solid_spacing(-1.0);
                                }
                            }
                        }
                    }
                }
            }
        }
        Some((force, convective))
    }

    /// The closure lag of a moving body's pressure correction: the
    /// corrector applies `p'` on the STEP apertures while the operator
    /// route reads the summed pressure on the END apertures, so the exact
    /// discrete force carries `Σ_c p'_c (W_step,c − W_end,c)` (a force on
    /// the body) that the route lacks. Zero for a body at rest.
    #[must_use]
    pub fn closure_lag(&self, p_prime: &[f64]) -> [f64; 3] {
        let g = self.grid;
        let area = [g.dy * g.dz, g.dx * g.dz, g.dx * g.dy];
        let mut lag = [0.0; 3];
        if self.step_apertures.is_none() {
            return lag;
        }
        for k in 0..g.nz {
            for j in 0..g.ny {
                for i in 0..g.nx {
                    let idx = g.cell(k, j, i);
                    if !self.cell_active(idx) {
                        continue;
                    }
                    let pp = p_prime[idx];
                    if pp == 0.0 {
                        continue;
                    }
                    // W_c = −Σ A_f n_f: x-part −(α_e − α_w) A_x, etc.
                    let (ue, uw) = (g.uface(k, j, i + 1), g.uface(k, j, i));
                    let (vn, vs) = (g.vface(k, j + 1, i), g.vface(k, j, i));
                    let (wt, wb) = (g.wface(k + 1, j, i), g.wface(k, j, i));
                    let d = [
                        (self.au_step(ue) - self.a_u(ue)) - (self.au_step(uw) - self.a_u(uw)),
                        (self.av_step(vn) - self.a_v(vn)) - (self.av_step(vs) - self.a_v(vs)),
                        (self.aw_step(wt) - self.a_w(wt)) - (self.aw_step(wb) - self.a_w(wb)),
                    ];
                    for c in 0..3 {
                        lag[c] += pp * (-d[c]) * area[c];
                    }
                }
            }
        }
        lag
    }

    /// The force the discrete momentum equation actually applied over the
    /// last step, read post-step: the pressure part on the end field, the
    /// implicit wall shear on the PREDICTED velocities `u*` (the corrector
    /// moves `u` without re-applying the shear) and the explicit wall
    /// exchange on the OLD velocities. On a body at rest at a steady state
    /// this equals `cut_wall_force`; on a moving body it is the number the
    /// box route should reproduce.
    pub fn cut_wall_force_applied(
        &self,
        body: &Body,
        f: &Field,
        mu: f64,
        rho: f64,
        t: f64,
    ) -> Option<[f64; 3]> {
        let mut star = f.clone();
        star.u.copy_from_slice(&f.u_star);
        star.v.copy_from_slice(&f.v_star);
        star.w.copy_from_slice(&f.w_star);
        let (p, s) = self.cut_wall_force_parts(body, &star, mu, t)?;
        let mut old = f.clone();
        old.u.copy_from_slice(&f.u_old);
        old.v.copy_from_slice(&f.v_old);
        old.w.copy_from_slice(&f.w_old);
        let x = self.cut_wall_exchange_force(body, &old, mu, rho, t, None)?;
        Some([p[0] + s[0] + x[0], p[1] + s[1] + x[1], p[2] + s[2] + x[2]])
    }

    /// S2-5 host prototype: per-face pressure-gradient weights `ω = h/δ`,
    /// `δ` the axis distance between the two cells' fluid centroids. A cut
    /// cell of fluid fraction `v` with unit wall normal `n` (into the body)
    /// has its centroid shifted by `−½(1 − v) h n` from the cell centre
    /// (exact for an axis-aligned cut); `δ` is clamped to `[¼h, 2h]`.
    pub fn compute_gradient_weights(&mut self) {
        let Some(cut) = self.cut.as_ref() else {
            return;
        };
        let g = self.grid;
        let h = [g.dx, g.dy, g.dz];
        let cells = g.cells();
        let mut shift = vec![[0.0f64; 3]; cells];
        for idx in 0..cells {
            if !self.cell_fluid[idx] {
                continue;
            }
            let w = cut.wall[idx];
            let a = (w[0] * w[0] + w[1] * w[1] + w[2] * w[2]).sqrt();
            if a == 0.0 {
                continue;
            }
            let v = self.vol(idx);
            for d in 0..3 {
                shift[idx][d] = -0.5 * (1.0 - v) * h[d] * w[d] / a;
            }
        }
        let weight = |d: usize, minus: usize, plus: usize| -> f64 {
            if !(self.cell_fluid[minus] && self.cell_fluid[plus]) {
                return 1.0;
            }
            let delta = h[d] + shift[plus][d] - shift[minus][d];
            h[d] / delta.clamp(0.25 * h[d], 2.0 * h[d])
        };
        let mut wu = vec![1.0; g.n_ufaces()];
        let mut wv = vec![1.0; g.n_vfaces()];
        let mut ww = vec![1.0; g.n_wfaces()];
        for k in 0..g.nz {
            for j in 0..g.ny {
                for i in 0..g.nx {
                    let c = g.cell(k, j, i);
                    if i > 0 {
                        wu[g.uface(k, j, i)] = weight(0, g.cell(k, j, i - 1), c);
                    }
                    if j > 0 {
                        wv[g.vface(k, j, i)] = weight(1, g.cell(k, j - 1, i), c);
                    }
                    if k > 0 {
                        ww[g.wface(k, j, i)] = weight(2, g.cell(k - 1, j, i), c);
                    }
                }
            }
        }
        self.grad_weights = Some((wu, wv, ww));
    }

    /// The pressure force the gradient weights add to the closure sum
    /// `Σ p_c W_c` (a force on the body): `Σ_f α_f A (ω_f − 1)(p_+ − p_−)`.
    #[must_use]
    pub fn gradient_weight_force(&self, p: &[f64], planes: Option<(usize, usize)>) -> [f64; 3] {
        let mut force = [0.0; 3];
        if self.grad_weights.is_none() {
            return force;
        }
        let g = self.grid;
        let area = [g.dy * g.dz, g.dx * g.dz, g.dx * g.dy];
        let (k0, k1) = planes.unwrap_or((0, g.nz));
        for k in k0..k1 {
            for j in 0..g.ny {
                for i in 0..g.nx {
                    let c = g.cell(k, j, i);
                    if i > 0 {
                        let f = g.uface(k, j, i);
                        force[0] += self.a_u(f)
                            * area[0]
                            * (self.grad_weight(0, f) - 1.0)
                            * (p[c] - p[g.cell(k, j, i - 1)]);
                    }
                    if j > 0 {
                        let f = g.vface(k, j, i);
                        force[1] += self.a_v(f)
                            * area[1]
                            * (self.grad_weight(1, f) - 1.0)
                            * (p[c] - p[g.cell(k, j - 1, i)]);
                    }
                    if k > 0 {
                        let f = g.wface(k, j, i);
                        force[2] += self.a_w(f)
                            * area[2]
                            * (self.grad_weight(2, f) - 1.0)
                            * (p[c] - p[g.cell(k - 1, j, i)]);
                    }
                }
            }
        }
        force
    }
}