bevy_map_scatter 0.5.0

#[cfg(feature = "ron")]
use std::collections::BTreeMap;
use std::result::Result;

use bevy::asset::io::Reader;
#[cfg(feature = "ron")]
use bevy::asset::io::{AsyncWriteExt, Writer};
#[cfg(feature = "ron")]
use bevy::asset::{saver::AssetSaver, saver::SavedAsset, AssetPath};
use bevy::asset::{AssetLoader, LoadContext, ReflectAsset};
use bevy::prelude::*;
use map_scatter::prelude::*;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

/// Asset describing a complete scatter [`Plan`] for `map_scatter`.
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Asset, Reflect, Clone, Debug)]
#[reflect(Asset)]
pub struct ScatterPlanAsset {
    /// Ordered list of layer definitions in the plan.
    pub layers: Vec<ScatterLayerDef>,
}

/// Layer definition within a [`ScatterPlanAsset`].
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, Debug, Reflect)]
pub struct ScatterLayerDef {
    /// Unique identifier for this layer.
    pub id: String,
    /// Kinds evaluated in this layer.
    pub kinds: Vec<ScatterKindDef>,
    /// Sampling strategy used for candidate generation.
    pub sampling: SamplingDef,
    /// Optional overlay mask size in pixels (width, height).
    pub overlay_mask_size_px: Option<(u32, u32)>,
    /// Optional overlay brush radius in pixels.
    pub overlay_brush_radius_px: Option<i32>,
    /// Strategy for selecting a kind when multiple are valid.
    pub selection_strategy: SelectionStrategyDef,
}

/// Kind definition.
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, Debug, Reflect)]
pub struct ScatterKindDef {
    /// Unique identifier for this kind.
    pub id: String,
    /// Field graph specification for this kind.
    ///
    /// This remains serialized through serde/RON, but is skipped by Bevy reflection because the
    /// core `map_scatter` crate intentionally does not depend on Bevy.
    #[reflect(ignore)]
    pub spec: FieldGraphSpec,
}

/// Selection strategy for layers.
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, Copy, Debug, Reflect)]
pub enum SelectionStrategyDef {
    WeightedRandom,
    HighestProbability,
}

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, Debug, Reflect)]
pub enum ParentDef {
    Count(
        /// Number of parent centers to generate.
        usize,
    ),
    Density(
        /// Parent density per unit area.
        f32,
    ),
}

#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Clone, Debug, Reflect)]
pub enum SamplingDef {
    UniformRandom {
        /// Number of candidate points to generate.
        count: usize,
    },
    Halton {
        /// Number of candidate points to generate.
        count: usize,
        /// Bases for the 2D Halton sequence.
        bases: (u32, u32),
        /// Starting index in the Halton sequence.
        start_index: u32,
        /// Apply Cranley-Patterson rotation.
        rotate: bool,
    },
    FibonacciLattice {
        /// Number of candidate points to generate.
        count: usize,
        /// Apply Cranley-Patterson rotation.
        rotate: bool,
    },
    StratifiedMultiJitter {
        /// Number of candidate points to generate.
        count: usize,
        /// Apply Cranley-Patterson rotation.
        rotate: bool,
    },
    BestCandidate {
        /// Number of candidate points to generate.
        count: usize,
        /// Trials per point for best-candidate selection.
        k: usize,
    },
    PoissonDisk {
        /// Minimum distance between points in world units.
        radius: f32,
    },
    JitterGrid {
        /// Jitter amount in [0, 1].
        jitter: f32,
        /// Cell size for the base grid in world units.
        cell_size: f32,
    },
    HexJitterGrid {
        /// Jitter amount in [0, 1].
        jitter: f32,
        /// Base spacing along X in world units.
        cell_size: f32,
    },
    ClusteredThomas {
        /// Parent placement configuration.
        parents: ParentDef,
        /// Mean number of children per parent.
        mean_children: f32,
        /// Gaussian sigma for child offsets.
        sigma: f32,
        /// Clamp children inside the domain bounds.
        clamp_inside: bool,
    },
    ClusteredNeymanScott {
        /// Parent placement configuration.
        parents: ParentDef,
        /// Mean number of children per parent.
        mean_children: f32,
        /// Disk radius for child offsets.
        radius: f32,
        /// Clamp children inside the domain bounds.
        clamp_inside: bool,
    },
}

impl From<&ScatterKindDef> for Kind {
    fn from(value: &ScatterKindDef) -> Self {
        Kind::new(value.id.clone(), value.spec.clone())
    }
}

impl From<ScatterKindDef> for Kind {
    fn from(value: ScatterKindDef) -> Self {
        Kind::new(value.id, value.spec)
    }
}

impl From<SelectionStrategyDef> for SelectionStrategy {
    fn from(value: SelectionStrategyDef) -> Self {
        match value {
            SelectionStrategyDef::WeightedRandom => SelectionStrategy::WeightedRandom,
            SelectionStrategyDef::HighestProbability => SelectionStrategy::HighestProbability,
        }
    }
}

impl From<&ScatterLayerDef> for Layer {
    fn from(def: &ScatterLayerDef) -> Self {
        let kinds: Vec<Kind> = def.kinds.iter().map(|k| k.into()).collect();
        let sampling: Box<dyn PositionSampling> = sampling_runtime(&def.sampling);
        let mut layer = Layer::new(def.id.clone(), kinds, sampling);

        if let (Some(size), Some(radius)) = (
            def.overlay_mask_size_px.as_ref(),
            def.overlay_brush_radius_px,
        ) {
            layer = layer.with_overlay(*size, radius);
        }

        layer.with_selection_strategy(def.selection_strategy.into())
    }
}

impl From<&ScatterPlanAsset> for Plan {
    fn from(asset: &ScatterPlanAsset) -> Self {
        let layers: Vec<Layer> = asset.layers.iter().map(|l| l.into()).collect();
        Plan::new().with_layers(layers)
    }
}

impl From<ScatterPlanAsset> for Plan {
    fn from(asset: ScatterPlanAsset) -> Self {
        (&asset).into()
    }
}

/// Convert a `SamplingDef` into a boxed runtime sampler.
fn sampling_runtime(def: &SamplingDef) -> Box<dyn PositionSampling> {
    match def {
        SamplingDef::UniformRandom { count } => Box::new(UniformRandomSampling { count: *count }),
        SamplingDef::Halton {
            count,
            bases,
            start_index,
            rotate,
        } => Box::new(HaltonSampling {
            count: *count,
            bases: *bases,
            start_index: *start_index,
            rotate: *rotate,
        }),
        SamplingDef::FibonacciLattice { count, rotate } => Box::new(FibonacciLatticeSampling {
            count: *count,
            rotate: *rotate,
        }),
        SamplingDef::StratifiedMultiJitter { count, rotate } => {
            Box::new(StratifiedMultiJitterSampling {
                count: *count,
                rotate: *rotate,
            })
        }
        SamplingDef::BestCandidate { count, k } => Box::new(BestCandidateSampling {
            count: *count,
            k: *k,
        }),
        SamplingDef::PoissonDisk { radius } => Box::new(PoissonDiskSampling { radius: *radius }),
        SamplingDef::JitterGrid { jitter, cell_size } => {
            Box::new(JitterGridSampling::new(*jitter, *cell_size))
        }
        SamplingDef::HexJitterGrid { jitter, cell_size } => {
            Box::new(HexJitterGridSampling::new(*jitter, *cell_size))
        }
        SamplingDef::ClusteredThomas {
            parents,
            mean_children,
            sigma,
            clamp_inside,
        } => {
            let base = match parents {
                ParentDef::Count(n) => {
                    ClusteredSampling::thomas_with_count(*n, *mean_children, *sigma)
                }
                ParentDef::Density(d) => {
                    ClusteredSampling::thomas_with_density(*d, *mean_children, *sigma)
                }
            };
            Box::new(base.with_clamp_inside(*clamp_inside))
        }
        SamplingDef::ClusteredNeymanScott {
            parents,
            mean_children,
            radius,
            clamp_inside,
        } => {
            let base = match parents {
                ParentDef::Count(n) => {
                    ClusteredSampling::neyman_scott_with_count(*n, *mean_children, *radius)
                }
                ParentDef::Density(d) => {
                    ClusteredSampling::neyman_scott_with_density(*d, *mean_children, *radius)
                }
            };
            Box::new(base.with_clamp_inside(*clamp_inside))
        }
    }
}

/// Asset loader for [`ScatterPlanAsset`] using RON files with `.scatter` extension.
#[derive(TypePath)]
pub struct ScatterPlanAssetLoader;

impl AssetLoader for ScatterPlanAssetLoader {
    type Asset = ScatterPlanAsset;
    type Settings = ();
    type Error = anyhow::Error;

    fn extensions(&self) -> &[&str] {
        &["scatter"]
    }

    async fn load(
        &self,
        reader: &mut dyn Reader,
        _settings: &Self::Settings,
        _context: &mut LoadContext<'_>,
    ) -> Result<Self::Asset, Self::Error> {
        let mut bytes = Vec::new();
        reader.read_to_end(&mut bytes).await?;
        #[cfg(feature = "ron")]
        {
            let asset: ScatterPlanAsset =
                ron::de::from_bytes(&bytes).map_err(|e| anyhow::anyhow!(e))?;
            Ok(asset)
        }
        #[cfg(not(feature = "ron"))]
        {
            let _ = bytes;
            Err(anyhow::anyhow!(
                "bevy_map_scatter: enable the `ron` feature to load .scatter assets"
            ))
        }
    }
}

impl FromWorld for ScatterPlanAssetLoader {
    fn from_world(_: &mut World) -> Self {
        ScatterPlanAssetLoader
    }
}

/// Asset saver for [`ScatterPlanAsset`] RON files with `.scatter` extension.
#[cfg(feature = "ron")]
#[derive(Default, TypePath)]
pub struct ScatterPlanAssetSaver;

#[cfg(feature = "ron")]
impl AssetSaver for ScatterPlanAssetSaver {
    type Asset = ScatterPlanAsset;
    type Settings = ();
    type OutputLoader = ScatterPlanAssetLoader;
    type Error = anyhow::Error;

    async fn save(
        &self,
        writer: &mut Writer,
        asset: SavedAsset<'_, '_, Self::Asset>,
        _settings: &Self::Settings,
        _asset_path: AssetPath<'_>,
    ) -> Result<(), Self::Error> {
        let ron = to_scatter_plan_ron(asset.get())?;
        writer.write_all(ron.as_bytes()).await?;
        Ok(())
    }
}

#[cfg(feature = "ron")]
fn to_scatter_plan_ron(asset: &ScatterPlanAsset) -> Result<String, ron::Error> {
    let stable = SerializableScatterPlanAsset::from(asset);
    let pretty = ron::ser::PrettyConfig::new()
        .new_line("\n")
        .indentor("  ")
        .struct_names(false);
    let mut ron = ron::ser::to_string_pretty(&stable, pretty)?;
    ron.push('\n');
    Ok(ron)
}

#[cfg(feature = "ron")]
#[derive(Serialize)]
struct SerializableScatterPlanAsset<'a> {
    layers: Vec<SerializableScatterLayerDef<'a>>,
}

#[cfg(feature = "ron")]
impl<'a> From<&'a ScatterPlanAsset> for SerializableScatterPlanAsset<'a> {
    fn from(asset: &'a ScatterPlanAsset) -> Self {
        Self {
            layers: asset.layers.iter().map(Into::into).collect(),
        }
    }
}

#[cfg(feature = "ron")]
#[derive(Serialize)]
struct SerializableScatterLayerDef<'a> {
    id: &'a str,
    kinds: Vec<SerializableScatterKindDef<'a>>,
    sampling: &'a SamplingDef,
    overlay_mask_size_px: &'a Option<(u32, u32)>,
    overlay_brush_radius_px: &'a Option<i32>,
    selection_strategy: &'a SelectionStrategyDef,
}

#[cfg(feature = "ron")]
impl<'a> From<&'a ScatterLayerDef> for SerializableScatterLayerDef<'a> {
    fn from(layer: &'a ScatterLayerDef) -> Self {
        Self {
            id: &layer.id,
            kinds: layer.kinds.iter().map(Into::into).collect(),
            sampling: &layer.sampling,
            overlay_mask_size_px: &layer.overlay_mask_size_px,
            overlay_brush_radius_px: &layer.overlay_brush_radius_px,
            selection_strategy: &layer.selection_strategy,
        }
    }
}

#[cfg(feature = "ron")]
#[derive(Serialize)]
struct SerializableScatterKindDef<'a> {
    id: &'a str,
    spec: SerializableFieldGraphSpec<'a>,
}

#[cfg(feature = "ron")]
impl<'a> From<&'a ScatterKindDef> for SerializableScatterKindDef<'a> {
    fn from(kind: &'a ScatterKindDef) -> Self {
        Self {
            id: &kind.id,
            spec: SerializableFieldGraphSpec::from(&kind.spec),
        }
    }
}

#[cfg(feature = "ron")]
#[derive(Serialize)]
struct SerializableFieldGraphSpec<'a> {
    nodes: BTreeMap<&'a str, &'a NodeSpec>,
    semantics: BTreeMap<&'a str, &'a FieldSemantics>,
}

#[cfg(feature = "ron")]
impl<'a> From<&'a FieldGraphSpec> for SerializableFieldGraphSpec<'a> {
    fn from(spec: &'a FieldGraphSpec) -> Self {
        Self {
            nodes: spec
                .nodes
                .iter()
                .map(|(id, node)| (id.as_str(), node))
                .collect(),
            semantics: spec
                .semantics
                .iter()
                .map(|(id, semantics)| (id.as_str(), semantics))
                .collect(),
        }
    }
}

#[cfg(all(test, feature = "ron"))]
mod tests {
    use super::*;

    #[test]
    fn serializes_simple_scatter_plan_asset_as_readable_ron() {
        let asset = simple_plan_asset();

        let ron = to_scatter_plan_ron(&asset).expect("scatter plan should serialize");

        assert!(ron.contains("layers"));
        assert!(ron.contains("JitterGrid"));
        assert!(ron.contains("\"probability\""));
        assert!(ron.ends_with('\n'));
    }

    #[test]
    fn serialized_scatter_plan_asset_loads_through_ron_parser() {
        let asset = simple_plan_asset();
        let ron = to_scatter_plan_ron(&asset).expect("scatter plan should serialize");

        let parsed: ScatterPlanAsset =
            ron::de::from_str(&ron).expect("serialized plan should parse through RON");

        assert_eq!(parsed.layers.len(), 1);
        let layer = &parsed.layers[0];
        assert_eq!(layer.id, "dots");
        assert_eq!(layer.kinds.len(), 1);
        assert_eq!(layer.kinds[0].id, "dot");
        assert!(matches!(
            layer.sampling,
            SamplingDef::JitterGrid {
                jitter: 1.0,
                cell_size: 1.0
            }
        ));
        assert!(matches!(
            layer.kinds[0].spec.semantics.get("probability"),
            Some(FieldSemantics::Probability)
        ));
        assert!(matches!(
            layer.kinds[0].spec.nodes.get("probability"),
            Some(NodeSpec::Constant { .. })
        ));
    }

    #[test]
    fn scatter_plan_asset_saver_type_is_available_with_ron_feature() {
        let _saver = ScatterPlanAssetSaver;
    }

    fn simple_plan_asset() -> ScatterPlanAsset {
        let mut spec = FieldGraphSpec::default();
        spec.add_with_semantics(
            "probability",
            NodeSpec::constant(1.0),
            FieldSemantics::Probability,
        );

        ScatterPlanAsset {
            layers: vec![ScatterLayerDef {
                id: "dots".to_string(),
                kinds: vec![ScatterKindDef {
                    id: "dot".to_string(),
                    spec,
                }],
                sampling: SamplingDef::JitterGrid {
                    jitter: 1.0,
                    cell_size: 1.0,
                },
                overlay_mask_size_px: None,
                overlay_brush_radius_px: None,
                selection_strategy: SelectionStrategyDef::WeightedRandom,
            }],
        }
    }
}