//! R8-a: the coupled state on disk — the fluid's fields, the flag's
//! kinematic state, the fluid's last centreline and the committed load —
//! so a march can continue on the same grid or be EXTRUDED onto the full
//! duct (the slab's z-average onto every plane: the 3D solver started on
//! the 2D problem's own state, rule 16).

use std::io::{Read as _, Write as _};
use std::path::Path;

use super::fluid::Line;
use rtx_cfd::solvers::incompressible::embedded3::{Field, Grid};

pub struct Saved {
    pub t: f64,
    pub dims: [usize; 3],
    pub h: f64,
    pub u: Vec<f64>,
    pub v: Vec<f64>,
    pub w: Vec<f64>,
    pub p: Vec<f64>,
    pub disp: Vec<f64>,
    pub vel: Vec<f64>,
    pub acc: Vec<f64>,
    pub line: Line,
    pub c_fluid: Vec<f64>,
    /// The committed nodal load (fx, fy per wetted node, in wetted order).
    pub nodal: Vec<f64>,
}

fn write_vec(dir: &Path, name: &str, v: &[f64]) -> std::io::Result<()> {
    let mut f = std::fs::File::create(dir.join(format!("{name}.f64")))?;
    let mut bytes = Vec::with_capacity(8 * v.len());
    for x in v {
        bytes.extend_from_slice(&x.to_le_bytes());
    }
    f.write_all(&bytes)
}

fn read_vec(dir: &Path, name: &str) -> std::io::Result<Vec<f64>> {
    let mut bytes = Vec::new();
    std::fs::File::open(dir.join(format!("{name}.f64")))?.read_to_end(&mut bytes)?;
    Ok(bytes
        .chunks_exact(8)
        .map(|c| f64::from_le_bytes(c.try_into().unwrap()))
        .collect())
}

impl Saved {
    pub fn save(&self, dir: &str) -> std::io::Result<()> {
        let d = Path::new(dir);
        std::fs::create_dir_all(d)?;
        let flat = |pts: &[[f64; 2]]| pts.iter().flat_map(|p| [p[0], p[1]]).collect::<Vec<_>>();
        write_vec(d, "u", &self.u)?;
        write_vec(d, "v", &self.v)?;
        write_vec(d, "w", &self.w)?;
        write_vec(d, "p", &self.p)?;
        write_vec(d, "disp", &self.disp)?;
        write_vec(d, "vel", &self.vel)?;
        write_vec(d, "acc", &self.acc)?;
        write_vec(d, "line_pts", &flat(&self.line.pts))?;
        write_vec(d, "line_vel", &flat(&self.line.vel))?;
        write_vec(d, "c_fluid", &self.c_fluid)?;
        write_vec(d, "nodal", &self.nodal)?;
        std::fs::write(
            d.join("meta.txt"),
            format!(
                "{:e} {} {} {} {:e} {:e}\n",
                self.t, self.dims[0], self.dims[1], self.dims[2], self.h, self.line.t
            ),
        )
    }

    pub fn load(dir: &str) -> std::io::Result<Self> {
        let d = Path::new(dir);
        let meta = std::fs::read_to_string(d.join("meta.txt"))?;
        let m: Vec<&str> = meta.split_whitespace().collect();
        let pairs = |v: Vec<f64>| v.chunks_exact(2).map(|c| [c[0], c[1]]).collect::<Vec<_>>();
        Ok(Self {
            t: m[0].parse().unwrap(),
            dims: [
                m[1].parse().unwrap(),
                m[2].parse().unwrap(),
                m[3].parse().unwrap(),
            ],
            h: m[4].parse().unwrap(),
            u: read_vec(d, "u")?,
            v: read_vec(d, "v")?,
            w: read_vec(d, "w")?,
            p: read_vec(d, "p")?,
            disp: read_vec(d, "disp")?,
            vel: read_vec(d, "vel")?,
            acc: read_vec(d, "acc")?,
            line: Line {
                t: m[5].parse().unwrap(),
                pts: pairs(read_vec(d, "line_pts")?),
                vel: pairs(read_vec(d, "line_vel")?),
            },
            c_fluid: read_vec(d, "c_fluid")?,
            nodal: read_vec(d, "nodal")?,
        })
    }

    /// The saved fields onto `field` (same nx, ny): the saved planes'
    /// z-average on every plane of the target (w = 0: the 2D problem's
    /// state); identical planes copy through when nz matches.
    pub fn fill(&self, field: &mut Field) {
        let g: Grid = field.grid;
        let [nx, ny, nzs] = self.dims;
        assert_eq!(
            (g.nx, g.ny),
            (nx, ny),
            "the saved state's grid differs in x or y"
        );
        assert!(
            (g.dx - self.h).abs() < 1e-12 * self.h,
            "the saved state's h differs"
        );
        let same = g.nz == nzs;
        let src = Grid::cubic(nx, ny, nzs, self.h);
        for j in 0..ny {
            for i in 0..=nx {
                let mean = (0..nzs).map(|k| self.u[src.uface(k, j, i)]).sum::<f64>() / nzs as f64;
                for k in 0..g.nz {
                    field.u[g.uface(k, j, i)] = if same {
                        self.u[src.uface(k, j, i)]
                    } else {
                        mean
                    };
                }
            }
        }
        for j in 0..=ny {
            for i in 0..nx {
                let mean = (0..nzs).map(|k| self.v[src.vface(k, j, i)]).sum::<f64>() / nzs as f64;
                for k in 0..g.nz {
                    field.v[g.vface(k, j, i)] = if same {
                        self.v[src.vface(k, j, i)]
                    } else {
                        mean
                    };
                }
            }
        }
        if same {
            field.w.copy_from_slice(&self.w);
        } else {
            field.w.iter_mut().for_each(|w| *w = 0.0);
        }
        for j in 0..ny {
            for i in 0..nx {
                let mean = (0..nzs).map(|k| self.p[src.cell(k, j, i)]).sum::<f64>() / nzs as f64;
                for k in 0..g.nz {
                    field.p[g.cell(k, j, i)] = if same {
                        self.p[src.cell(k, j, i)]
                    } else {
                        mean
                    };
                }
            }
        }
    }
}