import numpy as np

from .. import util
from ..points import PointCloud


def load_xyz(file_obj, delimiter=None, **kwargs):
    """
    Load an XYZ file into a PointCloud.

    Parameters
    ------------
    file_obj : an open file-like object
      Source data, ASCII XYZ
    delimiter : None or string
      Characters used to separate the columns of the file
      If not passed will use whitespace or commas

    Returns
    ----------
    kwargs : dict
      Data which can be passed to PointCloud constructor
    """
    # read the whole file into memory as a string
    raw = util.decode_text(file_obj.read()).strip()
    # get the first line to look at
    first = raw[: raw.find("\n")].strip()
    # guess the column count by looking at the first line
    columns = len(first.split())
    if columns < 3:
        raise ValueError("not enough columns in xyz file!")

    if delimiter is None and "," in first:
        # if no delimiter passed and file has commas
        delimiter = ","
    if delimiter is not None:
        # replace delimiter with whitespace so split works
        raw = raw.replace(delimiter, " ")

    # use string splitting to get array
    array = np.array(raw.split(), dtype=np.float64)
    # reshape to column count
    # if file has different numbers of values
    # per row this will fail as it should
    data = array.reshape((-1, columns))

    # start with no colors
    colors = None
    # vertices are the first three columns
    vertices = data[:, :3]
    if columns == 6:
        # RGB colors
        colors = np.array(data[:, 3:], dtype=np.uint8)
        colors = np.concatenate(
            (colors, np.ones((len(data), 1), dtype=np.uint8) * 255), axis=1
        )
    elif columns >= 7:
        # extract RGBA colors
        colors = np.array(data[:, 3:8], dtype=np.uint8)
    # add extracted colors and vertices to kwargs
    kwargs.update({"vertices": vertices, "colors": colors})

    return kwargs


def export_xyz(cloud, write_colors=True, delimiter=None):
    """
    Export a PointCloud object to an XYZ format string.

    Parameters
    -------------
    cloud : trimesh.PointCloud
      Geometry in space
    write_colors : bool
      Write colors or not
    delimiter : None or str
      What to separate columns with

    Returns
    --------------
    export : str
      Pointcloud in XYZ format
    """
    if not isinstance(cloud, PointCloud):
        raise ValueError("object must be PointCloud")

    # compile data into a blob
    data = cloud.vertices
    if write_colors and hasattr(cloud, "colors") and cloud.colors is not None:
        # stack colors and  vertices
        data = np.hstack((data, cloud.colors))

    # if delimiter not passed use whitespace
    if delimiter is None:
        delimiter = " "
    # stack blob into XYZ format
    export = util.array_to_string(data, col_delim=delimiter)

    return export


_xyz_loaders = {"xyz": load_xyz}
_xyz_exporters = {"xyz": export_xyz}