use egui::{Align2, FontId, Pos2, Rect, Sense, Stroke, Ui, vec2};
use facett_core::{Facet, FacetCaps, theme};
use crate::boxplot::paint_tooltip;
#[derive(Clone, Debug, PartialEq)]
pub struct Ohlc {
pub label: String,
pub open: f64,
pub high: f64,
pub low: f64,
pub close: f64,
}
impl Ohlc {
pub fn new(label: impl Into<String>, open: f64, high: f64, low: f64, close: f64) -> Self {
Self { label: label.into(), open, high, low, close }
}
pub fn bullish(&self) -> bool {
self.close >= self.open
}
}
#[derive(Clone, Copy, Debug)]
pub struct OhlcLayout {
plot: Rect,
pub price_min: f64,
pub price_max: f64,
pub count: usize,
pub slot: f32,
pub body_half: f32,
}
impl OhlcLayout {
pub fn new(plot: Rect, price_min: f64, price_max: f64, count: usize) -> Self {
let slot = if count == 0 { plot.width() } else { plot.width() / count as f32 };
let body_half = (slot * 0.3).clamp(1.0, 30.0);
Self { plot, price_min, price_max, count, slot, body_half }
}
pub fn x_center(&self, i: usize) -> f32 {
self.plot.left() + self.slot * (i as f32 + 0.5)
}
pub fn y_of(&self, price: f64) -> f32 {
let span = (self.price_max - self.price_min).max(1e-9);
self.plot.bottom() - ((price - self.price_min) / span) as f32 * self.plot.height()
}
pub fn bar_at_x(&self, x: f32) -> Option<usize> {
if x < self.plot.left() || x > self.plot.right() || self.count == 0 {
return None;
}
let i = ((x - self.plot.left()) / self.slot).floor() as usize;
if i < self.count { Some(i) } else { None }
}
}
pub struct CandlestickView {
pub title: String,
pub bars: Vec<Ohlc>,
}
impl CandlestickView {
pub fn new(title: impl Into<String>, bars: Vec<Ohlc>) -> Self {
Self { title: title.into(), bars }
}
pub fn price_range(&self) -> Option<(f64, f64)> {
let mut r: Option<(f64, f64)> = None;
for b in &self.bars {
r = Some(match r {
None => (b.low, b.high),
Some((lo, hi)) => (lo.min(b.low), hi.max(b.high)),
});
}
r
}
fn to_csv(&self) -> String {
let mut s = String::from("label,open,high,low,close\n");
for b in &self.bars {
s.push_str(&format!("{},{},{},{},{}\n", b.label, b.open, b.high, b.low, b.close));
}
s
}
}
impl Facet for CandlestickView {
fn title(&self) -> &str {
&self.title
}
fn caps(&self) -> FacetCaps {
FacetCaps::NONE.themeable().copyable().resizable()
}
fn copy(&mut self) -> Option<String> {
Some(self.to_csv())
}
fn ui(&mut self, ui: &mut Ui) {
let th = theme(ui);
let (rect, resp) = ui.allocate_exact_size(ui.available_size(), Sense::hover());
let painter = ui.painter_at(rect);
painter.rect_filled(rect, egui::CornerRadius::ZERO, th.bg);
let Some((pmin, pmax)) = self.price_range() else {
painter.text(rect.center(), Align2::CENTER_CENTER, "no data", FontId::proportional(13.0), th.text_dim);
return;
};
let plot = rect.shrink2(vec2(36.0, 16.0));
let lay = OhlcLayout::new(plot, pmin, pmax, self.bars.len());
let up = th.point;
let down = crate::ramp::lerp_color(th.text_dim, th.accent, 0.0);
for (i, b) in self.bars.iter().enumerate() {
let cx = lay.x_center(i);
let color = if b.bullish() { up } else { down };
painter.line_segment([Pos2::new(cx, lay.y_of(b.high)), Pos2::new(cx, lay.y_of(b.low))], Stroke::new(1.0, color));
let (top, bot) = (b.open.max(b.close), b.open.min(b.close));
let body = Rect::from_min_max(
Pos2::new(cx - lay.body_half, lay.y_of(top)),
Pos2::new(cx + lay.body_half, lay.y_of(bot)),
);
let body = if body.height() < 1.0 {
Rect::from_min_max(body.min, Pos2::new(body.right(), body.top() + 1.0))
} else {
body
};
if b.bullish() {
painter.rect_filled(body, egui::CornerRadius::same(1), color.gamma_multiply(0.7));
} else {
painter.rect_filled(body, egui::CornerRadius::same(1), th.bg);
painter.rect_stroke(body, egui::CornerRadius::same(1), Stroke::new(1.2, color), egui::StrokeKind::Middle);
}
}
if let Some(hp) = resp.hover_pos() {
if let Some(i) = lay.bar_at_x(hp.x) {
let b = &self.bars[i];
let cx = lay.x_center(i);
painter.line_segment([Pos2::new(cx, plot.top()), Pos2::new(cx, plot.bottom())], Stroke::new(0.5, th.text_dim.gamma_multiply(0.5)));
paint_tooltip(
&painter,
&th,
hp,
&format!("{}\nO {:.2}\nH {:.2}\nL {:.2}\nC {:.2}", b.label, b.open, b.high, b.low, b.close),
);
}
}
}
fn state_json(&self) -> serde_json::Value {
let (mn, mx) = self.price_range().unwrap_or((0.0, 0.0));
let up = self.bars.iter().filter(|b| b.bullish()).count();
serde_json::json!({
"kind": "candlestick",
"bars": self.bars.len(),
"bullish": up,
"bearish": self.bars.len() - up,
"low": mn, "high": mx,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn plot() -> Rect {
Rect::from_min_size(Pos2::new(0.0, 0.0), vec2(200.0, 100.0))
}
#[test]
fn bullish_classification() {
assert!(Ohlc::new("d", 10.0, 12.0, 9.0, 11.0).bullish());
assert!(!Ohlc::new("d", 11.0, 12.0, 9.0, 10.0).bullish());
assert!(Ohlc::new("d", 10.0, 12.0, 9.0, 10.0).bullish());
}
#[test]
fn y_mapping_is_inverted_and_spans_plot() {
let lay = OhlcLayout::new(plot(), 0.0, 100.0, 4);
assert!((lay.y_of(100.0) - plot().top()).abs() < 1e-3);
assert!((lay.y_of(0.0) - plot().bottom()).abs() < 1e-3);
assert!((lay.y_of(50.0) - plot().center().y).abs() < 1e-3);
}
#[test]
fn x_centers_are_evenly_spaced() {
let lay = OhlcLayout::new(plot(), 0.0, 1.0, 4);
let c0 = lay.x_center(0);
let c1 = lay.x_center(1);
assert!((c1 - c0 - lay.slot).abs() < 1e-4);
assert!((c0 - lay.slot * 0.5).abs() < 1e-4);
}
#[test]
fn bar_at_x_round_trips_through_centers() {
let lay = OhlcLayout::new(plot(), 0.0, 1.0, 5);
for i in 0..5 {
assert_eq!(lay.bar_at_x(lay.x_center(i)), Some(i));
}
assert_eq!(lay.bar_at_x(plot().left() - 5.0), None);
assert_eq!(lay.bar_at_x(plot().right() + 5.0), None);
}
#[test]
fn price_range_covers_low_and_high() {
let v = CandlestickView::new(
"t",
vec![Ohlc::new("a", 5.0, 9.0, 3.0, 7.0), Ohlc::new("b", 7.0, 12.0, 6.0, 8.0)],
);
assert_eq!(v.price_range(), Some((3.0, 12.0)));
}
#[test]
fn csv_header_and_rows() {
let v = CandlestickView::new("t", vec![Ohlc::new("a", 1.0, 2.0, 0.5, 1.5)]);
let csv = v.to_csv();
assert!(csv.starts_with("label,open,high,low,close\n"));
assert_eq!(csv.lines().count(), 2);
}
}