// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Turn a march fluid-field dump (`RTX_FSI{2,3}_FFLD`, see the FSI
//! harness's `MarchConfig::ffld_dir`) into clawview-readable legacy VTK:
//!
//! - one 2D snapshot per dump (`snap_STEP.vtk`): the MAC grid as
//!   triangles, ONE point scalar (`--field p|umag|vort`, default
//!   `vort` — clawview's legacy-VTK path reads exactly one) and the
//!   solver's own 0/1 fluid mask as integer CELL_DATA;
//! - optionally a space-time volume (`spacetime.vtk`, `--spacetime
//!   <field>`): selected frames stacked along z = time, each triangle
//!   extruded to a prism split into 3 tets, POINT_DATA `phi` = the
//!   chosen field averaged to nodes. clawview's slice-plane animation
//!   over z then plays the transient.
//!
//! Usage:
//!   ffld_to_vtk <dump_dir> [--out <dir>] [--spacetime p|umag|vort]
//!               [--stride N] [--zscale S] [--coarsen K]
//!
//! `--stride N` uses every Nth dump for the volume (default 1),
//! `--zscale` metres of z per second of t (default: domain height per
//! total time span), `--coarsen K` merges K×K cells per volume cell
//! (default 2 — the volume is for structure, the snapshots for detail).

use rtx_cfd::solvers::incompressible::FlowField;
use std::fmt::Write as _;
use std::path::{Path, PathBuf};

struct Frame {
    step: usize,
    t: f64,
    field: FlowField,
    /// Row-major `j * nx + i`, the solver's own mask.
    fluid: Vec<bool>,
}

fn read_index(dir: &Path) -> Vec<(usize, f64, PathBuf)> {
    let index = std::fs::read_to_string(dir.join("index.csv")).expect("index.csv in dump dir");
    index
        .lines()
        .filter(|l| !l.trim().is_empty())
        .map(|l| {
            let mut parts = l.splitn(3, ',');
            let step: usize = parts.next().unwrap().parse().expect("step");
            let t: f64 = parts.next().unwrap().parse().expect("t");
            let name = parts.next().expect("ffld name");
            (step, t, dir.join(name))
        })
        .collect()
}

fn read_mask(dir: &Path, step: usize, nx: usize, ny: usize) -> Vec<bool> {
    let text =
        std::fs::read_to_string(dir.join(format!("mask_{step:06}.txt"))).expect("mask sidecar");
    let mut fluid = Vec::with_capacity(nx * ny);
    for line in text.lines() {
        for c in line.chars() {
            fluid.push(c == '1');
        }
    }
    assert_eq!(fluid.len(), nx * ny, "mask size mismatch at step {step}");
    fluid
}

/// Corner (node) vorticity of the MAC field: `dv/dx - du/dy` at grid
/// node `(i, j)` from the four adjacent staggered faces; zero on the
/// domain boundary nodes.
fn node_vorticity(f: &FlowField) -> Vec<f64> {
    let (nx, ny, dx, dy) = f.grid_info();
    let mut w = vec![0.0; (nx + 1) * (ny + 1)];
    for j in 1..ny {
        for i in 1..nx {
            let dvdx = (f.v[(j, i)] - f.v[(j, i - 1)]) / dx;
            let dudy = (f.u[(j, i)] - f.u[(j - 1, i)]) / dy;
            w[j * (nx + 1) + i] = dvdx - dudy;
        }
    }
    w
}

/// Cell-centred scalar fields `(p, umag, vort)` of a frame.
fn cell_fields(frame: &Frame) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
    let f = &frame.field;
    let (nx, ny, _, _) = f.grid_info();
    let wn = node_vorticity(f);
    let mut p = Vec::with_capacity(nx * ny);
    let mut umag = Vec::with_capacity(nx * ny);
    let mut vort = Vec::with_capacity(nx * ny);
    for j in 0..ny {
        for i in 0..nx {
            if frame.fluid[j * nx + i] {
                p.push(f.p[(j, i)]);
                let uc = 0.5 * (f.u[(j, i)] + f.u[(j, i + 1)]);
                let vc = 0.5 * (f.v[(j, i)] + f.v[(j + 1, i)]);
                umag.push((uc * uc + vc * vc).sqrt());
                let s = wn[j * (nx + 1) + i]
                    + wn[j * (nx + 1) + i + 1]
                    + wn[(j + 1) * (nx + 1) + i]
                    + wn[(j + 1) * (nx + 1) + i + 1];
                vort.push(0.25 * s);
            } else {
                p.push(0.0);
                umag.push(0.0);
                vort.push(0.0);
            }
        }
    }
    (p, umag, vort)
}

fn write_scalar(out: &mut String, name: &str, values: &[f64]) {
    writeln!(out, "SCALARS {name} float 1").unwrap();
    writeln!(out, "LOOKUP_TABLE default").unwrap();
    for v in values {
        writeln!(out, "{v:.6e}").unwrap();
    }
}

/// One 2D snapshot as legacy VTK triangles, with `field_name` as the
/// single point scalar.
fn write_snapshot_vtk(frame: &Frame, field_name: &str, path: &Path) {
    let (nx, ny, dx, dy) = frame.field.grid_info();
    let (p, umag, _vort) = cell_fields(frame);
    let wn = node_vorticity(&frame.field);

    let mut out = String::new();
    writeln!(out, "# vtk DataFile Version 3.0").unwrap();
    writeln!(out, "rtx-cfd snapshot step {} t {:.6}", frame.step, frame.t).unwrap();
    writeln!(out, "ASCII").unwrap();
    writeln!(out, "DATASET UNSTRUCTURED_GRID").unwrap();
    writeln!(out, "POINTS {} float", (nx + 1) * (ny + 1)).unwrap();
    for j in 0..=ny {
        for i in 0..=nx {
            writeln!(out, "{:.6e} {:.6e} 0.0", i as f64 * dx, j as f64 * dy).unwrap();
        }
    }
    let ncells = 2 * nx * ny;
    writeln!(out, "CELLS {ncells} {}", 4 * ncells).unwrap();
    let node = |j: usize, i: usize| j * (nx + 1) + i;
    for j in 0..ny {
        for i in 0..nx {
            let (a, b, c, d) = (
                node(j, i),
                node(j, i + 1),
                node(j + 1, i + 1),
                node(j + 1, i),
            );
            writeln!(out, "3 {a} {b} {c}").unwrap();
            writeln!(out, "3 {a} {c} {d}").unwrap();
        }
    }
    writeln!(out, "CELL_TYPES {ncells}").unwrap();
    for _ in 0..ncells {
        writeln!(out, "5").unwrap();
    }
    // clawview's legacy-VTK path (found by the end-to-end check, not the
    // spec): CELL_DATA scalars are INTEGER markers, and it supports
    // exactly ONE point scalar (every POINT_DATA block appends into the
    // same `phi`). So: one selected field as the point scalar
    // (mask-aware node average; corner vorticity exact), and CELL_DATA
    // carries only the 0/1 fluid marker as integers.
    writeln!(out, "CELL_DATA {ncells}").unwrap();
    writeln!(out, "SCALARS fluid int 1").unwrap();
    writeln!(out, "LOOKUP_TABLE default").unwrap();
    for &f in &frame.fluid {
        let m = i32::from(f);
        writeln!(out, "{m}").unwrap();
        writeln!(out, "{m}").unwrap();
    }
    writeln!(out, "POINT_DATA {}", (nx + 1) * (ny + 1)).unwrap();
    let node_values = match field_name {
        "p" => node_average(&p, &frame.fluid, nx, ny),
        "umag" => node_average(&umag, &frame.fluid, nx, ny),
        "vort" => wn,
        other => panic!("unknown field {other:?} (use p|umag|vort)"),
    };
    write_scalar(&mut out, field_name, &node_values);
    std::fs::write(path, out).expect("write snapshot vtk");
}

/// Node values of a cell field: mask-aware average of adjacent cells.
fn node_average(cell: &[f64], fluid: &[bool], nx: usize, ny: usize) -> Vec<f64> {
    let mut node = vec![0.0; (nx + 1) * (ny + 1)];
    for j in 0..=ny {
        for i in 0..=nx {
            let mut sum = 0.0;
            let mut count = 0usize;
            let mut visit = |jj: isize, ii: isize| {
                if jj >= 0 && ii >= 0 && (jj as usize) < ny && (ii as usize) < nx {
                    let idx = jj as usize * nx + ii as usize;
                    if fluid[idx] {
                        sum += cell[idx];
                        count += 1;
                    }
                }
            };
            visit(j as isize - 1, i as isize - 1);
            visit(j as isize - 1, i as isize);
            visit(j as isize, i as isize - 1);
            visit(j as isize, i as isize);
            if count > 0 {
                node[j * (nx + 1) + i] = sum / count as f64;
            }
        }
    }
    node
}

/// The space-time volume: frames stacked along z, coarsened `k`x`k`,
/// each coarse triangle extruded to the next frame and split into 3
/// tets, POINT_DATA `phi` = the chosen field at the nodes.
#[allow(clippy::too_many_lines)]
fn write_spacetime_vtk(frames: &[&Frame], which: &str, zscale: f64, k: usize, path: &Path) {
    let (nx, ny, dx, dy) = frames[0].field.grid_info();
    let (cnx, cny) = (nx / k, ny / k);
    let nodes_per_frame = (cnx + 1) * (cny + 1);
    let t0 = frames[0].t;

    let mut points = String::new();
    let mut phi = Vec::new();
    for frame in frames {
        let (p, umag, vort) = cell_fields(frame);
        let cell = match which {
            "p" => &p,
            "umag" => &umag,
            "vort" => &vort,
            other => panic!("unknown spacetime field {other:?} (use p|umag|vort)"),
        };
        let node = node_average(cell, &frame.fluid, nx, ny);
        let z = (frame.t - t0) * zscale;
        for j in 0..=cny {
            for i in 0..=cnx {
                let (fi, fj) = ((i * k).min(nx), (j * k).min(ny));
                writeln!(
                    points,
                    "{:.6e} {:.6e} {:.6e}",
                    fi as f64 * dx,
                    fj as f64 * dy,
                    z
                )
                .unwrap();
                phi.push(node[fj * (nx + 1) + fi]);
            }
        }
    }

    let mut cells = String::new();
    let mut ncells = 0usize;
    let cnode = |f: usize, j: usize, i: usize| f * nodes_per_frame + j * (cnx + 1) + i;
    for f in 0..frames.len() - 1 {
        for j in 0..cny {
            for i in 0..cnx {
                let quad = [(j, i), (j, i + 1), (j + 1, i + 1), (j + 1, i)];
                for tri in [[0usize, 1, 2], [0, 2, 3]] {
                    let a = cnode(f, quad[tri[0]].0, quad[tri[0]].1);
                    let b = cnode(f, quad[tri[1]].0, quad[tri[1]].1);
                    let c = cnode(f, quad[tri[2]].0, quad[tri[2]].1);
                    let (a2, b2, c2) = (
                        a + nodes_per_frame,
                        b + nodes_per_frame,
                        c + nodes_per_frame,
                    );
                    // Prism (a, b, c | a2, b2, c2) as 3 tets.
                    for tet in [[a, b, c, a2], [b, c, a2, b2], [c, a2, b2, c2]] {
                        writeln!(cells, "4 {} {} {} {}", tet[0], tet[1], tet[2], tet[3]).unwrap();
                        ncells += 1;
                    }
                }
            }
        }
    }

    let npoints = frames.len() * nodes_per_frame;
    let mut out = String::new();
    writeln!(out, "# vtk DataFile Version 3.0").unwrap();
    writeln!(
        out,
        "rtx-cfd space-time volume ({which}), {} frames, z = (t - {t0:.4}) * {zscale:.4}",
        frames.len()
    )
    .unwrap();
    writeln!(out, "ASCII").unwrap();
    writeln!(out, "DATASET UNSTRUCTURED_GRID").unwrap();
    writeln!(out, "POINTS {npoints} float").unwrap();
    out.push_str(&points);
    writeln!(out, "CELLS {ncells} {}", 5 * ncells).unwrap();
    out.push_str(&cells);
    writeln!(out, "CELL_TYPES {ncells}").unwrap();
    for _ in 0..ncells {
        writeln!(out, "10").unwrap();
    }
    writeln!(out, "POINT_DATA {npoints}").unwrap();
    write_scalar(&mut out, "phi", &phi);
    std::fs::write(path, out).expect("write spacetime vtk");
    println!(
        "  spacetime: {} frames, {npoints} points, {ncells} tets -> {}",
        frames.len(),
        path.display()
    );
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let mut dump_dir: Option<PathBuf> = None;
    let mut out_dir: Option<PathBuf> = None;
    let mut spacetime: Option<String> = None;
    let mut field = "vort".to_string();
    let mut stride = 1usize;
    let mut zscale: Option<f64> = None;
    let mut coarsen = 2usize;
    let mut it = args.iter().skip(1);
    while let Some(a) = it.next() {
        match a.as_str() {
            "--out" => out_dir = Some(PathBuf::from(it.next().expect("--out value"))),
            "--field" => field = it.next().expect("--field value").clone(),
            "--spacetime" => spacetime = Some(it.next().expect("--spacetime value").clone()),
            "--stride" => stride = it.next().expect("--stride value").parse().expect("stride"),
            "--zscale" => {
                zscale = Some(it.next().expect("--zscale value").parse().expect("zscale"))
            }
            "--coarsen" => {
                coarsen = it
                    .next()
                    .expect("--coarsen value")
                    .parse()
                    .expect("coarsen")
            }
            other => dump_dir = Some(PathBuf::from(other)),
        }
    }
    let dump_dir = dump_dir.expect("usage: ffld_to_vtk <dump_dir> [--out d] [--spacetime f]");
    let out_dir = out_dir.unwrap_or_else(|| dump_dir.clone());
    std::fs::create_dir_all(&out_dir).expect("out dir");

    let index = read_index(&dump_dir);
    assert!(!index.is_empty(), "empty index.csv");
    let mut frames = Vec::new();
    for &(step, t, ref path) in &index {
        let field = FlowField::load(path).expect("ffld load");
        let (nx, ny, _, _) = field.grid_info();
        let fluid = read_mask(&dump_dir, step, nx, ny);
        frames.push(Frame {
            step,
            t,
            field,
            fluid,
        });
    }
    println!(
        "  {} frames, t in [{:.4}, {:.4}]",
        frames.len(),
        frames.first().unwrap().t,
        frames.last().unwrap().t
    );

    for frame in &frames {
        let path = out_dir.join(format!("snap_{:06}.vtk", frame.step));
        write_snapshot_vtk(frame, &field, &path);
    }
    println!("  {} snapshot VTKs -> {}", frames.len(), out_dir.display());

    if let Some(which) = spacetime {
        let picked: Vec<&Frame> = frames.iter().step_by(stride.max(1)).collect();
        assert!(picked.len() >= 2, "spacetime needs at least 2 frames");
        let (_, ny, _, dy) = picked[0].field.grid_info();
        let span = picked.last().unwrap().t - picked[0].t;
        let zscale = zscale.unwrap_or_else(|| ny as f64 * dy / span.max(1e-12));
        write_spacetime_vtk(
            &picked,
            &which,
            zscale,
            coarsen.max(1),
            &out_dir.join("spacetime.vtk"),
        );
    }
}