//! The apertured cut-cell wall (`WallScheme::CutCell`, item 10 — the
//! "AM-wall"): the cut geometry classifies the grid (a cell is fluid where
//! its fluid volume is positive; an interior face is an unknown where its
//! aperture is positive, prescribed the surface velocity otherwise — no
//! ghost faces), and the wall enters the operators through the apertures:
//! continuity `Σ_f A_f u_f·n_f + U_b·W_c = 0` per cell, the projection
//! coefficient `dt·A_f/δ`, the momentum control volume of an unknown face
//! `V_u = A_f h` with its own faces' apertures averaged from the two
//! adjacent cells (mass fluxes averaged, so the momentum volume conserves
//! mass exactly when the cells do), the wall closing it (`W = −Σ A n`),
//! an implicit wall shear `μ A_w (u_f − U_b)/d_f`, and an inertia floor of
//! 0.1 in the time derivative only. Every prescribed value is the limit of
//! the computed one as the aperture closes (the shear coefficient grows as
//! `1/A_f`), which is what makes the wall smooth in the interface position.

use super::Grid;
use super::body::Body;
use super::cut::CutGeometry;
use super::field::Field;
use super::step::{Boundaries, Side};
use super::wall::{FaceKind, Mask};

/// The inertia floor: the momentum volume's fraction in the time
/// derivative is at least this.
pub(super) const INERTIA_FLOOR: f64 = 0.1;
/// The wall-distance floor of a face, in units of the smallest spacing.
pub(super) const DISTANCE_FLOOR: f64 = 0.05;

/// Lattice addressing of faces and cells with the periodic wrap in z as
/// data: a face of component `c` at `p = [i, j, k]` (its own coordinate is
/// the face index, the others the cell's), a cell at `[i, j, k]`.
#[derive(Debug, Clone, Copy)]
pub(super) struct Lattice {
    pub(super) g: Grid,
    pub(super) periodic_z: bool,
}

impl Lattice {
    fn wrap_z(&self, k: i64, planes: i64) -> Option<usize> {
        if self.periodic_z {
            Some(k.rem_euclid(self.g.nz as i64) as usize)
        } else if (0..planes).contains(&k) {
            Some(k as usize)
        } else {
            None
        }
    }

    /// The index of the face of component `c` at `p`, `None` outside.
    pub(super) fn face(&self, c: usize, p: [i64; 3]) -> Option<usize> {
        let (nx, ny, nz) = (self.g.nx as i64, self.g.ny as i64, self.g.nz as i64);
        let (i, j) = (p[0], p[1]);
        let (ni, nj, nk) = match c {
            0 => (nx + 1, ny, nz),
            1 => (nx, ny + 1, nz),
            _ => (nx, ny, nz + 1),
        };
        if !(0..ni).contains(&i) || !(0..nj).contains(&j) {
            return None;
        }
        let k = self.wrap_z(p[2], nk)?;
        let (i, j) = (i as usize, j as usize);
        Some(match c {
            0 => self.g.uface(k, j, i),
            1 => self.g.vface(k, j, i),
            _ => self.g.wface(k, j, i),
        })
    }

    /// The index of the cell at `p`, `None` outside.
    pub(super) fn cell(&self, p: [i64; 3]) -> Option<usize> {
        let (nx, ny, nz) = (self.g.nx as i64, self.g.ny as i64, self.g.nz as i64);
        if !(0..nx).contains(&p[0]) || !(0..ny).contains(&p[1]) {
            return None;
        }
        let k = self.wrap_z(p[2], nz)?;
        Some(self.g.cell(k, p[1] as usize, p[0] as usize))
    }

    /// The centre of the face of component `c` at `p`.
    pub(super) fn face_position(&self, c: usize, p: [i64; 3]) -> [f64; 3] {
        let h = [self.g.dx, self.g.dy, self.g.dz];
        let mut x = [0.0; 3];
        for d in 0..3 {
            let off = if d == c { 0.0 } else { 0.5 };
            x[d] = (p[d] as f64 + off) * h[d];
        }
        x
    }
}

/// The geometry of an unknown face's momentum control volume.
#[derive(Debug, Clone, Copy)]
pub(super) struct CvGeometry {
    /// The face's own aperture.
    pub(super) alpha: f64,
    /// The control volume's face apertures `[direction][minus, plus]`.
    pub(super) ap: [[f64; 2]; 3],
    /// The wall's vector area closing the control volume (into the body).
    pub(super) wall: [f64; 3],
    /// The wall distance of the face (floored).
    pub(super) distance: f64,
}

impl Mask {
    /// Classify the grid against `body` at `t` by its cut geometry.
    pub fn build_cut(body: &Body, g: Grid, t: f64, b: Boundaries) -> Result<Self, String> {
        let cut = CutGeometry::build(body, g, t);
        let (nx, ny, nz) = (g.nx, g.ny, g.nz);
        let periodic = b.z0 == Side::Periodic;
        let allowed = |side: Side| matches!(side, Side::Velocity | Side::Periodic | Side::SlipWall);
        let mut cell_fluid = vec![true; g.cells()];
        let mut fluid_cells = 0;
        let mut anchor = None;
        for k in 0..nz {
            for j in 0..ny {
                for i in 0..nx {
                    let idx = g.cell(k, j, i);
                    let fluid = cut.vol[idx] > 0.0;
                    cell_fluid[idx] = fluid;
                    if fluid {
                        fluid_cells += 1;
                        if anchor.is_none() {
                            anchor = Some(idx);
                        }
                    } else {
                        let touches = (i == 0 && !allowed(b.x0))
                            || (i + 1 == nx && !allowed(b.x1))
                            || (j == 0 && !allowed(b.y0))
                            || (j + 1 == ny && !allowed(b.y1))
                            || (k == 0 && !allowed(b.z0))
                            || (k + 1 == nz && !allowed(b.z1));
                        if touches {
                            return Err(format!(
                                "embedded body reaches a domain side that is not a Velocity/Periodic side at cell ({k}, {j}, {i})"
                            ));
                        }
                    }
                }
            }
        }
        let Some(anchor) = anchor else {
            return Err("embedded body covers the whole domain".into());
        };
        let kind = |a: f64| {
            if a > 0.0 {
                FaceKind::Fluid
            } else {
                FaceKind::Solid
            }
        };
        let mut u_kind = vec![FaceKind::Fluid; g.n_ufaces()];
        let mut v_kind = vec![FaceKind::Fluid; g.n_vfaces()];
        let mut w_kind = vec![FaceKind::Fluid; g.n_wfaces()];
        for k in 0..nz {
            for j in 0..ny {
                for i in 1..nx {
                    let f = g.uface(k, j, i);
                    u_kind[f] = kind(cut.a_u[f]);
                }
            }
            for j in 1..ny {
                for i in 0..nx {
                    let f = g.vface(k, j, i);
                    v_kind[f] = kind(cut.a_v[f]);
                }
            }
        }
        let w_range = if periodic { 0..nz + 1 } else { 1..nz };
        for k in w_range {
            for j in 0..ny {
                for i in 0..nx {
                    let f = g.wface(k, j, i);
                    w_kind[f] = kind(cut.a_w[f]);
                }
            }
        }
        Ok(Self {
            grid: g,
            periodic_z: periodic,
            cell_fluid,
            u_kind,
            v_kind,
            w_kind,
            u_ghosts: Vec::new(),
            v_ghosts: Vec::new(),
            w_ghosts: Vec::new(),
            anchor,
            fluid_cells,
            cut: Some(cut),
            step_apertures: None,
            step_open: None,
        })
    }

    /// Set the step-averaged apertures and the space-time classification
    /// from the previous mask's geometry.
    pub fn set_step_apertures(&mut self, old: &Mask) {
        let (Some(cut), Some(old_cut)) = (self.cut.as_ref(), old.cut.as_ref()) else {
            return;
        };
        let avg = |a: &[f64], b: &[f64]| -> Vec<f64> {
            a.iter().zip(b).map(|(x, y)| 0.5 * (x + y)).collect()
        };
        let au = avg(&cut.a_u, &old_cut.a_u);
        let av = avg(&cut.a_v, &old_cut.a_v);
        let aw = avg(&cut.a_w, &old_cut.a_w);
        let open = |a: &[f64]| -> Vec<bool> { a.iter().map(|&x| x > 0.0).collect() };
        let active = self
            .cell_fluid
            .iter()
            .zip(&old.cell_fluid)
            .map(|(&n, &o)| n || o)
            .collect();
        self.step_open = Some((open(&au), open(&av), open(&aw), active));
        self.step_apertures = Some((au, av, aw));
    }

    pub(super) fn lattice(&self) -> Lattice {
        Lattice {
            g: self.grid,
            periodic_z: self.periodic_z,
        }
    }

    /// Aperture of the face of component `c` at lattice `p` (1 without a
    /// cut geometry), `None` outside the grid.
    pub(super) fn aperture(&self, c: usize, p: [i64; 3]) -> Option<f64> {
        let f = self.lattice().face(c, p)?;
        Some(match c {
            0 => self.a_u(f),
            1 => self.a_v(f),
            _ => self.a_w(f),
        })
    }

    /// The momentum control volume of the unknown face of component `c` at
    /// `p`: its face apertures are the averages of the two adjacent cells'
    /// (the own-direction faces at the cell centres average the face's and
    /// its own-direction neighbours' apertures), its wall closes them, its
    /// wall distance is the face centre's signed distance moved to the
    /// fluid part's centre, `φ + ½h(1 − α)`, floored.
    pub(super) fn cv_geometry(&self, c: usize, p: [i64; 3]) -> CvGeometry {
        let g = self.grid;
        let h = [g.dx, g.dy, g.dz];
        let area = [g.dy * g.dz, g.dx * g.dz, g.dx * g.dy];
        let mut 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 ec = e(c);
        let cell_minus = add(p, ec, -1);
        let cell_plus = p;
        let alpha = self.aperture(c, p).unwrap_or(1.0);
        let mut ap = [[1.0; 2]; 3];
        for d in 0..3 {
            let ed = e(d);
            if d == c {
                let am = self.aperture(c, add(p, ec, -1)).unwrap_or(alpha);
                let apl = self.aperture(c, add(p, ec, 1)).unwrap_or(alpha);
                ap[d] = [0.5 * (am + alpha), 0.5 * (alpha + apl)];
            } else {
                let minus = 0.5
                    * (self.aperture(d, cell_minus).unwrap_or(1.0)
                        + self.aperture(d, cell_plus).unwrap_or(1.0));
                let plus = 0.5
                    * (self.aperture(d, add(cell_minus, ed, 1)).unwrap_or(1.0)
                        + self.aperture(d, add(cell_plus, ed, 1)).unwrap_or(1.0));
                ap[d] = [minus, plus];
            }
        }
        let mut wall = [0.0; 3];
        for d in 0..3 {
            wall[d] = -(ap[d][1] - ap[d][0]) * area[d];
        }
        let h_min = g.dx.min(g.dy).min(g.dz);
        let phi_face = self.cut.as_ref().map_or(h_min, |cut| {
            let f = self
                .lattice()
                .face(c, p)
                .expect("unknown face inside the grid");
            match c {
                0 => cut.d_u[f],
                1 => cut.d_v[f],
                _ => cut.d_w[f],
            }
        });
        let distance = (phi_face + 0.5 * h[c] * (1.0 - alpha)).max(DISTANCE_FLOOR * h_min);
        CvGeometry {
            alpha,
            ap,
            wall,
            distance,
        }
    }

    /// The surface velocity component `c` at the foot of the normal from
    /// the face centre.
    pub fn surface_velocity_at(&self, body: &Body, x: [f64; 3], c: usize, t: f64) -> f64 {
        let g = self.grid;
        let eps = 1e-6 * g.dx.min(g.dy).min(g.dz);
        let s = body.phi(x[0], x[1], x[2], t);
        let (n1, n2, n3) = body.normal(x[0], x[1], x[2], t, eps);
        let v = body.surface_velocity(x[0] - s * n1, x[1] - s * n2, x[2] - s * n3, t);
        [v.0, v.1, v.2][c]
    }

    /// The volume fluxes of the surface velocity through every cell's wall
    /// into the body, `U_b·W_c` (zero for a body at rest; the porous
    /// manufactured surface's flux otherwise), made compatible: the net
    /// flux (the quadrature's defect on a closed surface — a rigid
    /// translation's is zero by closure) is redistributed over the wall
    /// cells by wall area, as the binary wall's ghost fluxes are. Returns
    /// the table and the correction (flux per unit wall area).
    pub fn wall_flux_table(&self, body: &Body, t: f64) -> (Vec<f64>, f64) {
        let mut table = vec![0.0; self.grid.cells()];
        let Some(cut) = self.cut.as_ref() else {
            return (table, 0.0);
        };
        let (mut net, mut area) = (0.0, 0.0);
        for idx in 0..table.len() {
            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;
            }
            table[idx] = self.wall_flux(body, idx, t);
            net += table[idx];
            area += a;
        }
        let correction = if area > 0.0 { net / area } else { 0.0 };
        if correction != 0.0 {
            for (idx, w) in cut.wall.iter().enumerate() {
                let a = (w[0] * w[0] + w[1] * w[1] + w[2] * w[2]).sqrt();
                table[idx] -= correction * a;
            }
        }
        (table, correction)
    }

    /// The moving rigid body's wall fluxes by the discrete geometric
    /// conservation law: `(V_c^{n+1} − V_c^n)/dt` per active cell (a dying
    /// cell's remaining volume leaves through its step-averaged apertures),
    /// the net (the cut geometry's closure defect) redistributed over the
    /// wall cells by wall area.
    pub fn gcl_flux_table(&self, old: &Mask, dt: f64) -> (Vec<f64>, f64) {
        let mut table = vec![0.0; self.grid.cells()];
        let (Some(cut), Some(old_cut)) = (self.cut.as_ref(), old.cut.as_ref()) else {
            return (table, 0.0);
        };
        let g = self.grid;
        let dv = g.dx * g.dy * g.dz;
        let (mut net, mut area) = (0.0, 0.0);
        for (idx, entry) in table.iter_mut().enumerate() {
            if !self.cell_active(idx) {
                continue;
            }
            *entry = (cut.vol[idx] - old_cut.vol[idx]) * dv / dt;
            net += *entry;
            let w = cut.wall[idx];
            area += (w[0] * w[0] + w[1] * w[1] + w[2] * w[2]).sqrt();
        }
        if std::env::var_os("RTX_E3_DEBUG").is_some() {
            let dead = (0..table.len())
                .filter(|&i| !self.cell_fluid[i] && old.cell_fluid[i])
                .count();
            let fresh = (0..table.len())
                .filter(|&i| self.cell_fluid[i] && !old.cell_fluid[i])
                .count();
            let (vn, vn1): (f64, f64) = (old_cut.vol.iter().sum(), cut.vol.iter().sum());
            eprintln!(
                "    gcl: dead {dead} fresh {fresh} net {net:.3e} area {area:.3e} ΣV old {vn:.6} new {vn1:.6} (Δ {:.3e})",
                vn1 - vn
            );
        }
        let correction = if area > 0.0 { net / area } else { 0.0 };
        if correction != 0.0 {
            for (idx, w) in cut.wall.iter().enumerate() {
                let a = (w[0] * w[0] + w[1] * w[1] + w[2] * w[2]).sqrt();
                table[idx] -= correction * a;
            }
        }
        (table, correction)
    }

    /// The volume flux of the surface velocity through a cell's wall into
    /// the body, `U_b·W_c`, uncorrected. Zero without a cut geometry.
    pub fn wall_flux(&self, body: &Body, idx: usize, t: f64) -> f64 {
        let Some(cut) = self.cut.as_ref() else {
            return 0.0;
        };
        let w = cut.wall[idx];
        if w == [0.0; 3] {
            return 0.0;
        }
        let g = self.grid;
        let (k, j, i) = g.kji(idx);
        let x = [
            (i as f64 + 0.5) * g.dx,
            (j as f64 + 0.5) * g.dy,
            (k as f64 + 0.5) * g.dz,
        ];
        let mut flux = 0.0;
        for (c, wc) in w.iter().enumerate() {
            flux += self.surface_velocity_at(body, x, c, t) * wc;
        }
        flux
    }

    /// 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)?;
        Some([p[0] + s[0], p[1] + s[1], p[2] + s[2]])
    }

    /// The cut-cell load route split into its pressure and shear parts.
    pub fn cut_wall_force_parts(
        &self,
        body: &Body,
        f: &Field,
        mu: f64,
        t: f64,
    ) -> Option<([f64; 3], [f64; 3])> {
        let cut = self.cut.as_ref()?;
        let g = self.grid;
        let (nx, ny, nz) = (g.nx, g.ny, g.nz);
        let mut pressure = [0.0; 3];
        let mut force = [0.0; 3];
        for (idx, w) in cut.wall.iter().enumerate() {
            if self.cell_fluid[idx] {
                for c in 0..3 {
                    pressure[c] += f.p[idx] * w[c];
                }
            }
        }
        let lat = self.lattice();
        let values: [&[f64]; 3] = [&f.u, &f.v, &f.w];
        let w_range = if self.periodic_z { 0..nz } else { 1..nz };
        for c in 0..3 {
            let (ir, jr, kr) = match c {
                0 => (1..nx, 0..ny, 0..nz),
                1 => (0..nx, 1..ny, 0..nz),
                _ => (0..nx, 0..ny, w_range.clone()),
            };
            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");
                        let kind = match c {
                            0 => self.u_kind[idx],
                            1 => self.v_kind[idx],
                            _ => self.w_kind[idx],
                        };
                        if kind != FaceKind::Fluid {
                            continue;
                        }
                        let cv = self.cv_geometry(c, p);
                        let a_w = (cv.wall[0] * cv.wall[0]
                            + cv.wall[1] * cv.wall[1]
                            + cv.wall[2] * cv.wall[2])
                            .sqrt();
                        if a_w == 0.0 {
                            continue;
                        }
                        let ub = self.surface_velocity_at(body, lat.face_position(c, p), c, t);
                        force[c] += mu * a_w * (values[c][idx] - ub) / cv.distance;
                    }
                }
            }
        }
        Some((pressure, force))
    }
}