//! 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;

impl Mask {
    /// 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 _ = 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];
        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 {
                            continue;
                        }
                        let cv = self.cv_geometry(c, p);
                        let u0 = vals[c][idx];
                        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,
                                )
                            };
                            // 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;
                                    let on_fluid = -rho * m_plus * (u_face - u0)
                                        + mu * cv.ap[d][1] * a_d * (un - u0) / h[d];
                                    force[c] -= on_fluid;
                                }
                            }
                            // 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;
                                    let on_fluid = rho * m_minus * (u_face - u0)
                                        + mu * cv.ap[d][0] * a_d * (ud - u0) / h[d];
                                    force[c] -= on_fluid;
                                }
                            }
                        }
                    }
                }
            }
        }
        Some(force)
    }

    /// 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]])
    }
}