PlyIO.h

#pragma once

#include <iterator>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>

namespace Euclid
{
// Forward declaration
class PlyReader;
class PlyWriter;

enum class PlyFormat
{
    ascii,                
    binary_little_endian, 
    binary_big_endian     
};

class PlyProperty
{
public:
    explicit PlyProperty(const std::string& name) : _name(name) {}

    PlyProperty(const std::string& name, bool is_list)
        : _name(name), _is_list(is_list)
    {}

    virtual ~PlyProperty() = default;

    PlyProperty(const PlyProperty&) = default;

    PlyProperty(PlyProperty&&) = default;

    PlyProperty& operator=(const PlyProperty&) = default;

    PlyProperty& operator=(PlyProperty&&) = default;

    const std::string& name() const
    {
        return _name;
    }

    bool is_list() const
    {
        return _is_list;
    }

    virtual std::string type_str() const = 0;

    virtual void apply(PlyReader& reader,
                       std::ifstream& stream,
                       PlyFormat format) const = 0;

    virtual void apply(PlyWriter& writer,
                       std::ofstream& stream,
                       PlyFormat format) const = 0;

private:
    std::string _name;
    bool _is_list = false;
};

class PlyDoubleProperty : public PlyProperty
{
public:
    using value_type = double;

public:
    explicit PlyDoubleProperty(const std::string& name) : PlyProperty(name) {}

    PlyDoubleProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "double";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyFloatProperty : public PlyProperty
{
public:
    using value_type = float;

public:
    explicit PlyFloatProperty(const std::string& name) : PlyProperty(name) {}

    PlyFloatProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "float";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyIntProperty : public PlyProperty
{
public:
    using value_type = int;

public:
    explicit PlyIntProperty(const std::string& name) : PlyProperty(name) {}

    PlyIntProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "int";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyUintProperty : public PlyProperty
{
public:
    using value_type = unsigned;

public:
    explicit PlyUintProperty(const std::string& name) : PlyProperty(name) {}

    PlyUintProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "uint";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyShortProperty : public PlyProperty
{
public:
    using value_type = int;

public:
    explicit PlyShortProperty(const std::string& name) : PlyProperty(name) {}

    PlyShortProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "short";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyUshortProperty : public PlyProperty
{
public:
    using value_type = unsigned;

public:
    explicit PlyUshortProperty(const std::string& name) : PlyProperty(name) {}

    PlyUshortProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "ushort";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyCharProperty : public PlyProperty
{
public:
    using value_type = int;

public:
    explicit PlyCharProperty(const std::string& name) : PlyProperty(name) {}

    PlyCharProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "char";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyUcharProperty : public PlyProperty
{
public:
    using value_type = unsigned;

public:
    explicit PlyUcharProperty(const std::string& name) : PlyProperty(name) {}

    PlyUcharProperty(const std::string& name, bool is_list)
        : PlyProperty(name, is_list)
    {}

    std::string type_str() const override
    {
        return "uchar";
    }

    void apply(PlyReader& reader,
               std::ifstream& stream,
               PlyFormat format) const override;

    void apply(PlyWriter& writer,
               std::ofstream& stream,
               PlyFormat format) const override;

    value_type get_ascii(std::ifstream& stream) const;

    value_type get_binary(std::ifstream& stream,
                          bool file_little_endian,
                          bool sys_little_endian) const;

    void put_ascii(std::ofstream& stream, value_type value) const;

    void put_binary(std::ofstream& stream,
                    value_type value,
                    bool file_little_endian,
                    bool sys_little_endian) const;
};

class PlyElement
{
public:
    class iterator
    {
    public:
        using iterator_category = std::forward_iterator_tag;
        using value_type = PlyProperty;
        using reference = value_type&;
        using pointer = value_type*;
        using difference_type = int;
        using unique_iter = std::vector<std::unique_ptr<PlyProperty>>::iterator;

        explicit iterator(unique_iter iter) : _iter(iter) {}

        iterator& operator++()
        {
            ++_iter;
            return *this;
        }

        iterator operator++(int)
        {
            auto cpy = *this;
            ++(*this);
            return cpy;
        }

        reference operator*()
        {
            return **_iter;
        }

        bool operator!=(const iterator& rhs)
        {
            return _iter != rhs._iter;
        }

    private:
        unique_iter _iter;
    };

    class const_iterator
    {
    public:
        using iterator_category = std::forward_iterator_tag;
        using value_type = const PlyProperty;
        using reference = value_type&;
        using pointer = value_type*;
        using difference_type = int;
        using unique_iter =
            std::vector<std::unique_ptr<PlyProperty>>::const_iterator;

        explicit const_iterator(unique_iter iter) : _iter(iter) {}

        const_iterator& operator++()
        {
            ++_iter;
            return *this;
        }

        const_iterator operator++(int)
        {
            auto cpy = *this;
            ++(*this);
            return cpy;
        }

        reference operator*()
        {
            return **_iter;
        }

        bool operator!=(const const_iterator& rhs)
        {
            return _iter != rhs._iter;
        }

    private:
        unique_iter _iter;
    };

    PlyElement(const std::string& name, unsigned count)
        : _name(name), _count(count)
    {}

    ~PlyElement() = default;

    PlyElement(const PlyElement& rhs) = delete;

    PlyElement& operator=(const PlyElement& rhs) = delete;

    PlyElement(PlyElement&& rhs)
    {
        _name = rhs._name;
        _count = rhs._count;
        std::copy(std::make_move_iterator(rhs._properties.begin()),
                  std::make_move_iterator(rhs._properties.end()),
                  std::back_inserter(_properties));
    }

    PlyElement& operator=(PlyElement&& rhs)
    {
        _name = rhs._name;
        _count = rhs._count;
        std::copy(std::make_move_iterator(rhs._properties.begin()),
                  std::make_move_iterator(rhs._properties.end()),
                  std::back_inserter(_properties));
        return *this;
    }

    const std::string& name() const
    {
        return _name;
    }

    unsigned count() const
    {
        return _count;
    }

    void add_property(std::unique_ptr<PlyProperty> property)
    {
        _properties.push_back(std::move(property));
    }

    PlyProperty* property(size_t i)
    {
        return _properties[i].get();
    }

    const PlyProperty* property(size_t i) const
    {
        return _properties[i].get();
    }

    size_t n_props() const
    {
        return _properties.size();
    }

    iterator begin()
    {
        return iterator(_properties.begin());
    }

    const_iterator begin() const
    {
        return const_iterator(_properties.cbegin());
    }

    iterator end()
    {
        return iterator(_properties.end());
    }

    const_iterator end() const
    {
        return const_iterator(_properties.cend());
    }

private:
    std::string _name;
    unsigned _count;
    std::vector<std::unique_ptr<PlyProperty>> _properties;
};

class PlyHeader
{
public:
    explicit PlyHeader(PlyFormat format) : _format(format) {}

    PlyHeader(const PlyHeader&) = delete;

    PlyHeader(PlyHeader&&) = default;

    PlyHeader& operator=(const PlyHeader&) = delete;

    PlyHeader& operator=(PlyHeader&&) = default;

    PlyFormat format() const
    {
        return _format;
    }

    void set_format(PlyFormat format)
    {
        _format = format;
    }

    PlyElement& element(size_t i)
    {
        return _elements[i];
    }

    const PlyElement& element(size_t i) const
    {
        return _elements[i];
    }

    void add_element(PlyElement&& element)
    {
        _elements.push_back(std::move(element));
    }

    size_t n_elems() const
    {
        return _elements.size();
    }

    decltype(auto) begin()
    {
        return _elements.cbegin();
    }

    decltype(auto) begin() const
    {
        return _elements.begin();
    }

    decltype(auto) end()
    {
        return _elements.end();
    }

    decltype(auto) end() const
    {
        return _elements.cend();
    }

private:
    PlyFormat _format;
    std::vector<PlyElement> _elements;
};

class PlyReader
{
public:
    PlyReader();

    virtual ~PlyReader() = default;

    virtual void on_read(const PlyHeader&) {}

    virtual void read(const PlyDoubleProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyFloatProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyIntProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyUintProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyShortProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyUshortProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyCharProperty*, std::ifstream&, PlyFormat) {}

    virtual void read(const PlyUcharProperty*, std::ifstream&, PlyFormat) {}

protected:
    bool _sys_little_endian;
};

class PlyWriter
{
public:
    PlyWriter();

    virtual ~PlyWriter() = default;

    virtual void on_write() {}

    virtual PlyHeader generate_header(PlyFormat format) const
    {
        PlyHeader header(format);
        return header;
    };

    virtual void write(const PlyDoubleProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyFloatProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyIntProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyUintProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyShortProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyUshortProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyCharProperty*, std::ofstream&, PlyFormat) {}

    virtual void write(const PlyUcharProperty*, std::ofstream&, PlyFormat) {}

protected:
    bool _sys_little_endian;
};

template<int VN, typename FloatType, typename IndexType, typename ColorType>
class CommonPlyReader : public PlyReader
{
public:
    CommonPlyReader(std::vector<FloatType>& positions,
                    std::vector<FloatType>* normals = nullptr,
                    std::vector<FloatType>* texcoords = nullptr,
                    std::vector<IndexType>* indices = nullptr,
                    std::vector<ColorType>* colors = nullptr)
        : _positions(positions), _normals(normals), _texcoords(texcoords),
          _indices(indices), _colors(colors)
    {}

    void on_read(const PlyHeader& header) override;

    void read(const PlyDoubleProperty* property,
              std::ifstream& stream,
              PlyFormat format) override;

    void read(const PlyFloatProperty* property,
              std::ifstream& stream,
              PlyFormat format) override;

    void read(const PlyIntProperty* property,
              std::ifstream& stream,
              PlyFormat format) override;

    void read(const PlyUintProperty* property,
              std::ifstream& stream,
              PlyFormat format) override;

    void read(const PlyUcharProperty* property,
              std::ifstream& stream,
              PlyFormat format) override;

private:
    template<typename TPlyProperty>
    void _store_float(const TPlyProperty* property,
                      std::ifstream& stream,
                      PlyFormat format);

    template<typename TPlyProperty>
    void _store_color(const TPlyProperty* property,
                      std::ifstream& stream,
                      PlyFormat format);

    template<typename TPlyProperty>
    void _store_indices(const TPlyProperty* property,
                        std::ifstream& stream,
                        PlyFormat format);

private:
    std::vector<FloatType>& _positions;
    std::vector<FloatType>* _normals = nullptr;
    std::vector<FloatType>* _texcoords = nullptr;
    std::vector<IndexType>* _indices = nullptr;
    std::vector<ColorType>* _colors = nullptr;
};

template<int VN, typename FloatType, typename IndexType, typename ColorType>
class CommonPlyWriter : public PlyWriter
{
public:
    CommonPlyWriter(const std::vector<FloatType>& positions,
                    const std::vector<FloatType>* normals = nullptr,
                    const std::vector<FloatType>* texcoords = nullptr,
                    const std::vector<IndexType>* indices = nullptr,
                    const std::vector<ColorType>* colors = nullptr)
        : _positions(positions), _normals(normals), _texcoords(texcoords),
          _indices(indices), _colors(colors)
    {
        auto n = positions.size() / 3;
        if (_normals != nullptr && _normals->size() != n * 3) {
            throw std::invalid_argument(
                "Size of normals not compatible with positions");
        }
        if (_texcoords != nullptr && _texcoords->size() != n * 2) {
            throw std::invalid_argument(
                "Size of texcoords not compatible with positions");
        }
        if (_colors != nullptr) {
            if (_colors->size() == n * 3) {
                _has_alpha = false;
            }
            else if (_colors->size() == n * 4) {
                _has_alpha = true;
            }
            else {
                throw std::invalid_argument(
                    "Size of colors not compatible with positions");
            }
        }
        if (_indices != nullptr && _indices->size() % VN != 0) {
            throw std::invalid_argument(
                "Size of indices not compatible with VN");
        }
    }

    void on_write() override;

    PlyHeader generate_header(PlyFormat format) const override;

    void write(const PlyDoubleProperty* property,
               std::ofstream& stream,
               PlyFormat format) override;

    void write(const PlyFloatProperty* property,
               std::ofstream& stream,
               PlyFormat format) override;

    void write(const PlyIntProperty* property,
               std::ofstream& stream,
               PlyFormat format) override;

    void write(const PlyUintProperty* property,
               std::ofstream& stream,
               PlyFormat format) override;

    void write(const PlyUcharProperty* property,
               std::ofstream& stream,
               PlyFormat format) override;

private:
    template<typename TPlyProperty>
    void _write_float(const TPlyProperty* property,
                      std::ofstream& stream,
                      PlyFormat format);

    template<typename TPlyProperty>
    void _write_color(const TPlyProperty* property,
                      std::ofstream& stream,
                      PlyFormat format);

    template<typename TPlyProperty>
    void _write_indices(const TPlyProperty* property,
                        std::ofstream& stream,
                        PlyFormat format);

private:
    const std::vector<FloatType>& _positions;
    const std::vector<FloatType>* _normals = nullptr;
    const std::vector<FloatType>* _texcoords = nullptr;
    const std::vector<IndexType>* _indices = nullptr;
    const std::vector<ColorType>* _colors = nullptr;
    bool _has_alpha = false;
    size_t _piter = 0;
    size_t _iiter = 0;
    size_t _niter = 0;
    size_t _titer = 0;
    size_t _citer = 0;
};

PlyHeader read_ply_header(const std::string& filename);

void read_ply(const std::string& filename, PlyReader& reader);

template<int VN,
         typename FloatType,
         typename IndexType,
         typename ColorType,
         typename Interface>
void read_ply(const std::string& filename,
              std::vector<FloatType>& positions,
              std::vector<FloatType>* normals,
              std::vector<FloatType>* texcoords,
              std::vector<IndexType>* indices,
              std::vector<ColorType>* colors);

void write_ply(const std::string& filename,
               PlyWriter& writer,
               PlyFormat format = PlyFormat::ascii);

template<int VN,
         typename FloatType,
         typename IndexType,
         typename ColorType,
         typename Interface>
void write_ply(const std::string& filename,
               const std::vector<FloatType>& positions,
               const std::vector<FloatType>* normals,
               const std::vector<FloatType>* texcoords,
               const std::vector<IndexType>* indices,
               const std::vector<ColorType>* colors,
               PlyFormat format = PlyFormat::ascii);

} // namespace Euclid

#include "src/PlyIO.cpp"