precomputed_io

class PrecomputedMeshReader(PrecomputedReader):
    def __init__(
        self,
        fmt: str = DEFAULT_FMT,
        attrs: Optional[Dict[str, Any]] = None,
        errors: str = "raise",
    ):
        super().__init__(
            fmt=fmt,
            attrs=attrs,
            file_ext="",
            name_fallback="mesh",
            read_binary=True,
            errors=errors,
        )

    @base.handle_errors
    def read_buffer(
        self, f: IO, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.MeshNeuron":
        """Read buffer into a MeshNeuron.

        Parameters
        ----------
        f :         IO
                    Readable buffer - must be bytes.
        attrs :     dict | None
                    Arbitrary attributes to include in the MeshNeuron.

        Returns
        -------
        core.MeshNeuron
        """
        if not isinstance(f.read(0), bytes):
            raise ValueError(f"Expected bytes, got {type(f.read(0))}")

        num_vertices = np.frombuffer(f.read(4), np.uint32)[0]
        vertices = np.frombuffer(f.read(int(3 * 4 * num_vertices)), np.float32).reshape(
            -1, 3
        )
        faces = np.frombuffer(f.read(), np.uint32).reshape(-1, 3)

        return core.MeshNeuron(
            {"vertices": vertices, "faces": faces},
            **(
                self._make_attributes(
                    {"name": self.name_fallback, "origin": "DataFrame"}, attrs
                )
            ),
        )

class PrecomputedReader(base.BaseReader):
    def is_valid_file(self, file):
        """Return True if file should be considered for reading."""
        if isinstance(file, zipfile.ZipInfo):
            file = str(file.filename)
        elif isinstance(file, Path):
            if not file.is_file():
                return False
            file = str(file.name)
        else:
            file = str(file)

        # Drop anything with a file extension or hidden files (e.g. ".DS_store")
        if "." in file:
            return False
        # Ignore the info file
        if file == "info":
            return False
        # Ignore manifests
        if file.endswith(":0"):
            return False
        return True

class PrecomputedSkeletonReader(PrecomputedReader):
    def __init__(
        self,
        fmt: str = DEFAULT_FMT,
        attrs: Optional[Dict[str, Any]] = None,
        info: Dict[str, Any] = {},
        errors: str = "raise",
    ):
        super().__init__(
            fmt=fmt,
            attrs=attrs,
            file_ext="",
            name_fallback="skeleton",
            read_binary=True,
            errors=errors,
        )
        self.info = info

    @base.handle_errors
    def read_buffer(
        self, f: IO, attrs: Optional[Dict[str, Any]] = None
    ) -> "core.TreeNeuron":
        """Read buffer into a TreeNeuron.

        Parameters
        ----------
        f :         IO
                    Readable buffer - must be bytes.
        attrs :     dict | None
                    Arbitrary attributes to include in the TreeNeuron.

        Returns
        -------
        core.TreeNeuron

        """
        if not isinstance(f.read(0), bytes):
            raise ValueError(f"Expected bytes, got {type(f.read(0))}")

        num_nodes = np.frombuffer(f.read(4), np.uint32)[0]
        num_edges = np.frombuffer(f.read(4), np.uint32)[0]
        nodes = np.frombuffer(f.read(int(3 * 4 * num_nodes)), np.float32).reshape(-1, 3)
        edges = np.frombuffer(f.read(int(2 * 4 * num_edges)), np.uint32).reshape(-1, 2)

        swc = self.make_swc(nodes, edges)

        # Check for malformed vertex attributes (should be list of dicts)
        if isinstance(self.info.get("vertex_attributes", None), dict):
            self.info["vertex_attributes"] = [self.info["vertex_attributes"]]

        # Parse additional vertex attributes if specified as per the info file
        for attr in self.info.get("vertex_attributes", []):
            dtype = np.dtype(attr["data_type"])
            n_comp = attr["num_components"]
            values = np.frombuffer(
                f.read(int(n_comp * dtype.itemsize * num_nodes)), dtype
            ).reshape(-1, n_comp)
            if n_comp == 1:
                swc[attr["id"]] = values.flatten()
            else:
                for i in range(n_comp):
                    swc[f"{attr['id']}_{i}"] = values[:, i]

        return core.TreeNeuron(
            swc,
            **(
                self._make_attributes(
                    {"name": self.name_fallback, "origin": "DataFrame"}, attrs
                )
            ),
        )

    def make_swc(self, nodes: np.ndarray, edges: np.ndarray) -> pd.DataFrame:
        """Make SWC table from nodes and edges.

        Parameters
        ----------
        nodes :     (N, 3) array
        edges :     (N, 2) array

        Returns
        -------
        pandas.DataFrame
        """
        swc = pd.DataFrame()
        swc["node_id"] = np.arange(len(nodes))
        swc["x"], swc["y"], swc["z"] = nodes[:, 0], nodes[:, 1], nodes[:, 2]

        edge_dict = dict(zip(edges[:, 1], edges[:, 0]))
        swc["parent_id"] = swc.node_id.map(lambda x: edge_dict.get(x, -1)).astype(
            np.int32
        )

        return swc

class PrecomputedWriter(base.Writer):
    """Writer class that also takes care of `info` files."""

    def write_any(self, x, filepath, write_info=True, **kwargs):
        """Write any to file. Default entry point."""
        # First write the actual neurons
        kwargs["write_info"] = False
        super().write_any(x, filepath=filepath, **kwargs)

        # Write info file to the correct directory/zipfile
        if write_info:
            add_props = {}
            if kwargs.get("radius", False):
                add_props["vertex_attributes"] = [
                    {"id": "radius", "data_type": "float32", "num_components": 1}
                ]

            if str(self.path).endswith(".zip"):
                with ZipFile(self.path, mode="a") as zf:
                    # Context-manager will remove temporary directory and its contents
                    with tempfile.TemporaryDirectory() as tempdir:
                        # Write info to zip
                        if write_info:
                            # Generate temporary filename
                            f = os.path.join(tempdir, "info")
                            write_info_file(x, f, add_props=add_props)
                            # Add file to zip
                            zf.write(f, arcname="info", compress_type=compression)
            else:
                fp = self.path
                # Find the first existing root directory
                while not fp.is_dir():
                    fp = fp.parent

                write_info_file(x, fp, add_props=add_props)