use std::collections::HashSet;
use std::sync::Arc;
use rstui::{App, AppContext, AppEvent, ControlFlow, Frame};
use crate::pdb::{Atom, Protein, Vec3};
use crate::raster::{Canvas, Color, mix, scale};
const BACKGROUND_TOP: Color = [5, 9, 17, 255];
const BACKGROUND_BOTTOM: Color = [12, 18, 29, 255];
const TEXT: Color = [220, 230, 241, 255];
const MUTED: Color = [135, 153, 174, 255];
const ACCENT: Color = [73, 215, 196, 255];
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Representation {
Backbone,
Atoms,
Combined,
}
impl Representation {
pub const fn name(self) -> &'static str {
match self {
Self::Backbone => "backbone",
Self::Atoms => "atoms",
Self::Combined => "combined",
}
}
fn next(self) -> Self {
match self {
Self::Backbone => Self::Atoms,
Self::Atoms => Self::Combined,
Self::Combined => Self::Backbone,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ColorScheme {
Chain,
Element,
Confidence,
}
impl ColorScheme {
pub const fn name(self) -> &'static str {
match self {
Self::Chain => "chain",
Self::Element => "element",
Self::Confidence => "pLDDT/B-factor",
}
}
fn next(self) -> Self {
match self {
Self::Chain => Self::Element,
Self::Element => Self::Confidence,
Self::Confidence => Self::Chain,
}
}
}
pub struct Viewer {
protein: Arc<Protein>,
yaw: f32,
pitch: f32,
zoom: f32,
representation: Representation,
color_scheme: ColorScheme,
dragging: Option<(u32, u32)>,
auto_rotate: bool,
last_tick_ms: u64,
}
impl Viewer {
pub fn new(
protein: Protein,
representation: Representation,
color_scheme: ColorScheme,
) -> Self {
Self {
protein: Arc::new(protein),
yaw: -0.72,
pitch: 0.38,
zoom: 1.0,
representation,
color_scheme,
dragging: None,
auto_rotate: false,
last_tick_ms: 0,
}
}
fn reset_view(&mut self) {
self.yaw = -0.72;
self.pitch = 0.38;
self.zoom = 1.0;
self.auto_rotate = false;
}
pub fn snapshot(&self, width: u32, height: u32) -> Frame {
self.render_scene(width.max(1), height.max(1), false)
}
fn render_scene(&self, width: u32, height: u32, fast: bool) -> Frame {
let mut canvas = Canvas::new(width, height, BACKGROUND_TOP);
canvas.vertical_gradient(BACKGROUND_TOP, BACKGROUND_BOTTOM);
let detailed = width >= 280 && height >= 105;
let header = if detailed { 22 } else { 2 };
let footer = if detailed { 22 } else { 2 };
let scene_height = height.saturating_sub(header + footer).max(1);
let scene_center = [
width as f32 * 0.5,
header as f32 + scene_height as f32 * 0.5,
];
let scene_scale = width.min(scene_height) as f32 * 0.43 * self.zoom;
let projected: Vec<Projected> = self
.protein
.atoms
.iter()
.map(|atom| {
let local = (atom.position - self.protein.center) * (1.0 / self.protein.radius);
let rotated = rotate(local, self.yaw, self.pitch);
Projected {
x: scene_center[0] + rotated.x * scene_scale,
y: scene_center[1] - rotated.y * scene_scale,
z: rotated.z,
}
})
.collect();
if matches!(
self.representation,
Representation::Backbone | Representation::Combined
) {
let thickness = (scene_scale * 0.012).clamp(1.3, 7.0);
for segment in &self.protein.trace {
let from = projected[segment.from];
let to = projected[segment.to];
let from_color = self.atom_color(&self.protein.atoms[segment.from], from.z);
let to_color = self.atom_color(&self.protein.atoms[segment.to], to.z);
canvas.line_3d(
[from.x, from.y, from.z],
[to.x, to.y, to.z],
thickness,
from_color,
to_color,
);
}
if self.representation == Representation::Backbone {
let mut seen = HashSet::new();
for segment in &self.protein.trace {
seen.insert(segment.from);
seen.insert(segment.to);
}
let radius = (scene_scale * 0.014).clamp(1.8, 7.0);
for index in seen {
let point = projected[index];
let color = self.atom_color(&self.protein.atoms[index], point.z);
canvas.sphere(point.x, point.y, point.z + 0.001, radius, color);
}
}
}
if matches!(
self.representation,
Representation::Atoms | Representation::Combined
) {
let maximum = if fast { 8_000 } else { 30_000 };
let stride = self.protein.atoms.len().div_ceil(maximum).max(1);
for (index, atom) in self.protein.atoms.iter().enumerate().step_by(stride) {
let point = projected[index];
let physical_radius = element_radius(&atom.element);
let radius = (physical_radius / self.protein.radius * scene_scale)
.clamp(if fast { 1.0 } else { 1.5 }, if fast { 7.0 } else { 12.0 });
let color = self.atom_color(atom, point.z);
canvas.sphere(point.x, point.y, point.z + 0.002, radius, color);
}
}
if detailed {
canvas.fill_rect(0, 0, width, 20, [3, 7, 13, 215]);
canvas.fill_rect(0, height.saturating_sub(20), width, 20, [3, 7, 13, 225]);
canvas.horizontal_line(20, [35, 55, 71, 255]);
canvas.horizontal_line(height.saturating_sub(21), [35, 55, 71, 255]);
let title_capacity = width.saturating_sub(16) as usize / 8;
canvas.text(8, 6, &truncate(&self.protein.title, title_capacity), TEXT);
let status = format!(
"{} atoms {} residues {} chains mode {} color {}{}",
self.protein.atoms.len(),
self.protein.residue_count,
self.protein.chains.len(),
self.representation.name(),
self.color_scheme.name(),
if self.auto_rotate { " AUTO" } else { "" }
);
canvas.text(8, height.saturating_sub(16), &status, ACCENT);
let help = "drag/arrows rotate wheel +/- zoom m mode c color space auto r reset q quit";
let help_x = width.saturating_sub(help.len() as u32 * 8 + 8);
if help_x > status.len() as u32 * 8 + 24 {
canvas.text(help_x, height.saturating_sub(16), help, MUTED);
}
}
Frame::rgba(width, height, canvas.pixels).expect("canvas dimensions are valid")
}
fn atom_color(&self, atom: &Atom, z: f32) -> Color {
let base = match self.color_scheme {
ColorScheme::Chain => chain_color(&atom.chain, &self.protein.chains),
ColorScheme::Element => element_color(&atom.element),
ColorScheme::Confidence => confidence_color(self.protein.confidence(atom.b_factor)),
};
scale(base, 0.82 + (z * 0.5 + 0.5).clamp(0.0, 1.0) * 0.18)
}
}
impl App for Viewer {
fn event(&mut self, event: AppEvent, _context: &AppContext) -> Result<ControlFlow, String> {
let render = match event {
AppEvent::Start | AppEvent::Resize { .. } => true,
AppEvent::Key { key, .. } => match key.as_str() {
"left" | "h" => {
self.yaw -= 0.12;
true
}
"right" | "l" => {
self.yaw += 0.12;
true
}
"up" | "k" => {
self.pitch = (self.pitch + 0.12).clamp(-1.55, 1.55);
true
}
"down" | "j" => {
self.pitch = (self.pitch - 0.12).clamp(-1.55, 1.55);
true
}
"+" | "=" => {
self.zoom = (self.zoom * 1.12).min(8.0);
true
}
"-" | "_" => {
self.zoom = (self.zoom / 1.12).max(0.15);
true
}
"m" => {
self.representation = self.representation.next();
true
}
"c" => {
self.color_scheme = self.color_scheme.next();
true
}
" " | "space" => {
self.auto_rotate = !self.auto_rotate;
true
}
"r" => {
self.reset_view();
true
}
_ => false,
},
AppEvent::Mouse {
kind,
column,
row,
pixel_x,
pixel_y,
button,
} => {
let x = pixel_x.unwrap_or(u32::from(column) * 8);
let y = pixel_y.unwrap_or(u32::from(row) * 16);
match (kind.as_str(), button.as_deref()) {
("down", Some("left")) => {
self.dragging = Some((x, y));
false
}
("drag", Some("left")) => {
if let Some((last_x, last_y)) = self.dragging.replace((x, y)) {
self.yaw += (x as f32 - last_x as f32) * 0.008;
self.pitch = (self.pitch - (y as f32 - last_y as f32) * 0.008)
.clamp(-1.55, 1.55);
true
} else {
false
}
}
("up", Some("left")) => {
self.dragging = None;
false
}
("scroll_up", _) => {
self.zoom = (self.zoom * 1.12).min(8.0);
true
}
("scroll_down", _) => {
self.zoom = (self.zoom / 1.12).max(0.15);
true
}
_ => false,
}
}
AppEvent::Tick { elapsed_ms } if self.auto_rotate => {
let delta = elapsed_ms.saturating_sub(self.last_tick_ms).min(100);
self.last_tick_ms = elapsed_ms;
self.yaw += delta as f32 * 0.00055;
true
}
AppEvent::Tick { elapsed_ms } => {
self.last_tick_ms = elapsed_ms;
false
}
};
Ok(if render {
ControlFlow::Render
} else {
ControlFlow::Continue
})
}
fn render(&mut self, context: &AppContext) -> Result<Frame, String> {
let scale = if context.quality == "fast"
&& matches!(context.protocol.as_str(), "kitty" | "iterm2")
{
0.58
} else {
1.0
};
let width = ((context.geometry.pixel_width as f32 * scale).round() as u32).max(1);
let height = ((context.geometry.pixel_height as f32 * scale).round() as u32).max(1);
Ok(self.render_scene(width, height, context.quality == "fast"))
}
}
#[derive(Clone, Copy)]
struct Projected {
x: f32,
y: f32,
z: f32,
}
fn rotate(point: Vec3, yaw: f32, pitch: f32) -> Vec3 {
let (sin_yaw, cos_yaw) = yaw.sin_cos();
let x = point.x * cos_yaw + point.z * sin_yaw;
let z = -point.x * sin_yaw + point.z * cos_yaw;
let (sin_pitch, cos_pitch) = pitch.sin_cos();
Vec3::new(
x,
point.y * cos_pitch - z * sin_pitch,
point.y * sin_pitch + z * cos_pitch,
)
}
fn chain_color(chain: &str, chains: &[String]) -> Color {
const PALETTE: &[Color] = &[
[73, 215, 196, 255],
[103, 155, 255, 255],
[240, 128, 154, 255],
[245, 192, 83, 255],
[177, 132, 245, 255],
[108, 214, 112, 255],
[244, 139, 70, 255],
[87, 200, 238, 255],
];
let index = chains.iter().position(|item| item == chain).unwrap_or(0);
PALETTE[index % PALETTE.len()]
}
fn element_color(element: &str) -> Color {
match element {
"H" => [235, 239, 245, 255],
"C" => [162, 174, 190, 255],
"N" => [76, 127, 255, 255],
"O" => [244, 82, 82, 255],
"S" => [250, 207, 65, 255],
"P" => [255, 146, 54, 255],
"F" | "CL" => [84, 214, 111, 255],
"BR" => [165, 72, 60, 255],
"I" => [148, 79, 201, 255],
"FE" => [224, 102, 51, 255],
"MG" | "CA" | "ZN" | "MN" | "CU" | "CO" | "NI" => [92, 217, 199, 255],
_ => [225, 120, 205, 255],
}
}
fn confidence_color(value: f32) -> Color {
if value >= 90.0 {
[33, 79, 186, 255]
} else if value >= 70.0 {
mix(
[41, 167, 225, 255],
[33, 79, 186, 255],
(value - 70.0) / 20.0,
)
} else if value >= 50.0 {
mix(
[246, 201, 68, 255],
[41, 167, 225, 255],
(value - 50.0) / 20.0,
)
} else {
mix(
[239, 74, 77, 255],
[246, 201, 68, 255],
value.clamp(0.0, 50.0) / 50.0,
)
}
}
fn element_radius(element: &str) -> f32 {
match element {
"H" => 1.20,
"C" => 1.70,
"N" => 1.55,
"O" => 1.52,
"F" => 1.47,
"P" => 1.80,
"S" => 1.80,
"CL" => 1.75,
"FE" | "MG" | "CA" | "ZN" | "MN" | "CU" | "CO" | "NI" => 1.65,
_ => 1.70,
}
}
fn truncate(value: &str, maximum: usize) -> String {
if value.chars().count() <= maximum {
return value.to_owned();
}
if maximum <= 3 {
return value.chars().take(maximum).collect();
}
let mut shortened: String = value.chars().take(maximum - 3).collect();
shortened.push_str("...");
shortened
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rotation_preserves_length() {
let point = Vec3::new(1.0, 2.0, 3.0);
let rotated = rotate(point, 0.7, -0.4);
assert!((point.length() - rotated.length()).abs() < 0.0001);
}
#[test]
fn confidence_palette_has_expected_extremes() {
assert_eq!(confidence_color(95.0), [33, 79, 186, 255]);
assert_eq!(confidence_color(0.0), [239, 74, 77, 255]);
}
}