Skip to main content

oxigdal_vrt/
dataset.rs

1//! VRT dataset definition
2
3use crate::band::VrtBand;
4use crate::error::{Result, VrtError};
5use crate::source::PixelRect;
6use oxigdal_core::types::{GeoTransform, RasterDataType};
7use serde::{Deserialize, Serialize};
8use std::path::PathBuf;
9
10/// VRT dataset
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
12pub struct VrtDataset {
13    /// Raster width in pixels
14    pub raster_x_size: u64,
15    /// Raster height in pixels
16    pub raster_y_size: u64,
17    /// GeoTransform (affine transform for georeferencing)
18    pub geo_transform: Option<GeoTransform>,
19    /// Spatial reference system (WKT or PROJ.4 string)
20    pub srs: Option<String>,
21    /// Bands
22    pub bands: Vec<VrtBand>,
23    /// Block size (default tile dimensions)
24    pub block_size: Option<(u32, u32)>,
25    /// Subclass (for special VRT types)
26    pub subclass: Option<VrtSubclass>,
27    /// VRT file path (for resolving relative paths)
28    pub vrt_path: Option<PathBuf>,
29}
30
31impl VrtDataset {
32    /// Creates a new VRT dataset
33    pub fn new(raster_x_size: u64, raster_y_size: u64) -> Self {
34        Self {
35            raster_x_size,
36            raster_y_size,
37            geo_transform: None,
38            srs: None,
39            bands: Vec::new(),
40            block_size: None,
41            subclass: None,
42            vrt_path: None,
43        }
44    }
45
46    /// Creates a new VRT dataset with extent from sources
47    ///
48    /// # Errors
49    /// Returns an error if no bands are provided
50    pub fn from_bands(bands: Vec<VrtBand>) -> Result<Self> {
51        if bands.is_empty() {
52            return Err(VrtError::invalid_structure(
53                "Dataset must have at least one band",
54            ));
55        }
56
57        // Calculate extent from first band's sources
58        let (width, height) = Self::calculate_extent_from_band(&bands[0])?;
59
60        Ok(Self {
61            raster_x_size: width,
62            raster_y_size: height,
63            geo_transform: None,
64            srs: None,
65            bands,
66            block_size: None,
67            subclass: None,
68            vrt_path: None,
69        })
70    }
71
72    /// Adds a band to the dataset
73    pub fn add_band(&mut self, band: VrtBand) {
74        self.bands.push(band);
75    }
76
77    /// Sets the GeoTransform
78    pub fn with_geo_transform(mut self, geo_transform: GeoTransform) -> Self {
79        self.geo_transform = Some(geo_transform);
80        self
81    }
82
83    /// Sets the spatial reference system
84    pub fn with_srs<S: Into<String>>(mut self, srs: S) -> Self {
85        self.srs = Some(srs.into());
86        self
87    }
88
89    /// Sets the block size
90    pub fn with_block_size(mut self, width: u32, height: u32) -> Self {
91        self.block_size = Some((width, height));
92        self
93    }
94
95    /// Sets the subclass
96    pub fn with_subclass(mut self, subclass: VrtSubclass) -> Self {
97        self.subclass = Some(subclass);
98        self
99    }
100
101    /// Sets the VRT file path
102    pub fn with_vrt_path<P: Into<PathBuf>>(mut self, path: P) -> Self {
103        self.vrt_path = Some(path.into());
104        self
105    }
106
107    /// Validates the dataset
108    ///
109    /// # Errors
110    /// Returns an error if the dataset is invalid
111    pub fn validate(&self) -> Result<()> {
112        if self.raster_x_size == 0 || self.raster_y_size == 0 {
113            return Err(VrtError::invalid_structure(
114                "Dataset dimensions must be > 0",
115            ));
116        }
117
118        if self.bands.is_empty() {
119            return Err(VrtError::invalid_structure(
120                "Dataset must have at least one band",
121            ));
122        }
123
124        // Validate all bands
125        for (idx, band) in self.bands.iter().enumerate() {
126            band.validate().map_err(|e| {
127                VrtError::invalid_structure(format!("Band {} validation failed: {}", idx + 1, e))
128            })?;
129
130            // Check that band numbers are sequential
131            if band.band != idx + 1 {
132                return Err(VrtError::invalid_structure(format!(
133                    "Band number mismatch: expected {}, got {}",
134                    idx + 1,
135                    band.band
136                )));
137            }
138        }
139
140        // Validate block size if present
141        if let Some((width, height)) = self.block_size {
142            if width == 0 || height == 0 {
143                return Err(VrtError::invalid_structure("Block size must be > 0"));
144            }
145        }
146
147        Ok(())
148    }
149
150    /// Gets the number of bands
151    pub fn band_count(&self) -> usize {
152        self.bands.len()
153    }
154
155    /// Gets a band by index (0-based)
156    pub fn get_band(&self, index: usize) -> Option<&VrtBand> {
157        self.bands.get(index)
158    }
159
160    /// Gets a mutable reference to a band by index (0-based)
161    pub fn get_band_mut(&mut self, index: usize) -> Option<&mut VrtBand> {
162        self.bands.get_mut(index)
163    }
164
165    /// Gets the extent as a pixel rectangle
166    pub fn extent(&self) -> PixelRect {
167        PixelRect::new(0, 0, self.raster_x_size, self.raster_y_size)
168    }
169
170    /// Gets the effective block size (uses dataset default or falls back to 256x256)
171    pub fn effective_block_size(&self) -> (u32, u32) {
172        self.block_size.unwrap_or((256, 256))
173    }
174
175    /// Calculates extent from a band's sources
176    fn calculate_extent_from_band(band: &VrtBand) -> Result<(u64, u64)> {
177        if band.sources.is_empty() {
178            return Err(VrtError::invalid_structure(
179                "Band has no sources to calculate extent from",
180            ));
181        }
182
183        // Find the bounding box of all destination rectangles
184        let mut min_x = u64::MAX;
185        let mut min_y = u64::MAX;
186        let mut max_x = 0u64;
187        let mut max_y = 0u64;
188
189        for source in &band.sources {
190            if let Some(dst_rect) = source.dst_rect() {
191                min_x = min_x.min(dst_rect.x_off);
192                min_y = min_y.min(dst_rect.y_off);
193                max_x = max_x.max(dst_rect.x_off + dst_rect.x_size);
194                max_y = max_y.max(dst_rect.y_off + dst_rect.y_size);
195            } else if let Some(ref props) = source.properties {
196                // If no dst_rect, use source properties
197                max_x = max_x.max(props.width);
198                max_y = max_y.max(props.height);
199            }
200        }
201
202        if max_x == 0 || max_y == 0 {
203            return Err(VrtError::invalid_structure(
204                "Cannot calculate extent: no valid source windows",
205            ));
206        }
207
208        Ok((max_x - min_x, max_y - min_y))
209    }
210
211    /// Merges GeoTransforms from multiple sources to create a unified transform
212    ///
213    /// # Errors
214    /// Returns an error if sources have incompatible GeoTransforms
215    pub fn merge_geo_transforms(&mut self) -> Result<()> {
216        if self.geo_transform.is_some() {
217            return Ok(()); // Already set
218        }
219
220        let mut transforms = Vec::new();
221
222        // Collect all GeoTransforms from band sources
223        for band in &self.bands {
224            for source in &band.sources {
225                if let Some(ref props) = source.properties {
226                    if let Some(ref gt) = props.geo_transform {
227                        transforms.push(*gt);
228                    }
229                }
230            }
231        }
232
233        if transforms.is_empty() {
234            return Ok(()); // No GeoTransforms to merge
235        }
236
237        // For simplicity, use the first transform
238        // In a real implementation, we would validate that all transforms are compatible
239        self.geo_transform = Some(transforms[0]);
240
241        Ok(())
242    }
243
244    /// Gets the data type of the first band
245    pub fn primary_data_type(&self) -> Option<RasterDataType> {
246        self.bands.first().map(|b| b.data_type)
247    }
248
249    /// Checks if all bands have the same data type
250    pub fn has_uniform_data_type(&self) -> bool {
251        if self.bands.is_empty() {
252            return true;
253        }
254
255        let first_type = self.bands[0].data_type;
256        self.bands.iter().all(|b| b.data_type == first_type)
257    }
258}
259
260/// VRT subclass types
261#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
262pub enum VrtSubclass {
263    /// Standard VRT
264    #[default]
265    Standard,
266    /// Warped VRT (for reprojection)
267    Warped,
268    /// Pansharpened VRT
269    Pansharpened,
270    /// Processed VRT (with pixel functions)
271    Processed,
272}
273
274/// VRT metadata
275#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
276pub struct VrtMetadata {
277    /// Metadata domain
278    pub domain: Option<String>,
279    /// Metadata items (key-value pairs)
280    pub items: Vec<(String, String)>,
281}
282
283impl VrtMetadata {
284    /// Creates new VRT metadata
285    pub fn new() -> Self {
286        Self {
287            domain: None,
288            items: Vec::new(),
289        }
290    }
291
292    /// Creates VRT metadata with a domain
293    pub fn with_domain<S: Into<String>>(domain: S) -> Self {
294        Self {
295            domain: Some(domain.into()),
296            items: Vec::new(),
297        }
298    }
299
300    /// Adds a metadata item
301    pub fn add_item<K: Into<String>, V: Into<String>>(&mut self, key: K, value: V) {
302        self.items.push((key.into(), value.into()));
303    }
304
305    /// Gets a metadata value by key
306    pub fn get(&self, key: &str) -> Option<&str> {
307        self.items
308            .iter()
309            .find(|(k, _)| k == key)
310            .map(|(_, v)| v.as_str())
311    }
312}
313
314impl Default for VrtMetadata {
315    fn default() -> Self {
316        Self::new()
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use crate::band::VrtBand;
324    use crate::source::{SourceFilename, VrtSource};
325    use oxigdal_core::types::RasterDataType;
326
327    #[test]
328    fn test_vrt_dataset_creation() {
329        let dataset = VrtDataset::new(512, 512);
330        assert_eq!(dataset.raster_x_size, 512);
331        assert_eq!(dataset.raster_y_size, 512);
332        assert_eq!(dataset.band_count(), 0);
333    }
334
335    #[test]
336    fn test_vrt_dataset_validation() {
337        let mut dataset = VrtDataset::new(512, 512);
338        let source = VrtSource::new(SourceFilename::absolute("/test.tif"), 1);
339        let band = VrtBand::simple(1, RasterDataType::UInt8, source);
340        dataset.add_band(band);
341
342        assert!(dataset.validate().is_ok());
343
344        let empty_dataset = VrtDataset::new(512, 512);
345        assert!(empty_dataset.validate().is_err());
346
347        let invalid_dataset = VrtDataset::new(0, 0);
348        assert!(invalid_dataset.validate().is_err());
349    }
350
351    #[test]
352    fn test_vrt_dataset_extent() {
353        let dataset = VrtDataset::new(1024, 768);
354        let extent = dataset.extent();
355        assert_eq!(extent.x_size, 1024);
356        assert_eq!(extent.y_size, 768);
357    }
358
359    #[test]
360    fn test_vrt_dataset_band_access() {
361        let mut dataset = VrtDataset::new(512, 512);
362        let source = VrtSource::new(SourceFilename::absolute("/test.tif"), 1);
363        let band = VrtBand::simple(1, RasterDataType::UInt8, source);
364        dataset.add_band(band);
365
366        assert_eq!(dataset.band_count(), 1);
367        assert!(dataset.get_band(0).is_some());
368        assert!(dataset.get_band(1).is_none());
369    }
370
371    #[test]
372    fn test_effective_block_size() {
373        let dataset = VrtDataset::new(512, 512);
374        assert_eq!(dataset.effective_block_size(), (256, 256));
375
376        let dataset_with_blocks = VrtDataset::new(512, 512).with_block_size(128, 128);
377        assert_eq!(dataset_with_blocks.effective_block_size(), (128, 128));
378    }
379
380    #[test]
381    fn test_vrt_metadata() {
382        let mut metadata = VrtMetadata::new();
383        metadata.add_item("author", "test");
384        metadata.add_item("version", "1.0");
385
386        assert_eq!(metadata.get("author"), Some("test"));
387        assert_eq!(metadata.get("version"), Some("1.0"));
388        assert_eq!(metadata.get("missing"), None);
389    }
390
391    #[test]
392    fn test_uniform_data_type() {
393        let mut dataset = VrtDataset::new(512, 512);
394
395        let source1 = VrtSource::new(SourceFilename::absolute("/test1.tif"), 1);
396        let band1 = VrtBand::simple(1, RasterDataType::UInt8, source1);
397        dataset.add_band(band1);
398
399        let source2 = VrtSource::new(SourceFilename::absolute("/test2.tif"), 1);
400        let band2 = VrtBand::simple(2, RasterDataType::UInt8, source2);
401        dataset.add_band(band2);
402
403        assert!(dataset.has_uniform_data_type());
404
405        let source3 = VrtSource::new(SourceFilename::absolute("/test3.tif"), 1);
406        let band3 = VrtBand::simple(3, RasterDataType::Float32, source3);
407        dataset.add_band(band3);
408
409        assert!(!dataset.has_uniform_data_type());
410    }
411}