plotiron/
figure.rs

1//! Figure management and SVG generation
2
3use crate::axes::Axes;
4use crate::colors::Color;
5
6/// Represents a figure that can contain multiple subplots
7#[derive(Debug)]
8pub struct Figure {
9    pub width: f64,
10    pub height: f64,
11    pub dpi: f64,
12    pub background_color: Color,
13    pub subplots: Vec<Axes>,
14    pub tight_layout: bool,
15}
16
17impl Figure {
18    /// Create a new figure with default settings
19    pub fn new() -> Self {
20        Figure {
21            width: 1200.0,
22            height: 900.0,
23            dpi: 100.0,
24            background_color: Color::WHITE,
25            subplots: Vec::new(),
26            tight_layout: true,
27        }
28    }
29
30    /// Create a new figure with specified dimensions
31    pub fn with_size(width: f64, height: f64) -> Self {
32        Figure {
33            width,
34            height,
35            dpi: 100.0,
36            background_color: Color::WHITE,
37            subplots: Vec::new(),
38            tight_layout: true,
39        }
40    }
41
42    /// Set the figure size
43    pub fn set_size(&mut self, width: f64, height: f64) -> &mut Self {
44        self.width = width;
45        self.height = height;
46        self
47    }
48
49    /// Set the DPI (dots per inch)
50    pub fn set_dpi(&mut self, dpi: f64) -> &mut Self {
51        self.dpi = dpi;
52        self
53    }
54
55    /// Set the background color
56    pub fn set_facecolor(&mut self, color: Color) -> &mut Self {
57        self.background_color = color;
58        self
59    }
60
61    /// Add a subplot and return a mutable reference to it
62    pub fn add_subplot(&mut self) -> &mut Axes {
63        let axes = Axes::new();
64        self.subplots.push(axes);
65        self.subplots.last_mut().unwrap()
66    }
67
68    /// Add a subplot with DOT graph content
69    pub fn add_dot_subplot(&mut self, dot_content: &str) -> Result<&mut Axes, String> {
70        self.add_dot_subplot_with_layout(dot_content, crate::dot::LayoutAlgorithm::Hierarchical)
71    }
72
73    /// Add a DOT subplot with specified layout algorithm
74    pub fn add_dot_subplot_with_layout(
75        &mut self,
76        dot_content: &str,
77        layout: crate::dot::LayoutAlgorithm,
78    ) -> Result<&mut Axes, String> {
79        let axes = self.add_subplot();
80
81        // Parse DOT content using the advanced renderer
82        let mut dot_graph = crate::dot::DotGraph::parse_dot(dot_content)?;
83        dot_graph.set_layout(layout);
84        dot_graph.apply_layout();
85
86        // Render the graph to the axes
87        dot_graph.render_to_axes(axes);
88
89        Ok(axes)
90    }
91
92    /// Get a mutable reference to a subplot by index
93    pub fn subplot(&mut self, index: usize) -> Option<&mut Axes> {
94        self.subplots.get_mut(index)
95    }
96
97    /// Generate SVG string for the entire figure
98    pub fn to_svg(&self) -> String {
99        let mut svg = String::new();
100
101        // SVG header
102        svg.push_str(&format!(
103            "<svg width=\"{}\" height=\"{}\" xmlns=\"http://www.w3.org/2000/svg\">\n",
104            self.width, self.height
105        ));
106
107        // Background
108        svg.push_str(&format!(
109            "<rect width=\"{}\" height=\"{}\" fill=\"{}\" />\n",
110            self.width,
111            self.height,
112            self.background_color.to_svg_string()
113        ));
114
115        // Render subplots
116        if self.subplots.len() == 1 {
117            // Single subplot takes the full figure
118            svg.push_str(&self.subplots[0].to_svg(self.width, self.height));
119        } else if !self.subplots.is_empty() {
120            // Multiple subplots - simple grid layout
121            let cols = (self.subplots.len() as f64).sqrt().ceil() as usize;
122            let rows = (self.subplots.len() + cols - 1) / cols;
123
124            let subplot_width = self.width / cols as f64;
125            let subplot_height = self.height / rows as f64;
126
127            for (i, subplot) in self.subplots.iter().enumerate() {
128                let col = i % cols;
129                let row = i / cols;
130                let x = col as f64 * subplot_width;
131                let y = row as f64 * subplot_height;
132
133                svg.push_str(&format!("<g transform=\"translate({},{})\">\n", x, y));
134                svg.push_str(&subplot.to_svg(subplot_width, subplot_height));
135                svg.push_str("</g>\n");
136            }
137        }
138
139        svg.push_str("</svg>");
140        svg
141    }
142
143    /// Display the figure (prints SVG to stdout for now)
144    pub fn show(&self) {
145        let svg = self.to_svg();
146        crate::viewer::show_svg(svg);
147    }
148
149    /// Clear all subplots
150    pub fn clear(&mut self) {
151        self.subplots.clear();
152    }
153
154    /// Set tight layout
155    pub fn tight_layout(&mut self, enable: bool) -> &mut Self {
156        self.tight_layout = enable;
157        self
158    }
159
160    /// Create a figure from DOT markup language
161    pub fn from_dot(dot_content: &str) -> Result<Self, String> {
162        let mut figure = Figure::new();
163        let axes = figure.add_subplot();
164
165        // Parse DOT content
166        let lines: Vec<&str> = dot_content.lines().collect();
167        let mut nodes = Vec::new();
168        let mut edges = Vec::new();
169
170        for line in lines {
171            let line = line.trim();
172            if line.is_empty()
173                || line.starts_with("//")
174                || line.starts_with("digraph")
175                || line.starts_with("graph")
176                || line == "{"
177                || line == "}"
178            {
179                continue;
180            }
181
182            if line.contains("->") {
183                // Edge definition
184                let parts: Vec<&str> = line.split("->").collect();
185                if parts.len() == 2 {
186                    let from = parts[0].trim().trim_matches('"');
187                    let to = parts[1].trim().trim_end_matches(';').trim_matches('"');
188                    edges.push((from.to_string(), to.to_string()));
189                }
190            } else if line.contains("--") {
191                // Undirected edge
192                let parts: Vec<&str> = line.split("--").collect();
193                if parts.len() == 2 {
194                    let from = parts[0].trim().trim_matches('"');
195                    let to = parts[1].trim().trim_end_matches(';').trim_matches('"');
196                    edges.push((from.to_string(), to.to_string()));
197                }
198            } else if line.contains('[') && line.contains(']') {
199                // Node with attributes
200                let node_name = line.split('[').next().unwrap().trim().trim_matches('"');
201                if !node_name.is_empty() {
202                    nodes.push(node_name.to_string());
203                }
204            } else if line.ends_with(';') {
205                // Simple node definition
206                let node_name = line.trim_end_matches(';').trim().trim_matches('"');
207                if !node_name.is_empty() {
208                    nodes.push(node_name.to_string());
209                }
210            }
211        }
212
213        // Collect all unique nodes from edges
214        for (from, to) in &edges {
215            if !nodes.contains(from) {
216                nodes.push(from.clone());
217            }
218            if !nodes.contains(to) {
219                nodes.push(to.clone());
220            }
221        }
222
223        if nodes.is_empty() {
224            return Err("No nodes found in DOT content".to_string());
225        }
226
227        // Create a simple layout for nodes
228        let node_count = nodes.len();
229        let mut x_coords = Vec::new();
230        let mut y_coords = Vec::new();
231
232        if node_count == 1 {
233            x_coords.push(0.5);
234            y_coords.push(0.5);
235        } else {
236            // Arrange nodes in a circle
237            for i in 0..node_count {
238                let angle = 2.0 * std::f64::consts::PI * i as f64 / node_count as f64;
239                let x = 0.5 + 0.3 * angle.cos();
240                let y = 0.5 + 0.3 * angle.sin();
241                x_coords.push(x);
242                y_coords.push(y);
243            }
244        }
245
246        // Plot nodes as scatter points
247        axes.scatter(x_coords.as_slice(), y_coords.as_slice());
248        if let Some(last_plot) = axes.plots.last_mut() {
249            last_plot.marker = crate::markers::Marker::Circle;
250            last_plot.marker_size = 10.0;
251            last_plot.color = crate::colors::Color::BLUE;
252        }
253
254        // Draw edges as lines
255        for (from, to) in edges {
256            if let (Some(from_idx), Some(to_idx)) = (
257                nodes.iter().position(|n| n == &from),
258                nodes.iter().position(|n| n == &to),
259            ) {
260                let x_line = vec![x_coords[from_idx], x_coords[to_idx]];
261                let y_line = vec![y_coords[from_idx], y_coords[to_idx]];
262                axes.plot(x_line, y_line);
263                if let Some(last_plot) = axes.plots.last_mut() {
264                    last_plot.color = crate::colors::Color::BLACK;
265                    last_plot.line_width = 1.0;
266                }
267            }
268        }
269
270        axes.set_title("Graph from DOT");
271        axes.set_xlabel("X");
272        axes.set_ylabel("Y");
273
274        Ok(figure)
275    }
276}
277
278impl Default for Figure {
279    fn default() -> Self {
280        Self::new()
281    }
282}