//! The 3D staggered field: u on `(nx + 1) × ny × nz` faces, v on
//! `nx × (ny + 1) × nz`, w on `nx × ny × (nz + 1)`, p on the cells; flat
//! storage (see [`super::Grid`] for the index conventions).

use super::Grid;

#[derive(Debug, Clone)]
pub struct Field {
    pub grid: Grid,
    pub u: Vec<f64>,
    pub v: Vec<f64>,
    pub w: Vec<f64>,
    pub p: Vec<f64>,
    pub u_old: Vec<f64>,
    pub v_old: Vec<f64>,
    pub w_old: Vec<f64>,
    pub u_star: Vec<f64>,
    pub v_star: Vec<f64>,
    pub w_star: Vec<f64>,
    pub p_prime: Vec<f64>,
    /// The continuity source of the projection.
    pub sp: Vec<f64>,
}

impl Field {
    #[must_use]
    pub fn new(grid: Grid) -> Self {
        let nu = (grid.nx + 1) * grid.ny * grid.nz;
        let nv = grid.nx * (grid.ny + 1) * grid.nz;
        let nw = grid.nx * grid.ny * (grid.nz + 1);
        let nc = grid.cells();
        Self {
            grid,
            u: vec![0.0; nu],
            v: vec![0.0; nv],
            w: vec![0.0; nw],
            p: vec![0.0; nc],
            u_old: vec![0.0; nu],
            v_old: vec![0.0; nv],
            w_old: vec![0.0; nw],
            u_star: vec![0.0; nu],
            v_star: vec![0.0; nv],
            w_star: vec![0.0; nw],
            p_prime: vec![0.0; nc],
            sp: vec![0.0; nc],
        }
    }

    pub fn update_old_values(&mut self) {
        self.u_old.copy_from_slice(&self.u);
        self.v_old.copy_from_slice(&self.v);
        self.w_old.copy_from_slice(&self.w);
    }

    pub fn copy_to_starred(&mut self) {
        self.u_star.copy_from_slice(&self.u);
        self.v_star.copy_from_slice(&self.v);
        self.w_star.copy_from_slice(&self.w);
    }

    /// `max |∇·u|` over the cells (per unit volume).
    #[must_use]
    pub fn max_divergence(&self) -> f64 {
        let g = self.grid;
        let mut worst = 0.0_f64;
        for k in 0..g.nz {
            for j in 0..g.ny {
                for i in 0..g.nx {
                    let div = (self.u[g.uface(k, j, i + 1)] - self.u[g.uface(k, j, i)]) / g.dx
                        + (self.v[g.vface(k, j + 1, i)] - self.v[g.vface(k, j, i)]) / g.dy
                        + (self.w[g.wface(k + 1, j, i)] - self.w[g.wface(k, j, i)]) / g.dz;
                    worst = worst.max(div.abs());
                }
            }
        }
        worst
    }

    /// Kinetic energy `½ Σ (u² + v² + w²) dV` over the cells (face values
    /// averaged to the cell).
    #[must_use]
    pub fn kinetic_energy(&self, rho: f64) -> f64 {
        let g = self.grid;
        let dv = g.dx * g.dy * g.dz;
        let mut e = 0.0;
        for k in 0..g.nz {
            for j in 0..g.ny {
                for i in 0..g.nx {
                    let uc = 0.5 * (self.u[g.uface(k, j, i)] + self.u[g.uface(k, j, i + 1)]);
                    let vc = 0.5 * (self.v[g.vface(k, j, i)] + self.v[g.vface(k, j + 1, i)]);
                    let wc = 0.5 * (self.w[g.wface(k, j, i)] + self.w[g.wface(k + 1, j, i)]);
                    e += 0.5 * rho * (uc * uc + vc * vc + wc * wc) * dv;
                }
            }
        }
        e
    }
}