medh5 2.0.1

//! §10 --- transforms and the frame graph (E5xx).

use std::collections::BTreeMap;

use indexmap::IndexMap;

use super::{
    child, children, float_list, grid_spatial, has_member, kinds_known, loc, read_f64, str_attr, strs_attr,
    sub_dataset, Context,
};
use crate::h5::attrs;
use crate::h5::ops::{self, Node};
use crate::json::{repr_int_tuple, repr_list, repr_str};
use crate::transforms::model::{affine_dimension, composite_cycles, DEFAULT_ORDER, TRANSFORM_KINDS, VECTOR_SPACES};
use crate::validate::Diagnostic;
use crate::Result;

/// Semantic: every transform against §10 --- its kind, frames and parameters
/// (E502), its units (E506), an affine's matrix (E504), a field's or a
/// lattice's grid (E101, E503), a composite's chain (E501), and each declared
/// inverse (E505).  A kind or vector space only a later minor defines is W913
/// in its files (1.1 §2.2).
pub fn check_transforms(ctx: &mut Context) -> Result<Vec<Diagnostic>> {
    let mut out = Vec::new();
    let Some(node) = ops::child_group(&ctx.root, "transforms") else { return Ok(out) };
    let grids = ops::child_group(&ctx.root, "grids");
    let transforms = children(&node)?;
    // Every composite's components, read once: a cycle is found for all of
    // them in one walk, where rebuilding the transforms and walking from each
    // composite in turn was quadratic in the number of composites.
    let mut graph: IndexMap<String, Vec<String>> = IndexMap::new();
    for (name, t) in &transforms {
        if str_attr(loc(t), "kind")?.as_deref() == Some("composite") {
            graph.insert(name.clone(), strs_attr(loc(t), "components")?);
        }
    }
    let cycles = composite_cycles(&graph)?;
    let mut declared: BTreeMap<String, (String, String)> = BTreeMap::new();
    for (name, t) in transforms {
        let location = format!("/transforms/{name}");
        let a = loc(&t);
        let Some(kind) = str_attr(a, "kind")? else {
            out.push(ctx.err("E502", location, "transform has no `kind` attribute"));
            continue;
        };
        if !TRANSFORM_KINDS.contains(&kind.as_str()) {
            out.push(ctx.unknown(
                "E502",
                location,
                format!("unknown transform kind {}; expected one of {}", repr_str(&kind), repr_list(&TRANSFORM_KINDS)),
            ));
            continue;
        }
        let missing: Vec<&str> = ["from_frame", "to_frame"].into_iter().filter(|k| !attrs::has(a, k)).collect();
        if !missing.is_empty() {
            out.push(ctx.err("E502", location, format!("missing {}", repr_list(&missing))));
            continue;
        }
        let source = str_attr(a, "from_frame")?.unwrap_or_default();
        let target = str_attr(a, "to_frame")?.unwrap_or_default();
        declared.insert(name.clone(), (source.clone(), target.clone()));
        if source == target {
            out.push(ctx.err(
                "E502",
                location.clone(),
                format!("maps frame {} to itself; grids sharing a frame need no transform (§3.4)", repr_str(&source)),
            ));
        }
        out.extend(check_units(ctx, &name, &t, grids.as_ref(), &source, &target)?);
        match kind.as_str() {
            "affine" => out.extend(check_affine(ctx, &name, &t, grids.as_ref(), &source, &target)?),
            "displacement" | "bspline" => {
                out.extend(check_field_transform(ctx, &name, &t, &kind, &source, grids.as_ref())?)
            }
            "composite" => out.extend(check_composite(ctx, &name, &t, &node, &source, &target, cycles.get(&name))?),
            _ => {}
        }
    }
    out.extend(check_inverses(ctx, &node, &declared)?);
    Ok(out)
}

/// `units` is present, and the units of every grid in either frame and of
/// `from_grid` and `to_grid` (§10.1, E506): a transform maps coordinates in
/// one unit, and none maps between two.
fn check_units(
    ctx: &Context,
    name: &str,
    group: &Node,
    grids: Option<&hdf5::Group>,
    source: &str,
    target: &str,
) -> Result<Vec<Diagnostic>> {
    let location = format!("/transforms/{name}");
    let a = loc(group);
    let Some(units) = str_attr(a, "units")? else {
        return Ok(vec![ctx.err("E506", location, "transform has no `units`; §10.1 requires the units it maps")]);
    };
    let Some(grids) = grids else { return Ok(Vec::new()) };
    let mut wrong = Vec::new();
    for (grid_id, g) in related_grids(group, grids, source, target)? {
        let found = str_attr(loc(&g), "units")?.unwrap_or_else(|| "mm".into());
        if found != units {
            wrong.push(format!("grid {} is in {}", repr_str(&grid_id), repr_str(&found)));
        }
    }
    if wrong.is_empty() {
        return Ok(Vec::new());
    }
    Ok(vec![ctx.err(
        "E506",
        location,
        format!(
            "is in {}, and {}: a transform maps coordinates in one unit, and none maps between two (§10.1)",
            repr_str(&units),
            wrong.join(", ")
        ),
    )])
}

/// The grids a transform relates, by id: those in either of its frames, and
/// `from_grid` and `to_grid` (§10.1).
fn related_grids(group: &Node, grids: &hdf5::Group, source: &str, target: &str) -> Result<Vec<(String, Node)>> {
    let mut named = Vec::new();
    for key in ["from_grid", "to_grid"] {
        named.extend(str_attr(loc(group), key)?);
    }
    let mut out = Vec::new();
    for (grid_id, g) in children(grids)? {
        let frame = str_attr(loc(&g), "frame_uid")?;
        if frame.as_deref().is_some_and(|f| f == source || f == target) || named.contains(&grid_id) {
            out.push((grid_id, g));
        }
    }
    Ok(out)
}

/// `matrix` is `(S+1, S+1)` for `S` the spatial axes of every grid the
/// affine relates, and its last row is `[0 … 0 1]` (§10.3).
fn check_affine(
    ctx: &Context,
    name: &str,
    group: &Node,
    grids: Option<&hdf5::Group>,
    source: &str,
    target: &str,
) -> Result<Vec<Diagnostic>> {
    let location = format!("/transforms/{name}");
    let Some(ds) = sub_dataset(group, "matrix") else {
        return Ok(vec![ctx.err("E502", location, "kind 'affine' requires a `matrix` dataset")]);
    };
    let s = match affine_dimension(&ds.shape()) {
        Ok(s) => s,
        Err(e) => return Ok(vec![ctx.err("E504", location, e.message())]),
    };
    let m = read_f64(&ds)?;
    let last: Vec<f64> = (0..=s).map(|c| m[[s, c]]).collect();
    let mut expected = vec![0.0; s + 1];
    expected[s] = 1.0;
    if !crate::geometry::linalg::allclose(&last, &expected, 1e-9, 1e-5) {
        return Ok(vec![ctx.err(
            "E504",
            location,
            format!("last row must be [0 \u{2026} 0 1], got {}", float_list(&last)),
        )]);
    }
    let Some(grids) = grids else { return Ok(Vec::new()) };
    let mut wrong = Vec::new();
    for (grid_id, g) in related_grids(group, grids, source, target)? {
        // A grid naming an axis kind this validator does not know has no
        // count of spatial axes to hold the matrix to; its own finding stands.
        if !kinds_known(grids, &grid_id)? {
            continue;
        }
        let n_spatial = strs_attr(loc(&g), "axis_kinds")?.iter().filter(|k| *k == "spatial").count();
        if n_spatial != s {
            wrong.push(format!("grid {} has {n_spatial} spatial axes", repr_str(&grid_id)));
        }
    }
    if wrong.is_empty() {
        return Ok(Vec::new());
    }
    Ok(vec![ctx.err(
        "E504",
        location,
        format!(
            "`matrix` is ({0}, {0}), an affine of {s} spatial axes, and {1}: an affine maps a point of the grids \
             it relates, one coordinate per spatial axis (§10.3)",
            s + 1,
            wrong.join(", ")
        ),
    )])
}

fn check_field_transform(
    ctx: &Context,
    name: &str,
    group: &Node,
    kind: &str,
    from_frame: &str,
    grids: Option<&hdf5::Group>,
) -> Result<Vec<Diagnostic>> {
    let mut out = Vec::new();
    let location = format!("/transforms/{name}");
    let (dataset, grid_attr) =
        if kind == "displacement" { ("field", "field_grid") } else { ("control_points", "cp_grid") };
    if !has_member(group, dataset) {
        out.push(ctx.err(
            "E502",
            location.clone(),
            format!("kind {} requires a {} dataset", repr_str(kind), repr_str(dataset)),
        ));
    }
    let Some(grid_id) = str_attr(loc(group), grid_attr)? else {
        out.push(ctx.err("E503", location, format!("kind {} requires {}", repr_str(kind), repr_str(grid_attr))));
        return Ok(out);
    };
    let Some(grids) = grids.filter(|g| ops::exists(g, &grid_id)) else {
        out.push(ctx.err("E101", location, format!("{grid_attr} {} does not exist", repr_str(&grid_id))));
        return Ok(out);
    };
    let g = grids.group(&grid_id)?;
    if let Some(frame) = str_attr(&g, "frame_uid")? {
        if frame != from_frame {
            out.push(ctx.err(
                "E503",
                location.clone(),
                format!(
                    "{grid_attr} {} is in frame {} but the transform starts in {}; the field is sampled in the source frame",
                    repr_str(&grid_id),
                    repr_str(&frame),
                    repr_str(from_frame)
                ),
            ));
        }
    }
    let space = str_attr(loc(group), "vector_space")?.unwrap_or_else(|| "world".into());
    if !VECTOR_SPACES.contains(&space.as_str()) {
        out.push(ctx.unknown("E502", location.clone(), format!("unknown vector_space {}", repr_str(&space))));
    }
    // The field and the lattice are measured in the grid's spatial axes,
    // which a kind this validator does not know leaves uncounted.
    if let (Some(ds), true) = (sub_dataset(group, dataset), kinds_known(grids, &grid_id)?) {
        let shape = ds.shape();
        let kinds = strs_attr(&g, "axis_kinds")?;
        let n_spatial = kinds.iter().filter(|k| *k == "spatial").count();
        let components = shape.first().copied().unwrap_or(0);
        if components != n_spatial {
            out.push(ctx.err(
                "E503",
                format!("{location}/{dataset}"),
                format!("{components} components on a {n_spatial}-D lattice; they must match"),
            ));
        }
        let lattice: Vec<i64> = shape.get(1..).unwrap_or(&[]).iter().map(|v| *v as i64).collect();
        if kind == "displacement" {
            let spatial = grid_spatial(grids, &grid_id)?;
            if lattice != spatial {
                out.push(ctx.err(
                    "E503",
                    format!("{location}/{dataset}"),
                    format!(
                        "field lattice {} != grid {} spatial shape {}",
                        repr_int_tuple(&lattice),
                        repr_str(&grid_id),
                        repr_int_tuple(&spatial)
                    ),
                ));
            }
        } else {
            // What the writer refuses (§10.5): a lattice of another rank than
            // its grid's, or one too small for the spline's order.
            // Saturating: `order` is whatever the file holds, and `order + 1`
            // overflowed at the largest.
            let order = attrs::get_i64(loc(group), "order")?.unwrap_or(DEFAULT_ORDER);
            let needed = order.saturating_add(1);
            if lattice.len() != n_spatial {
                out.push(ctx.err(
                    "E503",
                    format!("{location}/{dataset}"),
                    format!(
                        "control-point lattice {} has {} axes on a {n_spatial}-D grid",
                        repr_int_tuple(&lattice),
                        lattice.len()
                    ),
                ));
            } else if lattice.iter().any(|extent| *extent < needed) {
                out.push(ctx.err(
                    "E503",
                    format!("{location}/{dataset}"),
                    format!(
                        "an order-{order} B-spline needs at least {needed} control points per axis, got {}",
                        repr_int_tuple(&lattice)
                    ),
                ));
            }
        }
    }
    Ok(out)
}

fn check_composite(
    ctx: &Context,
    name: &str,
    group: &Node,
    node: &hdf5::Group,
    source: &str,
    target: &str,
    cycle: Option<&Vec<String>>,
) -> Result<Vec<Diagnostic>> {
    let mut out = Vec::new();
    let location = format!("/transforms/{name}");
    let a = loc(group);
    if !attrs::has(a, "components") {
        out.push(ctx.err("E501", location, "kind 'composite' requires `components`"));
        return Ok(out);
    }
    let components = strs_attr(a, "components")?;
    let unknown: Vec<&String> = components.iter().filter(|c| !ops::exists(node, c)).collect();
    if !unknown.is_empty() {
        out.push(ctx.err("E501", location, format!("names components {} that do not exist", repr_list(&unknown))));
        return Ok(out);
    }
    if let Some(cycle) = cycle {
        out.push(ctx.err(
            "E501",
            location,
            format!("contains itself through {}; a chain that never ends cannot be evaluated", cycle.join(" -> ")),
        ));
        return Ok(out);
    }
    let mut frames = Vec::new();
    for component in &components {
        let c = child(node, component).ok_or_else(|| crate::Error::Key(repr_str(component)))?;
        let (f, t) = (str_attr(loc(&c), "from_frame")?, str_attr(loc(&c), "to_frame")?);
        let (Some(f), Some(t)) = (f, t) else {
            out.push(ctx.err("E501", location, format!("component {} declares no frames", repr_str(component))));
            return Ok(out);
        };
        frames.push((f, t));
    }
    if let Some(first) = frames.first() {
        if first.0 != source {
            out.push(ctx.err(
                "E501",
                location.clone(),
                format!(
                    "first component starts in {} but the composite declares {}",
                    repr_str(&first.0),
                    repr_str(source)
                ),
            ));
        }
    }
    if let Some(last) = frames.last() {
        if last.1 != target {
            out.push(ctx.err(
                "E501",
                location.clone(),
                format!("last component ends in {} but the composite declares {}", repr_str(&last.1), repr_str(target)),
            ));
        }
    }
    for i in 0..frames.len().saturating_sub(1) {
        if frames[i].1 != frames[i + 1].0 {
            out.push(ctx.err(
                "E501",
                location.clone(),
                format!(
                    "{} ends in {} but {} starts in {}",
                    repr_str(&components[i]),
                    repr_str(&frames[i].1),
                    repr_str(&components[i + 1]),
                    repr_str(&frames[i + 1].0)
                ),
            ));
        }
    }
    if let Some(declared_units) = str_attr(a, "units")? {
        let mut mixed = Vec::new();
        for component in &components {
            let c = child(node, component).ok_or_else(|| crate::Error::Key(repr_str(component)))?;
            if let Some(u) = str_attr(loc(&c), "units")? {
                if u != declared_units {
                    mixed.push(format!("{} in {}", repr_str(component), repr_str(&u)));
                }
            }
        }
        if !mixed.is_empty() {
            out.push(ctx.err(
                "E501",
                location,
                format!(
                    "declares units {} but {} --- a chain whose legs are in different units does not compose",
                    repr_str(&declared_units),
                    mixed.join(", ")
                ),
            ));
        }
    }
    Ok(out)
}

fn check_inverses(
    ctx: &Context,
    node: &hdf5::Group,
    declared: &BTreeMap<String, (String, String)>,
) -> Result<Vec<Diagnostic>> {
    let mut out = Vec::new();
    for (name, (source, target)) in declared {
        let t = child(node, name).ok_or_else(|| crate::Error::Key(repr_str(name)))?;
        let Some(other) = str_attr(loc(&t), "inverse_id")? else { continue };
        let location = format!("/transforms/{name}");
        let Some((os, ot)) = declared.get(&other) else {
            // A transform that exists but could not be judged --- a kind this
            // validator does not know, no frames --- has its own finding there
            // (E502, or W913 in a later minor's file); it is not missing.
            if child(node, &other).is_none() {
                out.push(ctx.err(
                    "E505",
                    location,
                    format!("`inverse_id` names {}, which does not exist", repr_str(&other)),
                ));
            }
            continue;
        };
        if (os, ot) != (target, source) {
            out.push(ctx.err(
                "E505",
                location,
                format!(
                    "`inverse_id` names {}, which maps {} -> {}; an inverse must map {} -> {}",
                    repr_str(&other),
                    repr_str(os),
                    repr_str(ot),
                    repr_str(target),
                    repr_str(source)
                ),
            ));
            continue;
        }
        let o = child(node, &other).ok_or_else(|| crate::Error::Key(repr_str(&other)))?;
        if let Some(back) = str_attr(loc(&o), "inverse_id")? {
            if &back != name {
                out.push(ctx.err(
                    "E505",
                    location,
                    format!(
                        "{} names {} as its inverse, not {}; the relation must be mutual",
                        repr_str(&other),
                        repr_str(&back),
                        repr_str(name)
                    ),
                ));
            }
        }
    }
    Ok(out)
}