1use ratatui::buffer::Buffer;
6use std::time::{Duration, Instant};
7
8use ratatui::crossterm::event::{
9 self, Event, KeyCode, KeyEvent, KeyEventKind, KeyModifiers, KeyboardEnhancementFlags,
10 MouseEvent, PopKeyboardEnhancementFlags, PushKeyboardEnhancementFlags,
11};
12use ratatui::layout::{Constraint, Layout, Rect};
13use ratatui::style::{Color, Modifier, Style};
14use ratatui::text::{Line, Span};
15use ratatui::widgets::{Block, BorderType};
16use ratatui::Frame;
17
18use crate::qc::log_bin_key;
19
20pub const ACCENT: Color = Color::Rgb(217, 119, 87);
24
25pub const PLAIN: Style = Style::new();
27pub const DIM: Style = Style::new().add_modifier(Modifier::DIM);
29pub const ACCENTED: Style = Style::new().fg(ACCENT);
31pub const HIGHLIGHT: Style = Style::new().fg(ACCENT).add_modifier(Modifier::BOLD);
33
34pub trait Screen {
36 fn render(&mut self, frame: &mut Frame);
37 fn handle_key(&mut self, key: KeyEvent);
39 fn interrupt(&mut self);
40 fn done(&self) -> bool;
41 fn pending_work(&self) -> Option<String> {
44 None
45 }
46 fn do_work(&mut self) {}
47 fn tick(&mut self) -> bool {
50 false
51 }
52 fn takes_mouse(&self) -> bool {
55 false
56 }
57 fn mouse(&mut self, _event: MouseEvent) -> bool {
61 false
62 }
63 fn reports_chords(&self) -> bool {
67 false
68 }
69 fn take_notice(&mut self) -> Option<String> {
72 None
73 }
74}
75
76#[derive(Default)]
79struct Modes {
80 mouse: bool,
81 chords: bool,
82}
83
84const MOUSE_ON: &str = "\x1b[?1000h\x1b[?1006h";
87const MOUSE_OFF: &str = "\x1b[?1006l\x1b[?1000l";
88
89fn send(text: &str) -> bool {
91 let mut out = std::io::stdout();
92 std::io::Write::write_all(&mut out, text.as_bytes()).is_ok()
93 && std::io::Write::flush(&mut out).is_ok()
94}
95
96pub fn notify(text: &str) {
99 let text: String = text.chars().filter(|c| !c.is_control()).collect();
100 send(&format!("\x07\x1b]9;{text}\x07"));
101}
102
103impl Modes {
104 fn ask(&mut self, screen: &impl Screen) {
105 let mut out = std::io::stdout();
106 if screen.takes_mouse() && !self.mouse {
107 self.mouse = send(MOUSE_ON);
108 }
109 if screen.reports_chords() && !self.chords {
110 let flags = KeyboardEnhancementFlags::DISAMBIGUATE_ESCAPE_CODES;
111 self.chords =
112 ratatui::crossterm::execute!(out, PushKeyboardEnhancementFlags(flags)).is_ok();
113 }
114 }
115
116 fn release(&mut self) {
117 let mut out = std::io::stdout();
118 if std::mem::take(&mut self.mouse) {
119 send(MOUSE_OFF);
120 }
121 if std::mem::take(&mut self.chords) {
122 let _ = ratatui::crossterm::execute!(out, PopKeyboardEnhancementFlags);
123 }
124 }
125}
126
127pub const TICK: std::time::Duration = std::time::Duration::from_millis(200);
129
130struct HeldLogs;
132
133impl HeldLogs {
134 fn new() -> Self {
135 crate::aux::logging::hold_logs(true);
136 HeldLogs
137 }
138}
139
140impl Drop for HeldLogs {
141 fn drop(&mut self) {
142 crate::aux::logging::hold_logs(false);
143 }
144}
145
146pub fn run_screen<S: Screen>(screen: &mut S) -> anyhow::Result<()> {
154 run_screen_with(screen, S::handle_key)
155}
156
157pub fn run_screen_with<S: Screen>(
161 screen: &mut S,
162 mut on_key: impl FnMut(&mut S, KeyEvent),
163) -> anyhow::Result<()> {
164 ratatui::run(|terminal| -> anyhow::Result<()> {
165 let held = HeldLogs::new();
166 let mut modes = Modes::default();
167 let mut redraw = true;
168 let mut next_tick = Instant::now() + TICK;
171 while !screen.done() {
172 if redraw {
173 terminal.draw(|f| screen.render(f))?;
174 modes.ask(screen);
176 }
177 if let Some(notice) = screen.take_notice() {
178 notify(¬ice);
179 }
180 redraw = false;
181 if event::poll(next_tick.saturating_duration_since(Instant::now()))? {
182 loop {
183 redraw |= match event::read()? {
184 Event::Key(key) if key.kind != KeyEventKind::Press => false,
185 Event::Key(key)
186 if key.modifiers.contains(KeyModifiers::CONTROL)
187 && key.code == KeyCode::Char('c') =>
188 {
189 screen.interrupt();
190 true
191 }
192 Event::Key(key) => {
193 on_key(screen, key);
194 true
195 }
196 Event::Mouse(m) => screen.mouse(m),
197 Event::Resize(..) => true,
198 _ => false,
199 };
200 if screen.done()
201 || screen.pending_work().is_some()
202 || !event::poll(Duration::ZERO)?
203 {
204 break;
205 }
206 }
207 }
208 if Instant::now() >= next_tick {
209 redraw |= screen.tick();
210 next_tick = Instant::now() + TICK;
211 }
212 if let Some(message) = screen.pending_work() {
213 modes.release();
214 ratatui::restore();
215 crate::aux::logging::hold_logs(false);
216 eprintln!("{message}");
217 screen.do_work();
218 crate::aux::logging::hold_logs(true);
219 *terminal = ratatui::try_init()?;
220 redraw = true;
221 }
222 }
223 modes.release();
224 ratatui::restore();
226 drop(held);
227 Ok(())
228 })
229}
230
231pub fn header(badge: &str, title: &str, extra: &str) -> Line<'static> {
233 Line::from(vec![
234 Span::styled(
235 format!(" {badge} "),
236 HIGHLIGHT.add_modifier(Modifier::REVERSED),
237 ),
238 Span::raw(format!(" {title}")),
239 Span::styled(format!(" {extra}"), DIM),
240 ])
241}
242
243pub fn panel(title: String, focused: bool) -> Block<'static> {
246 Block::bordered()
247 .border_type(BorderType::Rounded)
248 .border_style(if focused { PLAIN } else { DIM })
249 .title(Line::from(title).style(Style::reset().patch(HIGHLIGHT)))
251}
252
253pub fn help_line(pairs: &[(&str, &str)]) -> Line<'static> {
255 let mut spans = vec![Span::raw(" ")];
256 for (key, what) in pairs {
257 spans.push(Span::styled(key.to_string(), HIGHLIGHT));
258 spans.push(Span::styled(format!(" {what} "), DIM));
259 }
260 Line::from(spans)
261}
262
263pub fn input_line(prompt: &str, text: &str, keys: &[(&str, &str)]) -> Line<'static> {
265 let mut spans = vec![
266 Span::raw(format!(" {prompt}")),
267 Span::styled(format!("{text}▏"), HIGHLIGHT),
268 Span::raw(" "),
269 ];
270 spans.extend(help_line(keys).spans);
271 Line::from(spans)
272}
273
274fn put(buf: &mut Buffer, x: u16, y: u16, symbol: &str, style: Style) {
276 buf[(x, y)]
277 .set_symbol(symbol)
278 .set_style(Style::reset().patch(style));
279}
280
281const GUTTER: u16 = 6;
283
284const TARGET_BINS: f64 = 50.0;
286
287#[derive(Debug, Clone, Copy, PartialEq, Eq)]
289pub enum Scale {
290 Log,
291 Sqrt,
292 Linear,
293}
294
295impl Scale {
296 pub fn next(self) -> Self {
297 match self {
298 Scale::Log => Scale::Sqrt,
299 Scale::Sqrt => Scale::Linear,
300 Scale::Linear => Scale::Log,
301 }
302 }
303
304 pub fn name(self) -> &'static str {
305 match self {
306 Scale::Log => "log",
307 Scale::Sqrt => "sqrt",
308 Scale::Linear => "linear",
309 }
310 }
311
312 fn apply(self, v: f64) -> f64 {
313 match self {
314 Scale::Log => (v + 1.0).log10(),
315 Scale::Sqrt => v.max(0.0).sqrt(),
316 Scale::Linear => v,
317 }
318 }
319
320 fn invert(self, t: f64) -> f64 {
321 match self {
322 Scale::Log => 10f64.powf(t) - 1.0,
323 Scale::Sqrt => t * t,
324 Scale::Linear => t,
325 }
326 }
327}
328
329#[derive(Debug, Clone, Copy)]
332pub struct Binning {
333 pub scale: Scale,
334 width: f64,
336}
337
338impl Binning {
339 pub fn new(scale: Scale, max: f64, integer: bool) -> Self {
342 let span = if integer { max + 1.0 } else { max };
343 let width = match scale {
344 Scale::Log => 0.1,
345 Scale::Linear if integer => (span / TARGET_BINS).ceil().max(1.0),
346 _ => scale.apply(span) / TARGET_BINS,
347 };
348 Self {
349 scale,
350 width: width.max(f64::MIN_POSITIVE),
351 }
352 }
353
354 pub fn with_width(scale: Scale, width: f64) -> Self {
358 Self {
359 scale,
360 width: width.max(f64::MIN_POSITIVE),
361 }
362 }
363
364 pub fn key(&self, x: f64) -> i32 {
365 match self.scale {
366 Scale::Log => log_bin_key(x),
367 _ => (self.scale.apply(x) / self.width).floor() as i32,
368 }
369 }
370
371 fn start(&self, k: i32) -> f64 {
374 match self.scale {
375 Scale::Log => (k as f64 - 0.5) * self.width,
376 _ => k as f64 * self.width,
377 }
378 }
379
380 pub fn lower_edge(&self, k: i32) -> usize {
383 if k <= 0 {
384 return 0;
385 }
386 let mut x = self.scale.invert(self.start(k)).ceil().max(0.0) as usize;
388 while self.key(x as f64) < k {
389 x += 1;
390 }
391 while x > 0 && self.key((x - 1) as f64) >= k {
392 x -= 1;
393 }
394 x
395 }
396
397 fn tick_value(&self, k: i32) -> f64 {
400 self.scale.invert(k as f64 * self.width)
401 }
402
403 fn tick_every(&self, nbins: usize) -> i32 {
405 match self.scale {
406 Scale::Log => 5,
407 _ => (nbins as i32 / 6).max(1),
408 }
409 }
410}
411
412pub struct Binned {
414 pub bins: Binning,
415 pub kmin: i32,
416 pub counts: Vec<usize>,
417}
418
419impl Binned {
420 pub fn new(sorted: &[f32], scale: Scale) -> Self {
422 let (min, max) = match (sorted.first(), sorted.last()) {
423 (Some(&lo), Some(&hi)) => (lo as f64, hi as f64),
424 _ => (0.0, 0.0),
425 };
426 let integer = sorted.iter().all(|v| v.fract() == 0.0);
427 let bins = Binning::new(scale, max, integer);
428 let kmin = bins.key(min);
429 let nbins = (bins.key(max) - kmin + 1).max(1) as usize;
430 let counts = count(&bins, kmin, nbins, sorted.iter().copied());
431 Self { bins, kmin, counts }
432 }
433
434 pub fn count(&self, values: impl Iterator<Item = f32>) -> Vec<usize> {
436 count(&self.bins, self.kmin, self.counts.len(), values)
437 }
438
439 pub fn kmax(&self) -> i32 {
440 self.kmin + self.counts.len() as i32 - 1
441 }
442}
443
444fn count(bins: &Binning, kmin: i32, nbins: usize, values: impl Iterator<Item = f32>) -> Vec<usize> {
447 let mut counts = vec![0; nbins];
448 for v in values {
449 let i = (bins.key(v as f64) - kmin).clamp(0, nbins as i32 - 1);
450 counts[i as usize] += 1;
451 }
452 counts
453}
454
455pub fn median(sorted: &[f32]) -> f32 {
457 crate::qc::median_of_sorted(sorted)
458}
459
460pub fn compact(v: f64) -> String {
463 if v != 0.0 && v.abs() < 10.0 && v.fract() != 0.0 {
464 format!("{:.2}", v)
465 .trim_end_matches('0')
466 .trim_end_matches('.')
467 .to_string()
468 } else if v < 1e3 {
469 format!("{}", v.round() as i64)
470 } else if v < 1e4 {
471 format!("{:.1}k", v / 1e3)
472 } else if v < 1e6 {
473 format!("{}k", (v / 1e3).round() as u64)
474 } else if v < 1e9 {
475 format!("{:.1}M", v / 1e6)
476 } else {
477 format!("{:.1}G", v / 1e9)
478 }
479}
480
481pub trait BarValue: Copy {
485 fn bar(self) -> f64;
486}
487
488impl BarValue for usize {
489 fn bar(self) -> f64 {
490 self as f64
491 }
492}
493
494impl BarValue for f64 {
495 fn bar(self) -> f64 {
496 self
497 }
498}
499
500pub struct HistPlot<'a, T: BarValue = usize> {
501 pub bins: Binning,
502 pub kmin: i32,
503 pub counts: &'a [T],
504 pub style: &'a dyn Fn(i32) -> Style,
506 pub subset: Option<&'a [T]>,
509 pub y_scale: Scale,
510 pub y_max: Option<f64>,
513 pub pointer: Option<i32>,
515 pub marks: Vec<(i32, &'static str, Style)>,
517 pub x_label: Option<&'a dyn Fn(i32) -> Option<String>>,
522 pub tick_every: Option<i32>,
524}
525
526const EIGHTHS: [&str; 8] = ["▁", "▂", "▃", "▄", "▅", "▆", "▇", "█"];
527
528impl<T: BarValue> HistPlot<'_, T> {
529 pub fn render(&self, buf: &mut Buffer, area: Rect) {
532 let [plot, axis, labels] = Layout::vertical([
533 Constraint::Min(1),
534 Constraint::Length(1),
535 Constraint::Length(1),
536 ])
537 .areas(area);
538 let [gutter, chart] =
539 Layout::horizontal([Constraint::Length(GUTTER), Constraint::Min(1)]).areas(plot);
540 if chart.width == 0 || chart.height == 0 {
541 return;
542 }
543 let nbins = self.counts.len();
544 let bw = (chart.width / nbins.max(1) as u16).clamp(1, 4);
545 let x_of = |k: i32| -> Option<u16> {
546 let i = k - self.kmin;
547 (i >= 0 && (i as usize) < nbins)
548 .then(|| chart.x + i as u16 * bw)
549 .filter(|&x| x < chart.right())
550 };
551
552 let height = |c: T| self.y_scale.apply(c.bar().max(0.0));
553 let tallest = self.counts.iter().map(|&c| height(c)).fold(0.0, f64::max);
554 let max_h = self
555 .y_max
556 .map_or(tallest, |m| self.y_scale.apply(m.max(0.0)).max(tallest));
557 let cells = chart.height as usize * 8;
558 let eighths = |c: T| {
559 if c.bar() <= 0.0 || max_h <= 0.0 {
560 0
561 } else {
562 ((height(c) / max_h * cells as f64).round() as usize).clamp(1, cells)
563 }
564 };
565
566 if let Some(x) = self.pointer.and_then(x_of) {
567 for y in chart.top()..chart.bottom() {
568 put(buf, x, y, "┊", ACCENTED);
569 }
570 }
571
572 let mut bars = |counts: &[T], behind: Option<&[T]>, dim: bool| {
573 for (i, &c) in counts.iter().enumerate() {
574 let x0 = chart.x + i as u16 * bw;
575 if x0 >= chart.right() {
576 break;
577 }
578 let style = if dim {
579 DIM
580 } else {
581 (self.style)(self.kmin + i as i32)
582 };
583 let (top, under) = (eighths(c), behind.map_or(0, |b| eighths(b[i])));
584 for (j, y) in (chart.top()..chart.bottom()).rev().enumerate() {
585 let mut fill = top.saturating_sub(j * 8).min(8);
586 if fill == 0 {
587 break;
588 }
589 if under >= (j + 1) * 8 {
593 fill = 8;
594 }
595 for x in x0..(x0 + bw).min(chart.right()) {
596 put(buf, x, y, EIGHTHS[fill - 1], style);
597 }
598 }
599 }
600 };
601 bars(self.counts, None, self.subset.is_some());
602 if let Some(subset) = self.subset {
603 bars(subset, Some(self.counts), false);
604 }
605
606 let gx = gutter.right() - 1;
608 for y in gutter.top()..gutter.bottom() {
609 put(buf, gx, y, "│", DIM);
610 }
611 let mut ylabel = |y: u16, v: f64| {
612 let s = compact(v);
613 let x = gx.saturating_sub(1 + s.len() as u16).max(gutter.x);
614 buf.set_string(x, y, &s, DIM);
615 put(buf, gx, y, "┤", DIM);
616 };
617 if max_h > 0.0 {
618 ylabel(gutter.top(), self.y_scale.invert(max_h));
619 if gutter.height >= 6 {
620 ylabel(
621 gutter.top() + gutter.height / 2,
622 self.y_scale.invert(max_h / 2.0),
623 );
624 }
625 }
626
627 for x in axis.left()..axis.right() {
629 let sym = match x.cmp(&gx) {
630 std::cmp::Ordering::Less => " ",
631 std::cmp::Ordering::Equal => "└",
632 std::cmp::Ordering::Greater => "─",
633 };
634 put(buf, x, axis.y, sym, DIM);
635 }
636 let every = self
637 .tick_every
638 .unwrap_or_else(|| self.bins.tick_every(nbins))
639 .max(1);
640 let mut next_free = labels.x;
641 let kmax = self.kmin + nbins as i32 - 1;
642 for k in (self.kmin..=kmax).filter(|k| k % every == 0) {
643 let Some(x) = x_of(k) else { continue };
644 let s = match self.x_label {
645 Some(label) => match label(k) {
646 Some(s) => s,
647 None => continue,
648 },
649 None => compact(self.bins.tick_value(k)),
650 };
651 put(buf, x, axis.y, "┴", DIM);
652 if x >= next_free && x + (s.len() as u16) <= labels.right() {
653 buf.set_string(x, labels.y, &s, DIM);
654 next_free = x + s.len() as u16 + 1;
655 }
656 }
657 let pointer = self.pointer.map(|k| (k, "▲", HIGHLIGHT));
658 for &(k, sym, style) in self.marks.iter().chain(pointer.iter()) {
659 if let Some(x) = x_of(k) {
660 put(buf, x, axis.y, sym, style);
661 }
662 }
663 }
664}
665
666pub struct MirrorSide<'a, T: BarValue = f64> {
668 pub counts: &'a [T],
669 pub subset: Option<&'a [T]>,
672 pub style: Style,
673 pub name: &'a str,
675}
676
677pub struct MirrorPlot<'a, T: BarValue = f64> {
682 pub up: MirrorSide<'a, T>,
683 pub down: MirrorSide<'a, T>,
684 pub y_scale: Scale,
685 pub y_max: Option<f64>,
687 pub y_labels: Option<[String; 3]>,
690 pub pointer: Option<usize>,
692 pub x_label: Option<&'a dyn Fn(usize) -> Option<String>>,
695}
696
697impl<T: BarValue> MirrorPlot<'_, T> {
698 pub fn render(&self, buf: &mut Buffer, area: Rect) {
702 let [plot, axis, labels] = Layout::vertical([
703 Constraint::Min(1),
704 Constraint::Length(1),
705 Constraint::Length(1),
706 ])
707 .areas(area);
708 let [gutter, chart] =
709 Layout::horizontal([Constraint::Length(GUTTER), Constraint::Min(1)]).areas(plot);
710 if chart.width == 0 || chart.height < 3 {
711 return;
712 }
713 let half = (chart.height - 1) / 2;
714 let zero = chart.top() + half;
715 let x_of = |i: usize| Some(chart.x + i as u16).filter(|&x| x < chart.right());
716
717 let height = |c: T| self.y_scale.apply(c.bar().max(0.0));
718 let all = self.up.counts.iter().chain(self.down.counts);
719 let tallest = all.map(|&c| height(c)).fold(0.0, f64::max);
720 let max_h = self
721 .y_max
722 .map_or(tallest, |m| self.y_scale.apply(m.max(0.0)).max(tallest));
723 let cells = half as usize * 2;
724 let halves = |c: T| {
725 if c.bar() <= 0.0 || max_h <= 0.0 {
726 0
727 } else {
728 ((height(c) / max_h * cells as f64).round() as usize).clamp(1, cells)
729 }
730 };
731
732 if let Some(x) = self.pointer.and_then(x_of) {
733 for y in chart.top()..chart.top() + 2 * half + 1 {
734 put(buf, x, y, "┊", ACCENTED);
735 }
736 }
737 for x in chart.left()..chart.right() {
738 put(buf, x, zero, "─", DIM);
739 }
740 for (side, up) in [(&self.up, true), (&self.down, false)] {
741 let (whole, part) = if up { ("█", "▄") } else { ("█", "▀") };
742 let mut bars = |counts: &[T], behind: Option<&[T]>, style: Style| {
743 for (i, &c) in counts.iter().enumerate() {
744 let Some(x) = x_of(i) else { break };
745 let (top, under) = (halves(c), behind.map_or(0, |b| halves(b[i])));
746 for k in 0..top.div_ceil(2) {
747 let y = if up {
748 zero - 1 - k as u16
749 } else {
750 zero + 1 + k as u16
751 };
752 let full = 2 * k + 2 <= top || under >= 2 * k + 2;
755 put(buf, x, y, if full { whole } else { part }, style);
756 }
757 }
758 };
759 match side.subset {
760 Some(subset) => {
761 bars(side.counts, None, DIM);
762 bars(subset, Some(side.counts), side.style);
763 }
764 None => bars(side.counts, None, side.style),
765 }
766 }
767 buf.set_string(chart.x, chart.top(), self.up.name, DIM);
768 buf.set_string(chart.x, chart.top() + 2 * half, self.down.name, DIM);
769
770 let gx = gutter.right() - 1;
772 for y in gutter.top()..gutter.bottom() {
773 put(buf, gx, y, "│", DIM);
774 }
775 let own = || {
776 let top = compact(self.y_scale.invert(max_h));
777 [top.clone(), "0".to_string(), top]
778 };
779 let ys = [chart.top(), zero, chart.top() + 2 * half];
780 for (y, s) in ys
781 .into_iter()
782 .zip(self.y_labels.clone().unwrap_or_else(own))
783 {
784 let x = gx.saturating_sub(1 + s.len() as u16).max(gutter.x);
785 buf.set_string(x, y, &s, DIM);
786 put(buf, gx, y, "┤", DIM);
787 }
788
789 for x in axis.left()..axis.right() {
791 let sym = match x.cmp(&gx) {
792 std::cmp::Ordering::Less => " ",
793 std::cmp::Ordering::Equal => "└",
794 std::cmp::Ordering::Greater => "─",
795 };
796 put(buf, x, axis.y, sym, DIM);
797 }
798 let n = self.up.counts.len().max(self.down.counts.len());
799 let mut next_free = labels.x;
800 for i in 0..n {
801 let Some(x) = x_of(i) else { break };
802 let Some(s) = self.x_label.and_then(|label| label(i)) else {
803 continue;
804 };
805 put(buf, x, axis.y, "┴", DIM);
806 if x >= next_free && x + (s.len() as u16) <= labels.right() {
807 buf.set_string(x, labels.y, &s, DIM);
808 next_free = x + s.len() as u16 + 1;
809 }
810 }
811 if let Some(x) = self.pointer.and_then(x_of) {
812 put(buf, x, axis.y, "▲", HIGHLIGHT);
813 }
814 }
815}
816
817#[cfg(test)]
818#[path = "tests/ui.rs"]
819mod tests;