imodfile 0.1.0

//! PyO3 bindings for the imodfile library.
//!
//! Provides Python classes that mirror the Rust data model with full
//! read/write access.  Point data is exchanged as NumPy arrays.

use std::sync::{Arc, Mutex};

use numpy::PyReadonlyArray2;
use pyo3::prelude::*;

use crate::model::*;

// ═════════════════════════════════════════════════════════════════════════════
//  Helper: convert between Vec<Ipoint> and numpy (N, 3) float32 arrays
// ═════════════════════════════════════════════════════════════════════════════

/// Build an (N, 3) float32 array from a slice of Ipoints.
fn points_to_numpy<'py>(py: Python<'py>, pts: &[Ipoint]) -> Bound<'py, numpy::PyArray2<f32>> {
    let rows: Vec<Vec<f32>> = pts.iter().map(|p| vec![p.x, p.y, p.z]).collect();
    numpy::PyArray::from_vec2(py, &rows).expect("shape (N,3)")
}

// ─── Model ───────────────────────────────────────────────────────────────────

/// An IMOD model — the root container for objects, views, and metadata.
///
/// ```python
/// model = imodfile.load("input.mod")
/// print(model.name, model.image_size)
/// model.save("output.mod")
/// ```
#[pyclass(name = "Model", skip_from_py_object)]
#[derive(Clone)]
pub struct PyModel {
    pub inner: Arc<Mutex<Imod>>,
}

#[pymethods]
impl PyModel {
    #[new]
    fn new() -> Self {
        PyModel {
            inner: Arc::new(Mutex::new(Imod::default())),
        }
    }

    /// Load a model from a file (auto-detects binary vs ASCII).
    #[staticmethod]
    fn load(path: &str) -> PyResult<Self> {
        let imod = crate::Imod::load(path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
        Ok(PyModel {
            inner: Arc::new(Mutex::new(imod)),
        })
    }

    /// Save to a file (binary unless path ends with .txt/.ascii).
    fn save(&self, path: &str) -> PyResult<()> {
        let guard = self.inner.lock().unwrap();
        guard
            .save(path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    }

    /// Write all contour points to a text file.
    fn save_points(&self, path: &str) -> PyResult<()> {
        let guard = self.inner.lock().unwrap();
        guard
            .save_points(path)
            .map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))
    }

    // ── Properties ──

    #[getter]
    fn name(&self) -> PyResult<String> {
        Ok(self.inner.lock().unwrap().name.clone())
    }
    #[setter]
    fn set_name(&self, val: String) -> PyResult<()> {
        self.inner.lock().unwrap().name = val;
        Ok(())
    }

    #[getter]
    fn image_size(&self) -> PyResult<(i32, i32, i32)> {
        let g = self.inner.lock().unwrap();
        Ok((g.xmax, g.ymax, g.zmax))
    }
    #[setter]
    fn set_image_size(&self, val: (i32, i32, i32)) -> PyResult<()> {
        let mut g = self.inner.lock().unwrap();
        g.xmax = val.0;
        g.ymax = val.1;
        g.zmax = val.2;
        Ok(())
    }

    #[getter]
    fn pixel_size(&self) -> PyResult<f32> {
        Ok(self.inner.lock().unwrap().pixsize)
    }
    #[setter]
    fn set_pixel_size(&self, val: f32) -> PyResult<()> {
        self.inner.lock().unwrap().pixsize = val;
        Ok(())
    }

    #[getter]
    fn objects(&self) -> PyResult<Vec<PyObject>> {
        let g = self.inner.lock().unwrap();
        Ok((0..g.obj.len())
            .map(|i| PyObject {
                model: self.inner.clone(),
                index: i,
            })
            .collect())
    }

    /// Number of objects.
    fn __len__(&self) -> PyResult<usize> {
        Ok(self.inner.lock().unwrap().obj.len())
    }

    /// Add a new empty object.
    fn add_object(&self) -> PyResult<PyObject> {
        let mut g = self.inner.lock().unwrap();
        g.obj.push(Iobj::default());
        g.objsize = g.obj.len() as i32;
        let index = g.obj.len() - 1;
        Ok(PyObject {
            model: self.inner.clone(),
            index,
        })
    }

    /// Remove an object by index.
    fn remove_object(&self, index: usize) -> PyResult<()> {
        let mut g = self.inner.lock().unwrap();
        if index >= g.obj.len() {
            return Err(pyo3::exceptions::PyIndexError::new_err(format!(
                "object index {index} out of range"
            )));
        }
        g.obj.remove(index);
        g.objsize = g.obj.len() as i32;
        Ok(())
    }

    /// Return all points as a single (N, 3) numpy array.
    fn points<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, numpy::PyArray2<f32>>> {
        let g = self.inner.lock().unwrap();
        let all: Vec<Ipoint> = g
            .obj
            .iter()
            .flat_map(|o| o.cont.iter().flat_map(|c| c.pts.iter().copied()))
            .collect();
        Ok(points_to_numpy(py, &all))
    }

    fn __repr__(&self) -> PyResult<String> {
        let g = self.inner.lock().unwrap();
        Ok(format!(
            "Model(name={:?}, objects={})",
            g.name,
            g.obj.len()
        ))
    }
}

// ─── Object ──────────────────────────────────────────────────────────────────

/// A model object with contours, meshes, colour, and material properties.
#[pyclass(name = "Object", skip_from_py_object)]
#[derive(Clone)]
pub struct PyObject {
    model: Arc<Mutex<Imod>>,
    index: usize,
}

#[pymethods]
impl PyObject {
    #[getter]
    fn name(&self) -> PyResult<String> {
        Ok(self.model.lock().unwrap().obj[self.index].name.clone())
    }
    #[setter]
    fn set_name(&self, val: String) -> PyResult<()> {
        self.model.lock().unwrap().obj[self.index].name = val;
        Ok(())
    }

    #[getter]
    fn color(&self) -> PyResult<(f32, f32, f32)> {
        let g = self.model.lock().unwrap();
        let o = &g.obj[self.index];
        Ok((o.red, o.green, o.blue))
    }
    #[setter]
    fn set_color(&self, val: (f32, f32, f32)) -> PyResult<()> {
        let mut m = self.model.lock().unwrap();
        let o = &mut m.obj[self.index];
        o.red = val.0;
        o.green = val.1;
        o.blue = val.2;
        Ok(())
    }

    #[getter]
    fn flags(&self) -> PyResult<u32> {
        Ok(self.model.lock().unwrap().obj[self.index].flags)
    }
    #[setter]
    fn set_flags(&self, val: u32) -> PyResult<()> {
        self.model.lock().unwrap().obj[self.index].flags = val;
        Ok(())
    }

    #[getter]
    fn contours(&self) -> PyResult<Vec<PyContour>> {
        let g = self.model.lock().unwrap();
        let obj = &g.obj[self.index];
        Ok((0..obj.cont.len())
            .map(|ci| PyContour {
                model: self.model.clone(),
                obj_index: self.index,
                index: ci,
            })
            .collect())
    }

    #[getter]
    fn meshes(&self) -> PyResult<Vec<PyMesh>> {
        let g = self.model.lock().unwrap();
        let obj = &g.obj[self.index];
        Ok((0..obj.mesh.len())
            .map(|mi| PyMesh {
                model: self.model.clone(),
                obj_index: self.index,
                index: mi,
            })
            .collect())
    }

    fn __len__(&self) -> PyResult<usize> {
        Ok(self.model.lock().unwrap().obj[self.index].cont.len())
    }

    /// Add a new empty contour.
    fn add_contour(&self) -> PyResult<PyContour> {
        let mut g = self.model.lock().unwrap();
        g.obj[self.index].cont.push(Icont::default());
        g.obj[self.index].contsize = g.obj[self.index].cont.len() as i32;
        let ci = g.obj[self.index].cont.len() - 1;
        Ok(PyContour {
            model: self.model.clone(),
            obj_index: self.index,
            index: ci,
        })
    }

    /// Remove a contour by index.
    fn remove_contour(&self, ci: usize) -> PyResult<()> {
        let mut g = self.model.lock().unwrap();
        let obj = &mut g.obj[self.index];
        if ci >= obj.cont.len() {
            return Err(pyo3::exceptions::PyIndexError::new_err(format!(
                "contour index {ci} out of range"
            )));
        }
        obj.cont.remove(ci);
        obj.contsize = obj.cont.len() as i32;
        Ok(())
    }

    /// Add a new empty mesh.
    fn add_mesh(&self) -> PyResult<PyMesh> {
        let mut g = self.model.lock().unwrap();
        g.obj[self.index].mesh.push(Imesh::default());
        g.obj[self.index].meshsize = g.obj[self.index].mesh.len() as i32;
        let mi = g.obj[self.index].mesh.len() - 1;
        Ok(PyMesh {
            model: self.model.clone(),
            obj_index: self.index,
            index: mi,
        })
    }

    /// Return contour points as a single (N, 3) numpy array.
    fn points<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, numpy::PyArray2<f32>>> {
        let g = self.model.lock().unwrap();
        let all: Vec<Ipoint> = g.obj[self.index]
            .cont
            .iter()
            .flat_map(|c| c.pts.iter().copied())
            .collect();
        Ok(points_to_numpy(py, &all))
    }

    fn __repr__(&self) -> PyResult<String> {
        let g = self.model.lock().unwrap();
        let o = &g.obj[self.index];
        Ok(format!(
            "Object(name={:?}, contours={}, meshes={})",
            o.name,
            o.cont.len(),
            o.mesh.len()
        ))
    }
}

// ─── Contour ────────────────────────────────────────────────────────────────

/// A contour — a series of 3-D points with optional surface/time metadata.
#[pyclass(name = "Contour", skip_from_py_object)]
#[derive(Clone)]
pub struct PyContour {
    model: Arc<Mutex<Imod>>,
    obj_index: usize,
    index: usize,
}

#[pymethods]
impl PyContour {
    #[getter]
    fn points<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, numpy::PyArray2<f32>>> {
        let g = self.model.lock().unwrap();
        let pts = &g.obj[self.obj_index].cont[self.index].pts;
        Ok(points_to_numpy(py, pts))
    }

    #[setter]
    fn set_points(&self, arr: PyReadonlyArray2<f32>) -> PyResult<()> {
        let view = arr.as_array();
        let shape = view.shape();
        if shape.len() != 2 || shape[1] != 3 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "expected array of shape (N, 3)",
            ));
        }
        let n = shape[0];
        let pts: Vec<Ipoint> = (0..n)
            .map(|i| Ipoint {
                x: view[[i, 0]],
                y: view[[i, 1]],
                z: view[[i, 2]],
            })
            .collect();
        let mut g = self.model.lock().unwrap();
        let cont = &mut g.obj[self.obj_index].cont[self.index];
        let psize = pts.len() as i32;
        cont.pts = pts;
        cont.psize = psize;
        Ok(())
    }

    #[getter]
    fn psize(&self) -> PyResult<i32> {
        Ok(self.model.lock().unwrap().obj[self.obj_index].cont[self.index].psize)
    }

    #[getter]
    fn surface(&self) -> PyResult<i32> {
        Ok(self.model.lock().unwrap().obj[self.obj_index].cont[self.index].surf)
    }
    #[setter]
    fn set_surface(&self, val: i32) -> PyResult<()> {
        self.model.lock().unwrap().obj[self.obj_index].cont[self.index].surf = val;
        Ok(())
    }

    #[getter]
    fn time(&self) -> PyResult<i32> {
        Ok(self.model.lock().unwrap().obj[self.obj_index].cont[self.index].time)
    }
    #[setter]
    fn set_time(&self, val: i32) -> PyResult<()> {
        self.model.lock().unwrap().obj[self.obj_index].cont[self.index].time = val;
        Ok(())
    }

    #[getter]
    fn flags(&self) -> PyResult<u32> {
        Ok(self.model.lock().unwrap().obj[self.obj_index].cont[self.index].flags)
    }
    #[setter]
    fn set_flags(&self, val: u32) -> PyResult<()> {
        self.model.lock().unwrap().obj[self.obj_index].cont[self.index].flags = val;
        Ok(())
    }

    fn __repr__(&self) -> PyResult<String> {
        let g = self.model.lock().unwrap();
        let c = &g.obj[self.obj_index].cont[self.index];
        Ok(format!("Contour(points={}, surface={})", c.psize, c.surf))
    }
}

// ─── Mesh ────────────────────────────────────────────────────────────────────

/// A mesh with vertices and an index list.
#[pyclass(name = "Mesh", skip_from_py_object)]
#[derive(Clone)]
pub struct PyMesh {
    model: Arc<Mutex<Imod>>,
    obj_index: usize,
    index: usize,
}

#[pymethods]
impl PyMesh {
    #[getter]
    fn vertices<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, numpy::PyArray2<f32>>> {
        let g = self.model.lock().unwrap();
        let pts = &g.obj[self.obj_index].mesh[self.index].vert;
        Ok(points_to_numpy(py, pts))
    }

    #[setter]
    fn set_vertices(&self, arr: PyReadonlyArray2<f32>) -> PyResult<()> {
        let view = arr.as_array();
        let shape = view.shape();
        if shape.len() != 2 || shape[1] != 3 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "expected array of shape (N, 3)",
            ));
        }
        let n = shape[0];
        let verts: Vec<Ipoint> = (0..n)
            .map(|i| Ipoint {
                x: view[[i, 0]],
                y: view[[i, 1]],
                z: view[[i, 2]],
            })
            .collect();
        let mut g = self.model.lock().unwrap();
        let m = &mut g.obj[self.obj_index].mesh[self.index];
        let vsize = verts.len() as i32;
        m.vert = verts;
        m.vsize = vsize;
        Ok(())
    }

    #[getter]
    fn indices(&self) -> PyResult<Vec<i32>> {
        Ok(self.model.lock().unwrap().obj[self.obj_index].mesh[self.index].list.clone())
    }
    #[setter]
    fn set_indices(&self, val: Vec<i32>) -> PyResult<()> {
        let mut g = self.model.lock().unwrap();
        let m = &mut g.obj[self.obj_index].mesh[self.index];
        let lsize = val.len() as i32;
        m.list = val;
        m.lsize = lsize;
        Ok(())
    }

    fn __repr__(&self) -> PyResult<String> {
        let g = self.model.lock().unwrap();
        let m = &g.obj[self.obj_index].mesh[self.index];
        Ok(format!("Mesh(vertices={}, indices={})", m.vsize, m.lsize))
    }
}

// ═════════════════════════════════════════════════════════════════════════════
//  Module registration
// ═════════════════════════════════════════════════════════════════════════════

/// IMOD model file decoder/encoder — Python bindings.
///
/// ```python
/// import imodfile
///
/// # Load a model
/// model = imodfile.load("cells.mod")
/// print(model.name)
///
/// # Read objects and their contours
/// for obj in model.objects:
///     print(obj.name)
///     for cont in obj.contours:
///         pts = cont.points  # (N, 3) numpy array
///
/// # Create new content
/// obj = model.add_object()
/// obj.name = "mitochondria"
/// cont = obj.add_contour()
/// cont.points = numpy.array([[0,0,0],[1,0,0],[0,1,0]], dtype=numpy.float32)
///
/// model.save("output.mod")
/// ```
#[pymodule]
fn imodfile(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyModel>()?;
    m.add_class::<PyObject>()?;
    m.add_class::<PyContour>()?;
    m.add_class::<PyMesh>()?;

    m.add_function(wrap_pyfunction!(load_py, m)?)?;

    Ok(())
}

#[pyfunction(name = "load")]
fn load_py(path: &str) -> PyResult<PyModel> {
    PyModel::load(path)
}