use diesel::{
debug_query,
pg::Pg,
sql_types::{Float, Integer, Uuid as SqlUuid},
CombineDsl, ExpressionMethods, QueryDsl, QueryResult, QueryableByName,
};
use diesel_async::{AsyncPgConnection, RunQueryDsl};
use log::debug;
use uuid::Uuid;
use crate::{
model::{
dto::{RelationDto, RelationSearchParameters, RelationsDto},
r#enum::relation_type::RelationType,
},
schema::relations,
};
pub const GRANULARITY: i32 = 10;
#[derive(Debug, Clone, QueryableByName)]
struct BoundingBox {
#[diesel(sql_type = Integer)]
x_min: i32,
#[diesel(sql_type = Integer)]
y_min: i32,
#[diesel(sql_type = Integer)]
x_max: i32,
#[diesel(sql_type = Integer)]
y_max: i32,
}
#[derive(Debug, Clone, QueryableByName)]
struct HeatMapElement {
#[diesel(sql_type = Float)]
score: f32,
#[diesel(sql_type = Integer)]
x: i32,
#[diesel(sql_type = Integer)]
y: i32,
}
#[allow(
clippy::cast_sign_loss, clippy::indexing_slicing, clippy::cast_possible_truncation, )]
pub async fn heatmap(
map_id: i32,
layer_id: Uuid,
plant_id: i32,
conn: &mut AsyncPgConnection,
) -> QueryResult<Vec<Vec<f32>>> {
let bounding_box_query =
diesel::sql_query("SELECT * FROM calculate_bbox($1)").bind::<Integer, _>(map_id);
debug!("{}", debug_query::<Pg, _>(&bounding_box_query));
let bounding_box = bounding_box_query.get_result::<BoundingBox>(conn).await?;
let query = diesel::sql_query("SELECT * FROM calculate_score($1, $2, $3, $4, $5, $6, $7, $8)")
.bind::<Integer, _>(map_id)
.bind::<SqlUuid, _>(layer_id)
.bind::<Integer, _>(plant_id)
.bind::<Integer, _>(GRANULARITY)
.bind::<Integer, _>(bounding_box.x_min)
.bind::<Integer, _>(bounding_box.y_min)
.bind::<Integer, _>(bounding_box.x_max)
.bind::<Integer, _>(bounding_box.y_max);
debug!("{}", debug_query::<Pg, _>(&query));
let result = query.load::<HeatMapElement>(conn).await?;
let num_cols =
(f64::from(bounding_box.x_max - bounding_box.x_min) / f64::from(GRANULARITY)).ceil();
let num_rows =
(f64::from(bounding_box.y_max - bounding_box.y_min) / f64::from(GRANULARITY)).ceil();
let mut heatmap = vec![vec![0.0; num_cols as usize]; num_rows as usize];
for HeatMapElement { score, x, y } in result {
heatmap[y as usize][x as usize] = score;
}
Ok(heatmap)
}
pub async fn find_relations(
search_query: RelationSearchParameters,
conn: &mut AsyncPgConnection,
) -> QueryResult<RelationsDto> {
let query = relations::table
.select((relations::plant2, relations::relation))
.filter(relations::plant1.eq(&search_query.plant_id))
.union(
relations::table
.select((relations::plant1, relations::relation))
.filter(relations::plant2.eq(&search_query.plant_id)),
);
debug!("{}", debug_query::<Pg, _>(&query));
let relations = query
.load::<(i32, RelationType)>(conn)
.await?
.into_iter()
.map(|(id, relation)| RelationDto { id, relation })
.collect();
Ok(RelationsDto {
id: search_query.plant_id,
relations,
})
}