# cython: language_level=3
# cython: language=c++
# cython: module_name=pytessel_core
# cython: c_string_type=unicode, c_string_encoding=utf8

from .pytessel_core cimport ScalarField, IsoSurface
from libcpp.string cimport string
from libcpp.memory cimport shared_ptr,make_shared
import numpy as np
import sys
import cython
import numpy.typing as npt

cdef class PyTessel:

    def __cinit__(self):
        pass

    @cython.embedsignature(True)
    def marching_cubes(
        self,
        vector[float] grid,
        vector[size_t] dimensions,
        vector[float] unitcell,
        float isovalue
    ) -> tuple[
        npt.NDArray[np.float64],
        npt.NDArray[np.float64],
        npt.NDArray[np.float64]
    ]:
        """
        Perform marching cubes algorithm to generate isosurface

        Parameters
        ----------
        grid : Iterable of floats
            Scalar field as a flattened array
        dimensions : Iterable of ints
            Dimensions of the scalar field grid (nx, ny, nz)
        unitcell : Iterable of floats
            Unitcell matrix (flattened)
        isovalue : float
            Isovalue of the isosurface
               
        Returns
        -------
        vertices : (Nx3) numpy array of floats
            Triangle vertices
        normals : (Nx3) numpy array of floats
            Triangle normals (at the vertices)
        indices : numpy array of ints
            Triangle indices
        
        Notes
        -----
        * The scalar field needs to be encoded such that the z-coordinate is the slowest moving
          index and the x-coordinate the fastest moving index.
        * The dimensions of the scalar field are encoded in the order :code:`(nx, ny, nz)`.
        * You can use :code:`reversed(grid.shape)` to pass the dimensions.
        * The output of the :code:`marching_cubes` function are three arrays, corresponding to the
          vertices, the normals, and the indices of the isosurface. The isosurface corresponds to
          a number of linked triangles (polygons) whose vertices are stored in the :code:`vertices`
          array. For each vertex, the normal vector is encoded in the :code:`normals` array. Finally,
          the triangles are stored as a triplet of indices in the :code:`indices` array. These indices
          refer to the position in the :code:`vertices` and :code:`normals` array. Because multiple
          triangles can use the same vertices, this is an efficient way to store the isosurface.
        * One rarely needs to perform any operations on the :code:`vertices`, :code:`normals` and
          :code:`indices` arrays. Typically, these arrays are constructed and immediately relayed
          to the :code:`write_ply` function to store them as a file which can be used in another
          program. 
        """
        cdef shared_ptr[ScalarField] scalarfield
        cdef shared_ptr[IsoSurface] isosurface
        cdef shared_ptr[IsoSurfaceMesh] isosurface_mesh

        # build scalar field
        scalarfield = make_shared[ScalarField](grid, dimensions, unitcell)

        # construct isosurface
        isosurface = make_shared[IsoSurface](scalarfield)
        isosurface.get().marching_cubes(isovalue)

        # extract isosurface mesh
        isosurface_mesh = make_shared[IsoSurfaceMesh](scalarfield, isosurface)
        isosurface_mesh.get().construct_mesh(False)

        # extract data
        vertices = np.array(isosurface_mesh.get().get_vertices(), dtype=np.float32).reshape(-1,3)
        normals = np.array(isosurface_mesh.get().get_normals(), dtype=np.float32).reshape(-1,3)
        indices = np.array(isosurface_mesh.get().get_indices(), dtype=np.uint32)

        return vertices, normals, indices

    def write_ply(self,
        filename: str,
        vertices: npt.NDArray[np.float64],
        normals: npt.NDArray[np.float64],
        indices: npt.NDArray[np.uint32],
    ) -> None:
        """
        Write a binary PLY file with vertices, normals, and triangular faces.
        """

        if vertices.shape != normals.shape:
            raise ValueError("vertices and normals must have the same shape")

        if vertices.shape[1] != 3:
            raise ValueError("vertices must be of shape (N, 3)")

        if indices.ndim != 1 or len(indices) % 3 != 0:
            raise ValueError("indices must be a flat array of length multiple of 3")

        n_vertices = vertices.shape[0]
        n_faces = len(indices) // 3

        endian = "binary_little_endian" if sys.byteorder == "little" else "binary_big_endian"

        header = (
            "ply\n"
            f"format {endian} 1.0\n"
            "comment generated by write_ply\n"
            f"element vertex {n_vertices}\n"
            "property float x\n"
            "property float y\n"
            "property float z\n"
            "property float nx\n"
            "property float ny\n"
            "property float nz\n"
            f"element face {n_faces}\n"
            "property list uchar uint vertex_indices\n"
            "end_header\n"
        ).encode("ascii")

        # Interleave vertices and normals: (N, 6)
        vertex_data = np.hstack((vertices, normals)).astype(np.float32, copy=False)

        # Face data: [3, i0, i1, i2]
        face_data = np.empty((n_faces, 4), dtype=np.uint32)
        face_data[:, 0] = 3
        face_data[:, 1:] = indices.reshape(-1, 3)

        # The first element must be uint8, the rest uint32
        face_bytes = np.empty(
            n_faces,
            dtype=[("n", "u1"), ("i", "u4", (3,))]
        )
        face_bytes["n"] = 3
        face_bytes["i"] = indices.reshape(-1, 3)

        with open(filename, "wb") as f:
            f.write(header)
            f.write(vertex_data.tobytes())
            f.write(face_bytes.tobytes())

    def write_stl(self,
        filename: str,
        vertices: npt.NDArray[np.float64],
        normals: npt.NDArray[np.float64],
        indices: npt.NDArray[np.uint32],
    ) -> None:
        """
        Write a binary STL file.

        Parameters
        ----------
        vertices : (N, 3) array of vertex positions
        indices  : (M,) flat array, length multiple of 3
        normals  : (N, 3) vertex normals
        """

        if vertices.shape != normals.shape:
            raise ValueError("vertices and normals must have the same shape")

        if vertices.shape[1] != 3:
            raise ValueError("vertices must be of shape (N, 3)")

        if indices.ndim != 1 or len(indices) % 3 != 0:
            raise ValueError("indices must be a flat array with length multiple of 3")

        # Build triangle vertices: (T, 3, 3)
        triangles = vertices[indices].reshape(-1, 3, 3)

        # Compute face normals from vertex normals
        tri_normals = normals[indices].reshape(-1, 3, 3)
        face_normals = tri_normals.sum(axis=1)

        # Normalize
        norm = np.linalg.norm(face_normals, axis=1)
        norm[norm == 0.0] = 1.0
        face_normals /= norm[:, None]

        n_triangles = triangles.shape[0]

        header = b"Generated by write_stl".ljust(80, b"\0")

        with open(filename, "wb") as f:
            f.write(header)
            f.write(np.uint32(n_triangles).tobytes())

            dtype = np.dtype([
                ("normal", "f4", (3,)),
                ("v1", "f4", (3,)),
                ("v2", "f4", (3,)),
                ("v3", "f4", (3,)),
                ("attr", "u2"),
            ])

            data = np.empty(n_triangles, dtype=dtype)
            data["normal"] = face_normals.astype(np.float32)
            data["v1"] = triangles[:, 0].astype(np.float32)
            data["v2"] = triangles[:, 1].astype(np.float32)
            data["v3"] = triangles[:, 2].astype(np.float32)
            data["attr"] = 0

            f.write(data.tobytes())
