1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

use crate::Error;
use evoxel_core::{VoxelDataColumnNames, VoxelGrid};
use nalgebra::Point3;

use polars::frame::DataFrame;

use crate::Error::LowerCornerMustBeBelowUpperCorner;
use polars::prelude::{all, col, len, ChunkCompare, IntoLazy};

pub fn aggregate_by_index(voxel_grid: &VoxelGrid) -> Result<VoxelGrid, Error> {
    let voxel_data = voxel_grid.voxel_data();
    let partition_columns = vec![
        VoxelDataColumnNames::X.as_str(),
        VoxelDataColumnNames::Y.as_str(),
        VoxelDataColumnNames::Z.as_str(),
    ];

    let partitioned: DataFrame = voxel_data
        .clone()
        .lazy()
        .group_by(partition_columns)
        .agg([all(), len()])
        //.limit(5)
        .collect()?;
    //dbg!("{:?}", &partitioned);

    /*println!("partitions: {}", partitioned.len());
    let parit_filtered: Vec<DataFrame> = partitioned
        .into_iter()
        .filter(|d| d.height() >= minimum)
        .collect();

    let partitioned_lazy: Vec<LazyFrame> =
        parit_filtered.iter().map(|d| d.clone().lazy()).collect();

    let merged_again = concat(partitioned_lazy, true, true)
        .unwrap()
        .collect()
        .unwrap();*/

    let filtered_voxel_grid = VoxelGrid::new(
        partitioned,
        voxel_grid.info().clone(),
        voxel_grid.reference_frames().clone(),
    )?;
    Ok(filtered_voxel_grid)
}

pub fn explode(voxel_grid: &VoxelGrid) -> Result<VoxelGrid, Error> {
    let voxel_data = voxel_grid.voxel_data();

    let column_names: Vec<&str> = voxel_data
        .get_columns()
        .iter()
        .filter(|s| s.dtype().inner_dtype().is_some()) // if contains inner, it's a list
        .map(|s| s.name())
        .collect();

    let df: DataFrame = voxel_data
        .clone()
        .lazy()
        .explode(column_names.into_iter().map(col).collect::<Vec<_>>())
        //.limit(5)
        .collect()?;

    let filtered_voxel_grid = VoxelGrid::new(
        df,
        voxel_grid.info().clone(),
        voxel_grid.reference_frames().clone(),
    )?;
    Ok(filtered_voxel_grid)
}

pub fn filter_by_count(voxel_grid: &VoxelGrid, minimum: usize) -> Result<VoxelGrid, Error> {
    let voxel_data = voxel_grid.voxel_data().clone();

    let mask = voxel_data
        .column(VoxelDataColumnNames::Count.as_str())?
        .gt_eq(minimum as i32)?;

    let filtered_voxel_data = voxel_data.filter(&mask)?;
    let filtered_voxel_grid = VoxelGrid::new(
        filtered_voxel_data,
        voxel_grid.info().clone(),
        voxel_grid.reference_frames().clone(),
    )?;

    Ok(filtered_voxel_grid)
}

pub fn filter_by_index_bounds(
    voxel_grid: &VoxelGrid,
    lower_corner: Point3<i64>,
    upper_corner: Point3<i64>,
) -> Result<VoxelGrid, Error> {
    if lower_corner >= upper_corner {
        return Err(LowerCornerMustBeBelowUpperCorner(""));
    }

    let voxel_data = voxel_grid.voxel_data().clone();

    let filtered_voxel_data = voxel_data
        .lazy()
        .filter(
            col(VoxelDataColumnNames::X.as_str())
                .gt_eq(lower_corner.x)
                .and(col(VoxelDataColumnNames::X.as_str()).lt_eq(upper_corner.x)),
        )
        .filter(
            col(VoxelDataColumnNames::Y.as_str())
                .gt_eq(lower_corner.y)
                .and(col(VoxelDataColumnNames::Y.as_str()).lt_eq(upper_corner.y)),
        )
        .filter(
            col(VoxelDataColumnNames::Z.as_str())
                .gt_eq(lower_corner.z)
                .and(col(VoxelDataColumnNames::Z.as_str()).lt_eq(upper_corner.z)),
        )
        .collect()?;
    let filtered_voxel_grid = VoxelGrid::new(
        filtered_voxel_data,
        voxel_grid.info().clone(),
        voxel_grid.reference_frames().clone(),
    )?;

    Ok(filtered_voxel_grid)
}