writer.cpp

//
// Created by egor on 1/23/24.
//

#include "writer.h"

using namespace std;

static size_t write_counter = 1;
static size_t spring_connector_counter = 1;
static size_t neck_counter = 1;

void write_spring_connectors(std::string const & dir, double r_part, vector<array<Eigen::Vector3d, 5>> const & spring_connectors, vector<bool> const & enabled_contacts) {
    stringstream path;
    path << dir << "/spring_connectors_" << spring_connector_counter << ".vtk";
    spring_connector_counter ++;

    ofstream ofs(path.str());

    if (!ofs.good()) {
        cerr << "ERROR: Unable to create dump file" << "\n";
        exit(EXIT_FAILURE);
    }

    size_t point_count = 0;
    for (auto val : enabled_contacts) {
        if (!val) point_count += 5;
    }

    ofs << "# vtk DataFile Version 4.0" << "\n";
    ofs << "Generated by libFractalCommon" << "\n";
    ofs << "ASCII" << "\n";
    ofs << "DATASET POLYDATA" << "\n";
    ofs << "POINTS " << point_count << " FLOAT" << "\n";

    // Coordinates are normalized to avoid issues with ParaView and small particles
    for (size_t i = 0; i < enabled_contacts.size(); i ++) {
        if (!enabled_contacts[i]) {
            // This is a necked contact
            for (size_t j = 0; j < 5; j ++) {
                Eigen::Vector3d p = spring_connectors[i][j];
                ofs << p[0] / r_part << " " << p[1] / r_part << " " << p[2] / r_part << " ";
            }
        }
    }

    // Flush the buffer
    ofs << std::endl;
}

void write_particles(const std::string & dir, std::vector<Eigen::Vector3d> const & x, std::vector<Eigen::Vector3d> const & theta, double r_part) {
    stringstream path;
    path << dir << "/particles_" << write_counter << ".vtk";
    write_counter ++;

    ofstream ofs(path.str());

    if (!ofs.good()) {
        cerr << "ERROR: Unable to create dump file" << "\n";
        exit(EXIT_FAILURE);
    }

    ofs << "# vtk DataFile Version 4.0" << "\n";
    ofs << "Generated by libFractalCommon" << "\n";
    ofs << "ASCII" << "\n";
    ofs << "DATASET POLYDATA" << "\n";
    ofs << "POINTS " << x.size() << " FLOAT" << "\n";

    // Coordinates are normalized to avoid issues with ParaView and small particles
    for (auto const & p : x) {
        ofs << p[0] / r_part << " " << p[1] / r_part << " " << p[2] / r_part << " ";
    }

    ofs << "\n" << "\n";
    ofs << "POINT_DATA " << x.size() << "\n";
    ofs << "FIELD FieldData 3" << "\n";
    ofs << "thetax 3 " << x.size() << " double" << "\n";
    for (auto const & t : theta) {
    Eigen::Quaternion<double> q;
        q = Eigen::AngleAxis(t[0], Eigen::Vector3d::UnitX())
            * Eigen::AngleAxis(t[1], Eigen::Vector3d::UnitY())
            * Eigen::AngleAxis(t[2], Eigen::Vector3d::UnitZ());

        auto unit = Eigen::Vector3d::UnitX();
        Eigen::Matrix3d m = q.matrix();
        auto orient = m*unit;

        ofs << orient[0] << " " << orient[1] << " " << orient[2] << " ";
    }
    ofs << "\n" << "thetay 3 " << x.size() << " double" << "\n";
    for (auto const & t : theta) {
        Eigen::Quaternion<double> q;
        q = Eigen::AngleAxis(t[0], Eigen::Vector3d::UnitX())
            * Eigen::AngleAxis(t[1], Eigen::Vector3d::UnitY())
            * Eigen::AngleAxis(t[2], Eigen::Vector3d::UnitZ());

        auto unit = Eigen::Vector3d::UnitY();
        Eigen::Matrix3d m = q.matrix();
        auto orient = m*unit;

        ofs << orient[0] << " " << orient[1] << " " << orient[2] << " ";
    }
    ofs << "\n" << "thetaz 3 " << x.size() << " double" << "\n";
    for (auto const & t : theta) {
        Eigen::Quaternion<double> q;
        q = Eigen::AngleAxis(t[0], Eigen::Vector3d::UnitX())
            * Eigen::AngleAxis(t[1], Eigen::Vector3d::UnitY())
            * Eigen::AngleAxis(t[2], Eigen::Vector3d::UnitZ());

        auto unit = Eigen::Vector3d::UnitZ();
        Eigen::Matrix3d m = q.matrix();
        auto orient = m*unit;

        ofs << orient[0] << " " << orient[1] << " " << orient[2] << " ";
    }
    // Flush the buffer
    ofs << std::endl;
}

void write_necks(std::string const & dir, std::vector<Eigen::Vector3d> const & x, double r_part,
                 std::vector<bool> const & enabled_contacts) {

    stringstream path;
    path << dir << "/necks_" << neck_counter << ".vtk";
    neck_counter ++;

    ofstream ofs(path.str());

    if (!ofs.good()) {
        cerr << "ERROR: Unable to create dump file" << endl;
        exit(EXIT_FAILURE);
    }

    std::vector<std::pair<size_t, size_t>> neck_indices;
    size_t n_part = x.size();
    for (size_t i = 0; i < n_part - 1; i ++) {
        for (size_t j = i + 1; j < n_part; j ++) {
            if (!enabled_contacts[i * n_part + j]) neck_indices.emplace_back(i, j);
        }
    }

    ofs << "# vtk DataFile Version 4.0" << endl;
    ofs << "Generated by libFractalCommon" << endl;
    ofs << "ASCII" << endl;
    ofs << "DATASET POLYDATA" << endl;
    ofs << "POINTS " << neck_indices.size() << " FLOAT" << endl;

    // Coordinates are normalized to avoid issues with ParaView and small particles
    for (auto [i, j] : neck_indices) {
        auto const & particleA = x[i];
        auto const & particleB = x[j];

        Eigen::Vector3d pos = (particleA + particleB) / 2.0;

        ofs << pos[0] /r_part << " " << pos[1] / r_part << " " << pos[2] / r_part << " ";
    }

    ofs << endl << endl;
    ofs << "POINT_DATA " << neck_indices.size() << endl;
    ofs << "FIELD FieldData 1" << endl;
    ofs << "normals 3 " << neck_indices.size() << " double" << endl;
    for (auto [i, j] : neck_indices) {
        auto const & particleA = x[i];
        auto const & particleB = x[j];

        Eigen::Vector3d normal = (particleA - particleB).normalized();

        ofs << normal[0] << " " << normal[1] << " " << normal[2] << " ";
    }
}