//! The two load routes of the 2D solver in 3D. Route A, surface stress
//! reconstruction: at each surface sample the pressure is linearly
//! extrapolated to the wall from probes at `h` and `2h` along the normal;
//! the wall-normal derivatives of the normal and two tangential velocity
//! components from the same probes with the surface velocity at the wall
//! (quadratic fit); the tangential derivatives of the normal velocity from
//! the surface-velocity function; `σ·n = (−p + 2μ ∂ₙuₙ) n + μ Σₜ (∂ₙuₜ +
//! ∂ₜuₙ) t`. Route B, a momentum balance over a box of whole cells
//! enclosing the body — it reads no near-wall value; the two are unrelated
//! readings of one solution.

use super::body::Body;
use super::field::Field;
use super::wall::{FaceKind, Mask, linear_fit, stencil_nodes, z_planes};

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SurfaceForce {
    pub f: [f64; 3],
    pub samples: usize,
    pub skipped: usize,
}

fn cross(a: [f64; 3], b: [f64; 3]) -> [f64; 3] {
    [
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0],
    ]
}

fn dot(a: [f64; 3], b: [f64; 3]) -> f64 {
    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
}

/// Two unit tangents orthonormal to `n`.
fn tangents(n: [f64; 3]) -> ([f64; 3], [f64; 3]) {
    let a = if n[0].abs() < 0.9 {
        [1.0, 0.0, 0.0]
    } else {
        [0.0, 1.0, 0.0]
    };
    let t1 = cross(n, a);
    let l = dot(t1, t1).sqrt();
    let t1 = [t1[0] / l, t1[1] / l, t1[2] / l];
    (t1, cross(n, t1))
}

impl Mask {
    /// Pressure at a point from the cell centres: trilinear when the
    /// surrounding cells are fluid, else the least-squares fit through the
    /// fluid ones; `None` if degenerate.
    pub fn pressure_at(&self, p: &[f64], x: f64, y: f64, z: f64) -> Option<f64> {
        let g = self.grid();
        let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz);
        let (gx, gy, gz) = (x / dx - 0.5, y / dy - 0.5, z / dz - 0.5);
        let i0 = gx.floor().clamp(0.0, (nx - 2) as f64) as usize;
        let j0 = gy.floor().clamp(0.0, (ny - 2) as f64) as usize;
        let fx = (gx - i0 as f64).clamp(0.0, 1.0);
        let fy = (gy - j0 as f64).clamp(0.0, 1.0);
        let (k0, k1, fz) = z_planes(gz, nz, self.periodic_z());
        // A query on a single weighted plane lives on it (the fit has no
        // z variation to determine).
        let z = if nz == 1 || fz == 0.0 {
            (k0 as f64 + 0.5) * dz
        } else if fz == 1.0 {
            (k1 as f64 + 0.5) * dz
        } else {
            z
        };
        let dirs = |m: usize, f: f64| {
            if m <= 1 {
                vec![(0usize, 1.0)]
            } else {
                vec![(0, 1.0 - f), (1, f)]
            }
        };
        let kz = |dk: usize| if dk == 0 { k0 } else { k1 };
        let mut nodes = Vec::with_capacity(8);
        for (dk, wk) in dirs(nz, fz) {
            for (dj, wj) in dirs(ny, fy) {
                for (di, wi) in dirs(nx, fx) {
                    nodes.push((g.cell(kz(dk), j0 + dj, i0 + di), wi * wj * wk));
                }
            }
        }
        // A virtually merged small cell carries its master's pressure
        // unknown, not its own: it is dropped from the fit.
        let usable = |idx: usize| self.is_fluid_cell(idx) && self.master(idx).is_none();
        if nodes.iter().all(|&(idx, _)| usable(idx)) {
            return Some(nodes.iter().map(|&(idx, w)| w * p[idx]).sum());
        }
        // Weighted by the z-direction weight (a z-invariant field then fits
        // as the 2D four-cell fit); the z weight of node `(dk, dj, di)` is
        // `wk`, recovered from the trilinear product.
        let zw = |wk: f64| wk;
        let mut pts: Vec<(f64, f64, f64, f64, f64)> = Vec::new();
        let mut it = nodes.iter();
        for (dk, wk) in dirs(nz, fz) {
            for _ in 0..dirs(ny, fy).len() * dirs(nx, fx).len() {
                let &(idx, _) = it.next().expect("node");
                let _ = dk;
                if usable(idx) {
                    let (k, j, i) = g.kji(idx);
                    pts.push((
                        (i as f64 + 0.5) * dx,
                        (j as f64 + 0.5) * dy,
                        (k as f64 + 0.5) * dz,
                        p[idx],
                        zw(wk),
                    ));
                }
            }
        }
        linear_fit(&pts, (x, y, z))
    }

    /// Velocity at a point: trilinear over a component's nodes when all are
    /// fluid faces, else the fit through the fluid ones and the point's own
    /// boundary intercept with its surface velocity.
    pub fn velocity_at(
        &self,
        body: &Body,
        f: &Field,
        x: f64,
        y: f64,
        z: f64,
        t: f64,
    ) -> Option<[f64; 3]> {
        let g = self.grid();
        let eps = 1e-6 * g.dx.min(g.dy).min(g.dz);
        let s = body.phi(x, y, z, t);
        let (n1, n2, n3) = body.normal(x, y, z, t, eps);
        let foot = (x - s * n1, y - s * n2, z - s * n3);
        let vf = body.surface_velocity(foot.0, foot.1, foot.2, t);
        let vel_foot = [vf.0, vf.1, vf.2];
        let mut out = [0.0; 3];
        for c in 0..3 {
            // A query on a single weighted plane of this component's nodes
            // lives on it (the fit has no z variation to determine).
            let (planes, offset) = if c == 2 {
                (if self.periodic_z() { g.nz } else { g.nz + 1 }, 0.0)
            } else {
                (g.nz, 0.5)
            };
            let gz = z / g.dz - offset;
            let (k0, k1, fz) = z_planes(gz, planes, self.periodic_z());
            let zq = if planes <= 1 || fz == 0.0 {
                (k0 as f64 + offset) * g.dz
            } else if fz == 1.0 {
                (k1 as f64 + offset) * g.dz
            } else {
                z
            };
            let foot_c = (foot.0, foot.1, if zq == z { foot.2 } else { zq });
            let nodes = stencil_nodes((x, y, zq), c, g, self.periodic_z(), |_| None);
            let values: &[f64] = [&f.u, &f.v, &f.w][c];
            let fluid = |idx: usize| match c {
                0 => self.u_kind(idx) == FaceKind::Fluid,
                1 => self.v_kind(idx) == FaceKind::Fluid,
                _ => self.w_kind(idx) == FaceKind::Fluid,
            };
            if nodes.iter().all(|n| fluid(n.idx)) {
                out[c] = nodes.iter().map(|n| n.weight * values[n.idx]).sum();
            } else {
                let mut pts: Vec<(f64, f64, f64, f64, f64)> = nodes
                    .iter()
                    .filter(|n| fluid(n.idx))
                    .map(|n| (n.x, n.y, n.z, values[n.idx], n.zw))
                    .collect();
                pts.push((foot_c.0, foot_c.1, foot_c.2, vel_foot[c], 1.0));
                out[c] = linear_fit(&pts, (x, y, zq))?;
            }
        }
        Some(out)
    }

    /// The reconstructed traction at one surface point with outward normal `n`.
    pub fn traction_at(
        &self,
        body: &Body,
        f: &Field,
        mu: f64,
        t: f64,
        x: [f64; 3],
        n: [f64; 3],
    ) -> Option<[f64; 3]> {
        let g = self.grid();
        let h = g.dx.min(g.dy).min(g.dz);
        let (d1, d2) = (h, 2.0 * h);
        let at = |d: f64| [x[0] + d * n[0], x[1] + d * n[1], x[2] + d * n[2]];
        let (x1, x2) = (at(d1), at(d2));
        let p1 = self.pressure_at(&f.p, x1[0], x1[1], x1[2])?;
        let p2 = self.pressure_at(&f.p, x2[0], x2[1], x2[2])?;
        let u1 = self.velocity_at(body, f, x1[0], x1[1], x1[2], t)?;
        let u2 = self.velocity_at(body, f, x2[0], x2[1], x2[2], t)?;
        let p_wall = p1 + (p1 - p2) * d1 / (d2 - d1);
        let (t1, t2) = tangents(n);
        let us = body.surface_velocity(x[0], x[1], x[2], t);
        let us = [us.0, us.1, us.2];
        let wall_gradient =
            |f1: f64, f2: f64| (f1 * d2 * d2 - f2 * d1 * d1) / (d1 * d2 * (d2 - d1));
        let dn =
            |dir: [f64; 3]| wall_gradient(dot(u1, dir) - dot(us, dir), dot(u2, dir) - dot(us, dir));
        let eps = 1e-6 * h;
        let dt_un = |dir: [f64; 3]| {
            let p = body.surface_velocity(
                x[0] + eps * dir[0],
                x[1] + eps * dir[1],
                x[2] + eps * dir[2],
                t,
            );
            let m = body.surface_velocity(
                x[0] - eps * dir[0],
                x[1] - eps * dir[1],
                x[2] - eps * dir[2],
                t,
            );
            ((p.0 - m.0) * n[0] + (p.1 - m.1) * n[1] + (p.2 - m.2) * n[2]) / (2.0 * eps)
        };
        let traction_n = -p_wall + 2.0 * mu * dn(n);
        let traction_t1 = mu * (dn(t1) + dt_un(t1));
        let traction_t2 = mu * (dn(t2) + dt_un(t2));
        Some([
            traction_n * n[0] + traction_t1 * t1[0] + traction_t2 * t2[0],
            traction_n * n[1] + traction_t1 * t1[1] + traction_t2 * t2[1],
            traction_n * n[2] + traction_t1 * t1[2] + traction_t2 * t2[2],
        ])
    }

    /// Route A: the surface integral of the reconstructed traction over the
    /// body's samples at spacing `ds`.
    pub fn surface_force(&self, body: &Body, f: &Field, mu: f64, t: f64, ds: f64) -> SurfaceForce {
        let samples = body.surface_samples(ds);
        let mut force = [0.0; 3];
        let mut skipped = 0;
        for s in &samples {
            match self.traction_at(body, f, mu, t, [s.x, s.y, s.z], [s.nx, s.ny, s.nz]) {
                Some(tr) => {
                    for c in 0..3 {
                        force[c] += tr[c] * s.area;
                    }
                }
                None => skipped += 1,
            }
        }
        SurfaceForce {
            f: force,
            samples: samples.len(),
            skipped,
        }
    }

    /// Route B: the momentum balance over the box of whole cells
    /// `[i0, i1) × [j0, j1) × [k0, k1)` (in the fluid on its boundary):
    /// `F = Σ_outer (σ·n − ρ u (u·n)) A − d/dt ∫ ρ u dV + ∫ f dV`.
    #[allow(clippy::too_many_arguments)]
    pub fn control_volume_force(
        &self,
        f: &Field,
        dt: f64,
        rho: f64,
        mu: f64,
        source: Option<&dyn Fn(f64, f64, f64) -> (f64, f64, f64)>,
        bx: (usize, usize, usize, usize, usize, usize),
    ) -> [f64; 3] {
        self.control_volume_force_with_walls(f, dt, rho, mu, source, bx, false)
    }

    /// Route B on a box spanning the whole z range between NO-SLIP z walls
    /// (`no_slip_z`): the walls' shear on the fluid inside the box (against
    /// a wall at rest) is an outer-face stress and enters the balance;
    /// without it (slip or periodic sides) the z faces carry nothing.
    #[allow(clippy::too_many_arguments)]
    pub fn control_volume_force_with_walls(
        &self,
        f: &Field,
        dt: f64,
        rho: f64,
        mu: f64,
        source: Option<&dyn Fn(f64, f64, f64) -> (f64, f64, f64)>,
        (i0, i1, j0, j1, k0, k1): (usize, usize, usize, usize, usize, usize),
        no_slip_z: bool,
    ) -> [f64; 3] {
        let g = self.grid();
        let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz);
        let (u, v, w, p) = (&f.u, &f.v, &f.w, &f.p);
        let mut force = [0.0; 3];
        let uc =
            |k: usize, j: usize, i: usize| 0.5 * (u[g.uface(k, j, i)] + u[g.uface(k, j, i + 1)]);
        let vc =
            |k: usize, j: usize, i: usize| 0.5 * (v[g.vface(k, j, i)] + v[g.vface(k, j + 1, i)]);
        let wc =
            |k: usize, j: usize, i: usize| 0.5 * (w[g.wface(k, j, i)] + w[g.wface(k + 1, j, i)]);
        // Central inside, one-sided at the domain edge.
        let dd = |prev: Option<f64>, here: f64, next: Option<f64>, h: f64| match (prev, next) {
            (Some(a), Some(b)) => (b - a) / (2.0 * h),
            (None, Some(b)) => (b - here) / h,
            (Some(a), None) => (here - a) / h,
            (None, None) => 0.0,
        };
        // x faces: u lives there.
        for k in k0..k1 {
            for j in j0..j1 {
                for (i, sign) in [(i1, 1.0), (i0, -1.0)] {
                    let un = u[g.uface(k, j, i)];
                    let p_f = 0.5 * (p[g.cell(k, j, i - 1)] + p[g.cell(k, j, i)]);
                    let dudx = (u[g.uface(k, j, i + 1)] - u[g.uface(k, j, i - 1)]) / (2.0 * dx);
                    let dvdx = (vc(k, j, i) - vc(k, j, i - 1)) / dx;
                    let dwdx = (wc(k, j, i) - wc(k, j, i - 1)) / dx;
                    let dudy = dd(
                        (j > 0).then(|| u[g.uface(k, j - 1, i)]),
                        un,
                        (j + 1 < ny).then(|| u[g.uface(k, j + 1, i)]),
                        dy,
                    );
                    let dudz = dd(
                        (k > 0).then(|| u[g.uface(k - 1, j, i)]),
                        un,
                        (k + 1 < nz).then(|| u[g.uface(k + 1, j, i)]),
                        dz,
                    );
                    let v_f = 0.5 * (vc(k, j, i - 1) + vc(k, j, i));
                    let w_f = 0.5 * (wc(k, j, i - 1) + wc(k, j, i));
                    let a = dy * dz;
                    force[0] += sign * ((-p_f + 2.0 * mu * dudx) - rho * un * un) * a;
                    force[1] += sign * (mu * (dudy + dvdx) - rho * v_f * un) * a;
                    force[2] += sign * (mu * (dudz + dwdx) - rho * w_f * un) * a;
                }
            }
        }
        // y faces: v lives there.
        for k in k0..k1 {
            for i in i0..i1 {
                for (j, sign) in [(j1, 1.0), (j0, -1.0)] {
                    let vn = v[g.vface(k, j, i)];
                    let p_f = 0.5 * (p[g.cell(k, j - 1, i)] + p[g.cell(k, j, i)]);
                    let dvdy = (v[g.vface(k, j + 1, i)] - v[g.vface(k, j - 1, i)]) / (2.0 * dy);
                    let dudy = (uc(k, j, i) - uc(k, j - 1, i)) / dy;
                    let dwdy = (wc(k, j, i) - wc(k, j - 1, i)) / dy;
                    let dvdx = dd(
                        (i > 0).then(|| v[g.vface(k, j, i - 1)]),
                        vn,
                        (i + 1 < nx).then(|| v[g.vface(k, j, i + 1)]),
                        dx,
                    );
                    let dvdz = dd(
                        (k > 0).then(|| v[g.vface(k - 1, j, i)]),
                        vn,
                        (k + 1 < nz).then(|| v[g.vface(k + 1, j, i)]),
                        dz,
                    );
                    let u_f = 0.5 * (uc(k, j - 1, i) + uc(k, j, i));
                    let w_f = 0.5 * (wc(k, j - 1, i) + wc(k, j, i));
                    let a = dx * dz;
                    force[0] += sign * (mu * (dudy + dvdx) - rho * u_f * vn) * a;
                    force[1] += sign * ((-p_f + 2.0 * mu * dvdy) - rho * vn * vn) * a;
                    force[2] += sign * (mu * (dvdz + dwdy) - rho * w_f * vn) * a;
                }
            }
        }
        // z faces: w lives there. A box spanning the whole z range has its
        // z faces on the domain's z sides: no momentum flux and no shear on
        // a slip wall, and the two faces cancel on a periodic pair — skipped.
        let full_span = k0 == 0 && k1 == nz;
        if full_span && no_slip_z {
            let a = dx * dy;
            for j in j0..j1 {
                for i in i0..i1 {
                    for k in [0, nz - 1] {
                        if !self.is_fluid_cell(g.cell(k, j, i)) {
                            continue;
                        }
                        force[0] -= mu * uc(k, j, i) / (0.5 * dz) * a;
                        force[1] -= mu * vc(k, j, i) / (0.5 * dz) * a;
                    }
                }
            }
        }
        for j in (j0..j1).filter(|_| !full_span) {
            for i in i0..i1 {
                for (k, sign) in [(k1, 1.0), (k0, -1.0)] {
                    let wn = w[g.wface(k, j, i)];
                    let p_f = 0.5 * (p[g.cell(k - 1, j, i)] + p[g.cell(k, j, i)]);
                    let dwdz = (w[g.wface(k + 1, j, i)] - w[g.wface(k - 1, j, i)]) / (2.0 * dz);
                    let dudz = (uc(k, j, i) - uc(k - 1, j, i)) / dz;
                    let dvdz = (vc(k, j, i) - vc(k - 1, j, i)) / dz;
                    let dwdx = dd(
                        (i > 0).then(|| w[g.wface(k, j, i - 1)]),
                        wn,
                        (i + 1 < nx).then(|| w[g.wface(k, j, i + 1)]),
                        dx,
                    );
                    let dwdy = dd(
                        (j > 0).then(|| w[g.wface(k, j - 1, i)]),
                        wn,
                        (j + 1 < ny).then(|| w[g.wface(k, j + 1, i)]),
                        dy,
                    );
                    let u_f = 0.5 * (uc(k - 1, j, i) + uc(k, j, i));
                    let v_f = 0.5 * (vc(k - 1, j, i) + vc(k, j, i));
                    let a = dx * dy;
                    force[0] += sign * (mu * (dudz + dwdx) - rho * u_f * wn) * a;
                    force[1] += sign * (mu * (dvdz + dwdy) - rho * v_f * wn) * a;
                    force[2] += sign * ((-p_f + 2.0 * mu * dwdz) - rho * wn * wn) * a;
                }
            }
        }
        // Unsteady term and source over the fluid cells of the box.
        let dv = dx * dy * dz;
        for k in k0..k1 {
            for j in j0..j1 {
                for i in i0..i1 {
                    let idx = g.cell(k, j, i);
                    if !self.is_fluid_cell(idx) {
                        continue;
                    }
                    let dv = dv * self.vol(idx);
                    let (fu0, fu1) = (g.uface(k, j, i), g.uface(k, j, i + 1));
                    let (fv0, fv1) = (g.vface(k, j, i), g.vface(k, j + 1, i));
                    let (fw0, fw1) = (g.wface(k, j, i), g.wface(k + 1, j, i));
                    let du = 0.5 * ((u[fu0] - f.u_old[fu0]) + (u[fu1] - f.u_old[fu1]));
                    let dvv = 0.5 * ((v[fv0] - f.v_old[fv0]) + (v[fv1] - f.v_old[fv1]));
                    let dw = 0.5 * ((w[fw0] - f.w_old[fw0]) + (w[fw1] - f.w_old[fw1]));
                    force[0] -= rho * du / dt * dv;
                    force[1] -= rho * dvv / dt * dv;
                    force[2] -= rho * dw / dt * dv;
                    if let Some(s) = source {
                        let (sx, sy, sz) = s(
                            (i as f64 + 0.5) * dx,
                            (j as f64 + 0.5) * dy,
                            (k as f64 + 0.5) * dz,
                        );
                        force[0] += sx * dv;
                        force[1] += sy * dv;
                        force[2] += sz * dv;
                    }
                }
            }
        }
        force
    }
}