1use bon::bon;
2
3use plotly::{Layout as LayoutPlotly, Ohlc as OhlcPlotly, Trace};
4
5use polars::frame::DataFrame;
6use serde::Serialize;
7
8use crate::{
9 common::{Layout, PlotHelper, Polar},
10 components::{Axis, Text},
11};
12
13#[derive(Clone, Serialize)]
65pub struct OhlcPlot {
66 traces: Vec<Box<dyn Trace + 'static>>,
67 layout: LayoutPlotly,
68}
69
70#[bon]
71impl OhlcPlot {
72 #[builder(on(String, into), on(Text, into))]
73 pub fn new(
74 data: &DataFrame,
75 dates: &str,
76 open: &str,
77 high: &str,
78 low: &str,
79 close: &str,
80 tick_width: Option<f64>,
81 plot_title: Option<Text>,
82 x_title: Option<Text>,
83 y_title: Option<Text>,
84 x_axis: Option<&Axis>,
85 y_axis: Option<&Axis>,
86 ) -> Self {
87 let z_title = None;
88 let y_title2 = None;
89 let z_axis = None;
90 let y_axis2 = None;
91 let legend_title = None;
92 let legend = None;
93
94 let layout = Self::create_layout(
95 plot_title,
96 x_title,
97 y_title,
98 y_title2,
99 z_title,
100 legend_title,
101 x_axis,
102 y_axis,
103 y_axis2,
104 z_axis,
105 legend,
106 None,
107 );
108
109 let traces = Self::create_traces(data, dates, open, high, low, close, tick_width);
110
111 Self { traces, layout }
112 }
113
114 fn create_traces(
115 data: &DataFrame,
116 dates_col: &str,
117 open_col: &str,
118 high_col: &str,
119 low_col: &str,
120 close_col: &str,
121 tick_width: Option<f64>,
122 ) -> Vec<Box<dyn Trace + 'static>> {
123 let mut traces: Vec<Box<dyn Trace + 'static>> = Vec::new();
124
125 let trace = Self::create_trace(
126 data, dates_col, open_col, high_col, low_col, close_col, tick_width,
127 );
128
129 traces.push(trace);
130 traces
131 }
132
133 fn create_trace(
134 data: &DataFrame,
135 dates_col: &str,
136 open_col: &str,
137 high_col: &str,
138 low_col: &str,
139 close_col: &str,
140 tick_width: Option<f64>,
141 ) -> Box<dyn Trace + 'static> {
142 let dates_data = Self::get_string_column(data, dates_col);
143 let open_data = Self::get_numeric_column(data, open_col);
144 let high_data = Self::get_numeric_column(data, high_col);
145 let low_data = Self::get_numeric_column(data, low_col);
146 let close_data = Self::get_numeric_column(data, close_col);
147
148 let open_values: Vec<f32> = open_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
150 let high_values: Vec<f32> = high_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
151 let low_values: Vec<f32> = low_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
152 let close_values: Vec<f32> = close_data.into_iter().map(|v| v.unwrap_or(0.0)).collect();
153 let dates_values: Vec<String> = dates_data
154 .into_iter()
155 .map(|v| v.unwrap_or_default())
156 .collect();
157
158 let mut trace = *OhlcPlotly::new(
159 dates_values,
160 open_values,
161 high_values,
162 low_values,
163 close_values,
164 );
165
166 if let Some(tick_w) = tick_width {
168 trace = trace.tick_width(tick_w);
169 }
170
171 Box::new(trace)
173 }
174}
175
176impl Layout for OhlcPlot {}
177impl Polar for OhlcPlot {}
178
179impl PlotHelper for OhlcPlot {
180 fn get_layout(&self) -> &LayoutPlotly {
181 &self.layout
182 }
183
184 fn get_traces(&self) -> &Vec<Box<dyn Trace + 'static>> {
185 &self.traces
186 }
187}