//! The wall's imposition on the velocity field (`impl Mask` continued
//! from `wall.rs`, split for the file-size rule): prescribed faces take
//! the surface velocity, ghost faces their reconstruction from the source
//! field minus the shared flux compatibility correction.

use super::body::Body;
use super::wall::{FaceKind, Mask};
use rayon::prelude::*;

/// The imposition band of a cut mask, in cells.
pub(crate) const IMPOSE_BAND_CELLS: f64 = 4.0;

impl Mask {
    /// The band (in length units) within which a cut mask's prescribed
    /// faces are imposed; `None` on the binary wall (every solid face).
    #[must_use]
    pub fn impose_band(&self) -> Option<f64> {
        self.cut
            .as_ref()
            .map(|_| IMPOSE_BAND_CELLS * self.grid.dx.min(self.grid.dy).min(self.grid.dz))
    }

    /// Impose the wall on `(u, v, w)` from the same field.
    pub fn impose(&self, body: &Body, u: &mut [f64], v: &mut [f64], w: &mut [f64], t: f64) -> f64 {
        let (us, vs, ws) = (u.to_vec(), v.to_vec(), w.to_vec());
        self.impose_from(body, &us, &vs, &ws, u, v, w, t)
    }

    /// Solid faces: the surface velocity; ghost faces: the reconstruction
    /// from the SOURCE field, minus the shared flux compatibility
    /// correction over the flux-carrying ghosts. Returns the correction.
    #[allow(clippy::too_many_arguments)]
    pub fn impose_from(
        &self,
        body: &Body,
        u_src: &[f64],
        v_src: &[f64],
        w_src: &[f64],
        u: &mut [f64],
        v: &mut [f64],
        w: &mut [f64],
        t: f64,
    ) -> f64 {
        let g = self.grid;
        let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz);
        // On a cut mask only the prescribed faces within the band are
        // imposed: deeper ones are never read (the predictor reaches two
        // faces past an open one) and the surface velocity is the costly
        // part of a moving body's step.
        let band = self.impose_band();
        let near = |d: &[f64], idx: usize| band.is_none_or(|b| d[idx].abs() <= b);
        let (d_u, d_v, d_w): (&[f64], &[f64], &[f64]) = match &self.cut {
            Some(c) => (&c.d_u, &c.d_v, &c.d_w),
            None => (&[], &[], &[]),
        };
        // One z plane per task: every solid face is written from the body's
        // surface velocity alone (no reduction), so the parallel loop is
        // digit-identical to the serial one (PERF-3 P1-1: 1.95 s of a 9.4 s
        // step at 11.6 M cells, serial).
        u.par_chunks_mut(ny * (nx + 1))
            .take(nz)
            .enumerate()
            .for_each(|(k, plane)| {
                for j in 0..ny {
                    for i in 1..nx {
                        let idx = g.uface(k, j, i);
                        if self.u_kind[idx] == FaceKind::Solid && near(d_u, idx) {
                            plane[j * (nx + 1) + i] = body
                                .surface_velocity(
                                    i as f64 * dx,
                                    (j as f64 + 0.5) * dy,
                                    (k as f64 + 0.5) * dz,
                                    t,
                                )
                                .0;
                        }
                    }
                }
            });
        v.par_chunks_mut((ny + 1) * nx)
            .take(nz)
            .enumerate()
            .for_each(|(k, plane)| {
                for j in 1..ny {
                    for i in 0..nx {
                        let idx = g.vface(k, j, i);
                        if self.v_kind[idx] == FaceKind::Solid && near(d_v, idx) {
                            plane[j * nx + i] = body
                                .surface_velocity(
                                    (i as f64 + 0.5) * dx,
                                    j as f64 * dy,
                                    (k as f64 + 0.5) * dz,
                                    t,
                                )
                                .1;
                        }
                    }
                }
            });
        w.par_chunks_mut(ny * nx)
            .take(nz + 1)
            .enumerate()
            .for_each(|(k, plane)| {
                for j in 0..ny {
                    for i in 0..nx {
                        let idx = g.wface(k, j, i);
                        if self.w_kind[idx] == FaceKind::Solid && near(d_w, idx) {
                            plane[j * nx + i] = body
                                .surface_velocity(
                                    (i as f64 + 0.5) * dx,
                                    (j as f64 + 0.5) * dy,
                                    k as f64 * dz,
                                    t,
                                )
                                .2;
                        }
                    }
                }
            });
        let u_vals: Vec<f64> = self
            .u_ghosts
            .iter()
            .map(|gh| gh.reconstruct(u_src))
            .collect();
        let v_vals: Vec<f64> = self
            .v_ghosts
            .iter()
            .map(|gh| gh.reconstruct(v_src))
            .collect();
        let w_vals: Vec<f64> = self
            .w_ghosts
            .iter()
            .map(|gh| gh.reconstruct(w_src))
            .collect();
        let (au, av, aw) = (dy * dz, dx * dz, dx * dy);
        let mut net = 0.0;
        let mut area = 0.0;
        for (gh, &val) in self.u_ghosts.iter().zip(&u_vals) {
            if gh.flux_sign != 0.0 {
                net += gh.flux_sign * val * au;
                area += au;
            }
        }
        for (gh, &val) in self.v_ghosts.iter().zip(&v_vals) {
            if gh.flux_sign != 0.0 {
                net += gh.flux_sign * val * av;
                area += av;
            }
        }
        for (gh, &val) in self.w_ghosts.iter().zip(&w_vals) {
            if gh.flux_sign != 0.0 {
                net += gh.flux_sign * val * aw;
                area += aw;
            }
        }
        let correction = if area > 0.0 { net / area } else { 0.0 };
        for (gh, &val) in self.u_ghosts.iter().zip(&u_vals) {
            u[gh.idx] = val - gh.flux_sign * correction;
        }
        for (gh, &val) in self.v_ghosts.iter().zip(&v_vals) {
            v[gh.idx] = val - gh.flux_sign * correction;
        }
        for (gh, &val) in self.w_ghosts.iter().zip(&w_vals) {
            w[gh.idx] = val - gh.flux_sign * correction;
        }
        // The periodic seam: the w face at k = nz is the face at k = 0.
        for j in 0..ny {
            for i in 0..nx {
                let (f0, fn_) = (g.wface(0, j, i), g.wface(nz, j, i));
                if self.w_kind[f0] != FaceKind::Fluid && self.w_kind[fn_] == self.w_kind[f0] {
                    w[fn_] = w[f0];
                }
            }
        }
        correction
    }
}