terrio 0.2.0

//! E57 (ASTM E2807): laser scans, read as PDAL's `readers.e57` reads them and written with the
//! pure-Rust [`e57`] crate.
//!
//! The points of every scan are read, in order, with each scan's pose applied. The dimensions
//! are PDAL's: `x`, `y`, `z` (from spherical coordinates when the scans have only those),
//! `spherical_range`, `spherical_azimuth`, `spherical_elevation`, `normal_x`, `normal_y`,
//! `normal_z`, `red`, `green`, `blue`, `intensity` and `classification`. As PDAL does, colours and
//! intensities become 16-bit values and classifications 8-bit ones, scaled by `max / (limit max
//! - limit min)` (the scan's limits, or its prototype's), so that colours of 0 to 255 become 0 to
//! 65535.

use super::crs::LasCrs;
use super::las::LasHeader;
use super::read::{FormatReader, Source, apply_options, made_up_header, stride_for};
use super::values::{Dimension, ValueType, Values};
use super::{PointCloud, PointFilter, PointReadOptions, PointWriteOptions};
use crate::{Error, Result};
use e57::{RecordDataType, RecordName, RecordValue};
use std::io::{Read, Seek, Write};

/// The E57 fields read, PDAL's order, with terrio's name and type.
const FIELDS: &[(&str, &str, ValueType)] = &[
    ("cartesianX", "x", ValueType::F64),
    ("cartesianY", "y", ValueType::F64),
    ("cartesianZ", "z", ValueType::F64),
    ("sphericalRange", "spherical_range", ValueType::F64),
    ("sphericalAzimuth", "spherical_azimuth", ValueType::F64),
    ("sphericalElevation", "spherical_elevation", ValueType::F64),
    ("normalX", "normal_x", ValueType::F64),
    ("normalY", "normal_y", ValueType::F64),
    ("normalZ", "normal_z", ValueType::F64),
    ("colorRed", "red", ValueType::U16),
    ("colorGreen", "green", ValueType::U16),
    ("colorBlue", "blue", ValueType::U16),
    ("intensity", "intensity", ValueType::U16),
    ("classification", "classification", ValueType::U8),
];

/// The namespace terrio writes the dimensions E57 has no field for under.
const NAMESPACE: &str = "terrio";

/// A record's field name, without its namespace for the normals and terrio's own fields.
fn field_name(name: &RecordName) -> String {
    match name {
        RecordName::CartesianX => "cartesianX".into(),
        RecordName::CartesianY => "cartesianY".into(),
        RecordName::CartesianZ => "cartesianZ".into(),
        RecordName::SphericalRange => "sphericalRange".into(),
        RecordName::SphericalAzimuth => "sphericalAzimuth".into(),
        RecordName::SphericalElevation => "sphericalElevation".into(),
        RecordName::Intensity => "intensity".into(),
        RecordName::ColorRed => "colorRed".into(),
        RecordName::ColorGreen => "colorGreen".into(),
        RecordName::ColorBlue => "colorBlue".into(),
        RecordName::Unknown { namespace, name } if namespace.is_empty() || namespace == "nor" || namespace == NAMESPACE => name.clone(),
        RecordName::Unknown { namespace, name } => format!("{namespace}:{name}"),
        other => format!("{other:?}"),
    }
}

/// A limit value as a float.
fn limit(v: &Option<RecordValue>, dt: Option<&RecordDataType>) -> Option<f64> {
    match v.as_ref()? {
        RecordValue::Single(s) => Some(*s as f64),
        RecordValue::Double(d) => Some(*d),
        RecordValue::Integer(i) => Some(*i as f64),
        RecordValue::ScaledInteger(i) => match dt {
            Some(RecordDataType::ScaledInteger { scale, offset, .. }) => Some(*i as f64 * scale + offset),
            _ => Some(*i as f64),
        },
    }
}

/// The minimum and maximum of a prototype record's type.
fn type_limits(dt: &RecordDataType) -> Option<(f64, f64)> {
    match dt {
        RecordDataType::Integer { min, max } => Some((*min as f64, *max as f64)),
        RecordDataType::ScaledInteger { min, max, scale, offset } => Some((*min as f64 * scale + offset, *max as f64 * scale + offset)),
        RecordDataType::Single { min: Some(a), max: Some(b) } => Some((*a as f64, *b as f64)),
        RecordDataType::Double { min: Some(a), max: Some(b) } => Some((*a, *b)),
        _ => None,
    }
}

/// How to read one scan.
struct ScanPlan {
    /// For each prototype record, the column it goes to and the factor it is scaled by.
    targets: Vec<Option<(usize, f64)>>,
    rotation: [[f64; 3]; 3],
    translation: [f64; 3],
    pose: bool,
    spherical: bool,
}

/// An E57 file.
pub(crate) struct E57File {
    reader: e57::E57Reader<Source>,
    scans: Vec<e57::PointCloud>,
    header: LasHeader,
    crs: LasCrs,
    dims: Vec<(String, ValueType)>,
    cursor: Option<E57Cursor>,
    _local: Option<crate::vfs::LocalPath>,
}

struct E57Cursor {
    scan: usize,
    done: u64,
    stride: u64,
    filter: PointFilter,
}

impl E57File {
    pub fn new(source: Source, local: Option<crate::vfs::LocalPath>) -> Result<Self> {
        let reader = e57::E57Reader::new(source).map_err(|e| Error::InvalidData(format!("E57: {e}")))?;
        let scans = reader.pointclouds();
        // The dimensions any scan has, in PDAL's order; x, y and z first, also when they come
        // from spherical coordinates.
        let has = |f: &str| scans.iter().any(|s| s.prototype.iter().any(|r| field_name(&r.name) == f));
        let spherical = ["sphericalRange", "sphericalAzimuth", "sphericalElevation"].iter().all(|f| has(f));
        let mut dims: Vec<(String, ValueType)> = Vec::new();
        for (f, name, t) in FIELDS {
            let xyz = matches!(*f, "cartesianX" | "cartesianY" | "cartesianZ");
            if has(f) || (xyz && spherical) {
                dims.push((name.to_string(), *t));
            }
        }
        // terrio's own fields, written by write_e57.
        for s in &scans {
            for r in &s.prototype {
                if let RecordName::Unknown { namespace, name } = &r.name
                    && namespace == NAMESPACE
                    && !dims.iter().any(|(d, _)| d == name)
                {
                    dims.push((name.clone(), ValueType::F64));
                }
            }
        }
        // The bounds: the scans' Cartesian bounds, posed (all eight corners), when they all have
        // them.
        let mut bounds: Option<[f64; 6]> = None;
        let mut known = !scans.is_empty();
        for s in &scans {
            let Some(b) = s.cartesian_bounds.as_ref() else {
                known = false;
                break;
            };
            let (Some(x0), Some(x1), Some(y0), Some(y1), Some(z0), Some(z1)) = (b.x_min, b.x_max, b.y_min, b.y_max, b.z_min, b.z_max)
            else {
                known = false;
                break;
            };
            let plan = plan(s, &dims);
            for &x in &[x0, x1] {
                for &y in &[y0, y1] {
                    for &z in &[z0, z1] {
                        let p = if plan.pose { transform(&plan, [x, y, z]) } else { [x, y, z] };
                        bounds = Some(match bounds {
                            None => [p[0], p[1], p[2], p[0], p[1], p[2]],
                            Some(a) => [a[0].min(p[0]), a[1].min(p[1]), a[2].min(p[2]), a[3].max(p[0]), a[4].max(p[1]), a[5].max(p[2])],
                        });
                    }
                }
            }
        }
        let count = scans.iter().map(|s| s.records).sum();
        let mut f = Self { reader, scans, header: LasHeader::default(), crs: LasCrs::default(), dims, cursor: None, _local: local };
        if !known {
            bounds = None;
            f.start(&PointReadOptions::default(), 0)?;
            while let Some(chunk) = f.next_chunk()? {
                if let Some(b) = PointCloud::new(chunk, Default::default()).bounds() {
                    bounds = Some(match bounds {
                        None => b,
                        Some(a) => [a[0].min(b[0]), a[1].min(b[1]), a[2].min(b[2]), a[3].max(b[3]), a[4].max(b[4]), a[5].max(b[5])],
                    });
                }
            }
        }
        f.header = made_up_header(&f.dims, count, bounds, &Default::default(), None);
        Ok(f)
    }
}

/// How to read a scan: where each field goes, PDAL's scale factors, and the pose.
fn plan(scan: &e57::PointCloud, dims: &[(String, ValueType)]) -> ScanPlan {
    let column = |name: &str| dims.iter().position(|(d, _)| d == name);
    let targets = scan
        .prototype
        .iter()
        .map(|r| {
            let f = field_name(&r.name);
            let (name, t) = match FIELDS.iter().find(|(e, ..)| *e == f) {
                Some((_, n, t)) => (n.to_string(), *t),
                None => match &r.name {
                    RecordName::Unknown { namespace, name } if namespace == NAMESPACE => (name.clone(), ValueType::F64),
                    _ => return None,
                },
            };
            let col = column(&name)?;
            // PDAL scales colours, intensities and classifications to the full range of their
            // type, by a factor computed in single precision.
            let scaled_max = match (f.as_str(), t) {
                ("colorRed" | "colorGreen" | "colorBlue" | "intensity", _) => Some(65535.0),
                ("classification", _) => Some(255.0),
                _ => None,
            };
            let factor = match scaled_max {
                None => 1.0,
                Some(top) => {
                    // Scaled-integer limits take the scale of the field's.
                    let dt = Some(&r.data_type);
                    let limits = match f.as_str() {
                        "colorRed" => scan.color_limits.as_ref().map(|l| (limit(&l.red_min, dt), limit(&l.red_max, dt))),
                        "colorGreen" => scan.color_limits.as_ref().map(|l| (limit(&l.green_min, dt), limit(&l.green_max, dt))),
                        "colorBlue" => scan.color_limits.as_ref().map(|l| (limit(&l.blue_min, dt), limit(&l.blue_max, dt))),
                        "intensity" => scan.intensity_limits.as_ref().map(|l| (limit(&l.intensity_min, dt), limit(&l.intensity_max, dt))),
                        _ => None,
                    };
                    let range = match limits {
                        Some((Some(a), Some(b))) => Some((a, b)),
                        _ => type_limits(&r.data_type),
                    };
                    match range {
                        // Equal limits give PDAL an infinite factor, and values it cannot store.
                        Some((a, b)) => (top / (b - a)) as f32 as f64,
                        None => 1.0,
                    }
                }
            };
            Some((col, factor))
        })
        .collect();
    let has = |f: &str| scan.prototype.iter().any(|r| field_name(&r.name) == f);
    let spherical = ["sphericalRange", "sphericalAzimuth", "sphericalElevation"].iter().all(|f| has(f));
    let (rotation, translation, pose) = match &scan.transform {
        Some(t) => {
            let q = [t.rotation.w, t.rotation.x, t.rotation.y, t.rotation.z];
            let (q11, q22, q33) = (q[1] * q[1], q[2] * q[2], q[3] * q[3]);
            let (q03, q13, q23, q02, q12, q01) = (q[0] * q[3], q[1] * q[3], q[2] * q[3], q[0] * q[2], q[1] * q[2], q[0] * q[1]);
            (
                [
                    [1.0 - 2.0 * (q22 + q33), 2.0 * (q12 - q03), 2.0 * (q13 + q02)],
                    [2.0 * (q12 + q03), 1.0 - 2.0 * (q11 + q33), 2.0 * (q23 - q01)],
                    [2.0 * (q13 - q02), 2.0 * (q23 + q01), 1.0 - 2.0 * (q11 + q22)],
                ],
                [t.translation.x, t.translation.y, t.translation.z],
                true,
            )
        }
        None => ([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [0.0; 3], false),
    };
    ScanPlan { targets, rotation, translation, pose, spherical }
}

/// A point posed as PDAL poses it.
fn transform(plan: &ScanPlan, p: [f64; 3]) -> [f64; 3] {
    let (r, t) = (&plan.rotation, &plan.translation);
    [
        p[0] * r[0][0] + p[1] * r[0][1] + p[2] * r[0][2] + t[0],
        p[0] * r[1][0] + p[1] * r[1][1] + p[2] * r[1][2] + t[1],
        p[0] * r[2][0] + p[1] * r[2][1] + p[2] * r[2][2] + t[2],
    ]
}

impl FormatReader for E57File {
    fn header(&self) -> &LasHeader {
        &self.header
    }

    fn crs(&self) -> &LasCrs {
        &self.crs
    }

    fn dimensions(&self) -> Vec<(String, ValueType)> {
        self.dims.clone()
    }

    fn format_name(&self) -> &'static str {
        "E57"
    }

    fn start(&mut self, options: &PointReadOptions, _chunk: usize) -> Result<()> {
        let n = self.scans.iter().map(|s| s.records).sum();
        self.cursor = Some(E57Cursor { scan: 0, done: 0, stride: stride_for(n, options.max_points), filter: options.filter.clone() });
        Ok(())
    }

    /// The points of the next scan.
    fn next_chunk(&mut self) -> Result<Option<Vec<Dimension>>> {
        let Some(mut c) = self.cursor.take() else { return Ok(None) };
        if c.scan >= self.scans.len() {
            return Ok(None);
        }
        let scan = self.scans[c.scan].clone();
        let plan = plan(&scan, &self.dims);
        let m = scan.records as usize;
        let mut cols: Vec<Vec<f64>> = vec![Vec::with_capacity(m); self.dims.len()];
        let col = |name: &str| self.dims.iter().position(|(d, _)| d == name);
        let (cx, cy, cz) = (col("x"), col("y"), col("z"));
        let (cr, ca, ce) = (col("spherical_range"), col("spherical_azimuth"), col("spherical_elevation"));
        let mut row = vec![0.0; self.dims.len()];
        let raw = self.reader.pointcloud_raw(&scan).map_err(|e| Error::InvalidData(format!("E57: {e}")))?;
        for values in raw {
            let values = values.map_err(|e| Error::InvalidData(format!("E57: {e}")))?;
            row.iter_mut().for_each(|v| *v = 0.0);
            for ((v, r), target) in values.iter().zip(&scan.prototype).zip(&plan.targets) {
                if let Some((k, factor)) = target {
                    let x = v.to_f64(&r.data_type).map_err(|e| Error::InvalidData(format!("E57: {e}")))?;
                    row[*k] = factor * x;
                }
            }
            if plan.spherical
                && let (Some(r), Some(a), Some(e), Some(x), Some(y), Some(z)) = (cr, ca, ce, cx, cy, cz)
            {
                let (range, az, el) = (row[r], row[a], row[e]);
                row[x] = range * el.cos() * az.cos();
                row[y] = range * el.cos() * az.sin();
                row[z] = range * el.sin();
            }
            if plan.pose
                && let (Some(x), Some(y), Some(z)) = (cx, cy, cz)
            {
                let p = transform(&plan, [row[x], row[y], row[z]]);
                row[x] = p[0];
                row[y] = p[1];
                row[z] = p[2];
            }
            for (c, v) in cols.iter_mut().zip(&row) {
                c.push(*v);
            }
        }
        // The values in their types, rounded as PDAL rounds them.
        let dims: Vec<Dimension> = self
            .dims
            .iter()
            .zip(cols)
            .map(|((name, t), v)| {
                let values = match t {
                    ValueType::F64 => Values::F64(v),
                    t => {
                        // PDAL rounds, and leaves 0 for values out of the type's range.
                        let max = if *t == ValueType::U8 { 255.0 } else { 65535.0 };
                        let mut out = Values::with_capacity(*t, v.len());
                        for x in v {
                            let r = x.round();
                            out.push_f64(if (0.0..=max).contains(&r) { r } else { 0.0 });
                        }
                        out
                    }
                };
                Dimension::new(name, values)
            })
            .collect();
        let first = c.done;
        c.done += m as u64;
        c.scan += 1;
        let out = apply_options(dims, &c.filter, c.stride, first);
        self.cursor = Some(c);
        Ok(Some(out))
    }
}

/// A GUID made of the time and a counter (E57 files and scans need one).
fn guid() -> String {
    use std::sync::atomic::{AtomicU64, Ordering};
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    // The clock is not available in WebAssembly builds without a host.
    #[cfg(not(target_arch = "wasm32"))]
    let t = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_nanos()).unwrap_or(0) as u64;
    #[cfg(target_arch = "wasm32")]
    let t = 0x5eed_u64;
    let c = COUNTER.fetch_add(1, Ordering::Relaxed);
    let a = super::copc::scramble(t ^ c.rotate_left(32));
    let b = super::copc::scramble(a ^ 0x9e37_79b9_7f4a_7c15);
    format!("{{{:08x}-{:04x}-{:04x}-{:04x}-{:012x}}}", a >> 32, (a >> 16) & 0xffff, a & 0xffff, b >> 48, b & 0xffff_ffff_ffff)
}

/// Writes a point cloud as an E57 file of one scan: `x`, `y` and `z` as doubles, colours and
/// intensities as integers of 0 to 65535 (with limits of 0 to 65535, so that PDAL reads them
/// unchanged), normals in the `nor` extension, and the other dimensions as doubles in terrio's
/// own extension.
pub(crate) fn write_e57<T: Read + Write + Seek>(cloud: &PointCloud, out: T, options: &PointWriteOptions) -> Result<()> {
    let err = |e: e57::Error| Error::InvalidData(format!("E57: {e}"));
    let pick: Vec<&Dimension> = if options.dimensions.is_empty() {
        cloud.dimensions.iter().collect()
    } else {
        options
            .dimensions
            .iter()
            .map(|n| cloud.dim(n).ok_or_else(|| Error::InvalidArgument(format!("no dimension '{n}'"))))
            .collect::<Result<_>>()?
    };
    for c in ["x", "y", "z"] {
        if !pick.iter().any(|d| d.name == c) {
            return Err(Error::InvalidArgument(format!("E57 needs '{c}'")));
        }
    }
    // E57 has colours as all of red, green and blue, or none.
    let colours = ["red", "green", "blue"].iter().all(|c| pick.iter().any(|d| d.name == *c));
    let mut w = e57::E57Writer::new(out, &guid()).map_err(err)?;
    let mut prototype = Vec::new();
    let mut sources = Vec::new();
    let mut normals = false;
    let mut own = false;
    let int16 = RecordDataType::Integer { min: 0, max: 65535 };
    for d in &pick {
        let name = match d.name.as_str() {
            "x" => RecordName::CartesianX,
            "y" => RecordName::CartesianY,
            "z" => RecordName::CartesianZ,
            "spherical_range" => RecordName::SphericalRange,
            "spherical_azimuth" => RecordName::SphericalAzimuth,
            "spherical_elevation" => RecordName::SphericalElevation,
            "red" if colours => RecordName::ColorRed,
            "green" if colours => RecordName::ColorGreen,
            "blue" if colours => RecordName::ColorBlue,
            "intensity" => RecordName::Intensity,
            "normal_x" | "normal_y" | "normal_z" => {
                normals = true;
                let n = match d.name.as_str() {
                    "normal_x" => "normalX",
                    "normal_y" => "normalY",
                    _ => "normalZ",
                };
                RecordName::Unknown { namespace: "nor".into(), name: n.into() }
            }
            other => {
                if !other.chars().all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
                    || other.to_ascii_lowercase().starts_with("xml")
                {
                    continue;
                }
                own = true;
                RecordName::Unknown { namespace: NAMESPACE.into(), name: other.into() }
            }
        };
        let data_type = match name {
            RecordName::ColorRed | RecordName::ColorGreen | RecordName::ColorBlue | RecordName::Intensity => int16.clone(),
            _ => RecordDataType::Double { min: None, max: None },
        };
        prototype.push(e57::Record { name, data_type });
        sources.push(*d);
    }
    if normals {
        w.register_extension(e57::Extension { namespace: "nor".into(), url: "http://www.libe57.org/E57_NOR_surface_normals.txt".into() })
            .map_err(err)?;
    }
    if own {
        w.register_extension(e57::Extension { namespace: NAMESPACE.into(), url: "https://github.com/opengeos/terrio".into() })
            .map_err(err)?;
    }
    {
        let mut pc = w.add_pointcloud(&guid(), prototype.clone()).map_err(err)?;
        for i in 0..cloud.len() {
            let values = prototype
                .iter()
                .zip(&sources)
                .map(|(r, d)| match r.data_type {
                    RecordDataType::Integer { .. } => RecordValue::Integer((d.get(i).round() as i64).clamp(0, 65535)),
                    _ => RecordValue::Double(d.get(i)),
                })
                .collect();
            pc.add_point(values).map_err(err)?;
        }
        pc.finalize().map_err(err)?;
    }
    w.finalize().map_err(err)?;
    Ok(())
}