plotpy 1.23.0

use super::{generate_nested_list_3, matrix_to_array, AsMatrix, GraphMaker};
use num_traits::Num;
use std::fmt::Write;

/// Generates an image plot (imshow)
///
/// Uses Matplotlib's [imshow](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html))
///
/// Displays a 2D matrix as a heatmap with configurable colormap.
/// Supports both scalar data (via [`draw`](Self::draw)) and RGB/RGBA channel data
/// (via [`draw_rgb_or_rgba`](Self::draw_rgb_or_rgba)).
///
/// # Examples
///
/// ```
/// use plotpy::{Image, Plot, StrError};
///
/// fn main() -> Result<(), StrError> {
///     // set values
///     let data = [
///         [0.8, 2.4, 2.5, 3.9, 0.0, 4.0, 0.0],
///         [2.4, 0.0, 4.0, 1.0, 2.7, 0.0, 0.0],
///         [1.1, 2.4, 0.8, 4.3, 1.9, 4.4, 0.0],
///         [0.6, 0.0, 0.3, 0.0, 3.1, 0.0, 0.0],
///         [0.7, 1.7, 0.6, 2.6, 2.2, 6.2, 0.0],
///         [1.3, 1.2, 0.0, 0.0, 0.0, 3.2, 5.1],
///         [0.1, 2.0, 0.0, 1.4, 0.0, 1.9, 6.3],
///     ];
///
///     // image plot and options
///     let mut img = Image::new();
///     img.set_colormap_name("hsv").draw(&data);
///
///     // save figure
///     let mut plot = Plot::new();
///     plot.add(&img);
///     plot.save("/tmp/plotpy/doc_tests/doc_image_1.svg")?;
///     Ok(())
/// }
/// ```
///
/// ![doc_image_1.svg](https://raw.githubusercontent.com/cpmech/plotpy/main/figures/doc_image_1.svg)
///
/// See also integration test in the **tests** directory.
pub struct Image {
    colormap_name: String, // Colormap name
    extra: String,         // Extra commands (comma separated)
    buffer: String,        // buffer
}

impl Image {
    /// Creates a new Image object
    pub fn new() -> Self {
        Image {
            colormap_name: String::new(),
            extra: String::new(),
            buffer: String::new(),
        }
    }

    /// (imshow) Displays scalar data as a colored image / heatmap
    ///
    /// # Input
    ///
    /// * `data` -- 2D matrix of scalar values (rows × columns)
    ///
    /// See also: [`draw_rgb_or_rgba`](Self::draw_rgb_or_rgba) for RGB/RGBA images
    pub fn draw<'a, T, U>(&mut self, data: &'a T)
    where
        T: AsMatrix<'a, U>,
        U: 'a + std::fmt::Display + Num,
    {
        matrix_to_array(&mut self.buffer, "data", data);
        let opt = self.options();
        write!(&mut self.buffer, "plt.imshow(data{})\n", &opt).unwrap();
    }

    /// (imshow) Displays data as an image with RGB or RGB(A) values
    ///
    /// # Input
    ///
    /// * `data` - 3D vector with shape (height, width, 3) for RGB or (height, width, 4) for RGBA
    ///   The inner-most vector contains the color channels.
    ///
    /// See <https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html>
    pub fn draw_rgb_or_rgba<T>(&mut self, data: &Vec<Vec<Vec<T>>>)
    where
        T: std::fmt::Display + Num,
    {
        generate_nested_list_3(&mut self.buffer, "data", data);
        let opt = self.options();
        write!(&mut self.buffer, "plt.imshow(data{})\n", &opt).unwrap();
    }

    /// Sets a colormap by index from a built-in list
    ///
    /// Indices wrap around:
    ///
    /// | Index | Colormap |
    /// |-------|----------|
    /// | 0     | bwr      |
    /// | 1     | RdBu     |
    /// | 2     | hsv      |
    /// | 3     | jet      |
    /// | 4     | terrain  |
    /// | 5     | pink     |
    /// | 6     | Greys    |
    ///
    /// See also: [`set_colormap_name`](Self::set_colormap_name)
    pub fn set_colormap_index(&mut self, index: usize) -> &mut Self {
        const CMAP: [&str; 7] = ["bwr", "RdBu", "hsv", "jet", "terrain", "pink", "Greys"];
        self.colormap_name = CMAP[index % 7].to_string();
        self
    }

    /// Sets the colormap by name (e.g., `"terrain"`, `"jet"`, `"viridis"`)
    ///
    /// See the [Matplotlib colormap reference](https://matplotlib.org/stable/tutorials/colors/colormaps.html)
    /// for all available options.
    ///
    /// See also: [`set_colormap_index`](Self::set_colormap_index)
    pub fn set_colormap_name(&mut self, name: &str) -> &mut Self {
        self.colormap_name = String::from(name);
        self
    }

    /// Sets extra matplotlib commands (comma separated)
    ///
    /// **Important:** The extra commands must be comma separated. For example:
    ///
    /// ```text
    /// param1=123,param2='hello'
    /// ```
    ///
    /// [See Matplotlib's documentation](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html)
    pub fn set_extra(&mut self, extra: &str) -> &mut Self {
        self.extra = extra.to_string();
        self
    }

    /// Returns options for image
    fn options(&self) -> String {
        let mut opt = String::new();
        if self.colormap_name != "" {
            write!(&mut opt, ",cmap=plt.get_cmap('{}')", self.colormap_name).unwrap();
        }
        if self.extra != "" {
            write!(&mut opt, ",{}", self.extra).unwrap();
        }
        opt
    }
}

impl GraphMaker for Image {
    fn get_buffer<'a>(&'a self) -> &'a String {
        &self.buffer
    }
    fn clear_buffer(&mut self) {
        self.buffer.clear();
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#[cfg(test)]
mod tests {
    use super::Image;
    use crate::GraphMaker;

    #[test]
    fn new_works() {
        let img = Image::new();
        assert_eq!(img.colormap_name.len(), 0);
        assert_eq!(img.extra.len(), 0);
        assert_eq!(img.buffer.len(), 0);
    }

    #[test]
    fn draw_works_1() {
        let xx = [[1, 2], [3, 2]];
        let mut img = Image::new();
        img.set_colormap_index(0).set_colormap_name("terrain").draw(&xx);
        let b: &str = "data=np.array([[1,2,],[3,2,],])\n\
                       plt.imshow(data,cmap=plt.get_cmap('terrain'))\n";
        assert_eq!(img.buffer, b);
        img.clear_buffer();
        assert_eq!(img.buffer, "");
    }
}