WheelSimulator.cpp

// WheelSimulator.cpp
#include "WheelSimulator.h"
#include "Utils.h"
#include "Constants.h"
#include "include/json.hpp"
using json = nlohmann::json;

#include <chrono>
#include <cmath>
#include <cstdio>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <random>
#include <sstream>
#include <stdexcept>
#include <unordered_map>

using namespace deme;

// #define DEBUG
#ifdef DEBUG
#include <ostream>

// Helper to print maps
static std::ostream& print_material_properties(std::ostream& os,
                                               const std::unordered_map<std::string, float>& m) {
    os << "{";
    bool first = true;
    for (const auto& [k, v] : m) {
        if (!first) os << ", ";
        first = false;
        os << k << ": " << v;
    }
    os << "}";
    return os;
}

// If float3 has x, y, z members; adjust if your type is different
static std::ostream& print_float3(std::ostream& os, const float3& v) {
    return os << "(" << v.x << ", " << v.y << ", " << v.z << ")";
}

std::ostream& operator<<(std::ostream& os, const Wheel& w) {
    os << "Wheel{"
       << "r_effective=" << w.r_effective
       << ", r_outer=" << w.r_outer
       << ", width=" << w.width
       << ", mass=" << w.mass
       << ", total_mass=" << w.total_mass
       << ", IXX=" << w.IXX
       << ", IYY=" << w.IYY
       << ", IZZ=" << w.IZZ
       << ", mesh_file_path=" << w.mesh_file_path
       << ", material_properties=";
    print_material_properties(os, w.material_properties);
    os << "}";
    return os;
}

std::ostream& operator<<(std::ostream& os, const Terrain& t) {
    os << "Terrain{"
       << "terrain_filepath=" << t.terrain_filepath
       << ", world_size=(" << t.world_size_x << ", "
                           << t.world_size_y << ", "
                           << t.world_size_z << ")"
       << ", world_bottom=" << t.world_bottom
       << ", terrain_density=" << t.terrain_density
       << ", volume1=" << t.volume1
       << ", volume2=" << t.volume2
       << ", scales=[";
    for (std::size_t i = 0; i < t.scales.size(); ++i) {
        if (i > 0) os << ", ";
        os << t.scales[i];
    }
    os << "]"
       << ", MOI1=";
    print_float3(os, t.MOI1);
    os << ", MOI2=";
    print_float3(os, t.MOI2);
    os << ", material_properties=";
    print_material_properties(os, t.material_properties);
    os << "}";
    return os;
}

std::ostream& operator<<(std::ostream& os, const SimParams& s) {
    os << "SimParams{"
       << "slip=" << s.slip
       << ", sim_endtime=" << s.sim_endtime
       << ", batch_dir=\"" << s.batch_dir << "\""
       << ", data_drivepath=" << s.data_drivepath
       << ", rotational_velocity=" << s.rotational_velocity
       << ", step_size=" << s.step_size
       << ", angle_deg=" << s.angle_deg
       << ", offset=(" << s.offset_x << ", "
                       << s.offset_y << ", "
                       << s.offset_z << ")"
       << ", settling_time=" << s.settling_time
       << "}";
    return os;
}
#endif

WheelSimulator::WheelSimulator(Wheel wheel, Terrain terrain, 
                                SimParams simparams, const json param)
    : param_(param),
      fps_(Constants::FPS),
      out_steps_(static_cast<unsigned int>(1.0 / (Constants::FPS * Constants::INITIAL_STEP_SIZE))),
      report_steps_(static_cast<unsigned int>(1.0 / (Constants::REPORT_PERTIMESTEP * Constants::INITIAL_STEP_SIZE))),
      curr_step_(0),
      currframe_(0),
      frame_time_(1.0 / Constants::FPS),
      total_pressure_(0.0f),
      added_pressure_(0.0f),
      mat_type_terrain_(DEMSim_.LoadMaterial(terrain.material_properties)),
      wheel_(wheel), // initializes wheel
      terrain_(terrain),
      simparams_(simparams)
      
      
{
    // Constructor body. Can remain empty or initialize additional members if necessary
#ifdef DEBUG    
    std::cout << wheel_ << '\n';
    std::cout << terrain_ << '\n';
    std::cout << simparams_ << '\n';
#endif
}

void WheelSimulator::PrepareSimulation() {
    std::cout << "terrain" << terrain_.terrain_filepath <<std::endl;
    std::cout << "data dir" << simparams_.data_drivepath <<std::endl;
    std::cout << "outer radius" << wheel_.r_outer <<std::endl;
    std::cout << "effective radius" << wheel_.r_effective <<std::endl;
    
    std::cout << "Intializing Output Directories" <<std::endl;
    InitializeOutputDirectories();
    std::cout << "Writing Sim Parameters" <<std::endl;
    WriteSimulationParameters();
    std::cout << "Intializing Output Files" <<std::endl;
    InitializeOutputFiles();
    std::cout << "Configuring DEM Solver" <<std::endl;
    ConfigureDEMSolver();
    std::cout << "Preparing Particles" <<std::endl;
    PrepareParticles();
    std::cout << "Configuring Wheel" <<std::endl;
    ConfigureWheel();
    std::cout << "Setting up prescribed motions" <<std::endl;
    SetupPrescribedMotions();
    std::cout << "Setting up inspectors" <<std::endl;
    SetupInspectors();
    std::cout << "Initializing DEMsim..." << std::endl;
    DEMSim_.Initialize();
}

void WheelSimulator::RunSimulation() {
    std::cout << "Letting wheel sink..." << std::endl;
    PerformInitialSink();
    std::cout << "Applying Wheel Forward Motion..." << std::endl;
    ApplyWheelForwardMotion();
    std::cout << "Running main sim loop..." << std::endl;
    RunSimulationLoop();
}

// Create the output folder structure
void WheelSimulator::InitializeOutputDirectories() {
    out_dir_ = simparams_.data_drivepath / simparams_.batch_dir / (Utils::getCurrentTimeStamp() + "_SkidSteerSim_" + std::to_string(simparams_.slip));
    rover_dir_ = out_dir_ / "rover";
    particles_dir_ = out_dir_ / "particles";

    std::error_code ec;
    if (!std::filesystem::create_directories(rover_dir_, ec) ||
        !std::filesystem::create_directories(particles_dir_, ec)) {
        throw std::runtime_error("Failed to create output directories.");
    }

    std::cout << "Output directory: " << out_dir_ << std::endl;
}

// Write params.json
void WheelSimulator::WriteSimulationParameters() {
    std::filesystem::path output_params_path = out_dir_ / "params.json";
    output_params_.open(output_params_path);
    if (!output_params_) {
        throw std::runtime_error("Failed to open params.json for writing.");
    }
    output_params_ << param_.dump(4);
    output_params_.close();
}

// Initialize output.csv
void WheelSimulator::InitializeOutputFiles() {
    std::filesystem::path output_datafile_path = out_dir_ / "output.csv";
    output_datafile_.open(output_datafile_path);
    if (!output_datafile_) {
        throw std::runtime_error("Failed to open output.csv for writing.");
    }
    output_datafile_ << "t,f_x,f_y,f_z,d_c,v_max,pos_x,pos_y,pos_z,oriq_x,oriq_y,oriq_z,oriq_w,vel_x,vel_y,vel_z" << std::endl;
}

// Configure DEM Solver settings
void WheelSimulator::ConfigureDEMSolver() {
    DEMSim_.SetVerbosity(INFO);
    DEMSim_.SetOutputFormat(OUTPUT_FORMAT::CSV);
    DEMSim_.SetOutputContent(OUTPUT_CONTENT::ABSV);
    DEMSim_.SetOutputContent(OUTPUT_CONTENT::VEL);
    DEMSim_.SetOutputContent(OUTPUT_CONTENT::ABS_ACC);
    DEMSim_.SetOutputContent(OUTPUT_CONTENT::FAMILY);
    DEMSim_.SetMeshOutputFormat(MESH_FORMAT::VTK);
    DEMSim_.SetContactOutputContent({"OWNER", "FORCE", "POINT"});

    // Family Settings
    DEMSim_.SetFamilyFixed(Family::FIXED);
    DEMSim_.DisableContactBetweenFamilies(Family::FIXED, Family::FIXED);

    // Force Recording
    DEMSim_.SetNoForceRecord();

    // World Settings
    DEMSim_.SetInitTimeStep(simparams_.step_size);
    DEMSim_.SetGravitationalAcceleration(make_float3(0.0f, 0.0f, -Constants::GRAVITY_MAGNITUDE));
    DEMSim_.SetMaxVelocity(Constants::MAX_VELOCITY);
    DEMSim_.SetErrorOutVelocity(Constants::ERROR_OUT_VELOCITY);
    DEMSim_.SetExpandSafetyMultiplier(Constants::EXPAND_SAFETY_MULTIPLIER);
    DEMSim_.SetCDUpdateFreq(Constants::CD_UPDATE_FREQ);
}

// Load in the particles
void WheelSimulator::PrepareParticles() {
    // Prepare Terrain Particles
    // Load clump types and properties
    DEMSim_.SetMaterialPropertyPair("mu", DEMSim_.LoadMaterial(wheel_.material_properties), mat_type_terrain_, 0.8);

    std::cout << "Defining World..." << std::endl;

    // Define world dimensions
    double world_size_x = terrain_.world_size_x;
    double world_size_y = terrain_.world_size_y;
    double world_size_z = terrain_.world_size_z;
    DEMSim_.InstructBoxDomainDimension(world_size_x, world_size_y, world_size_z);
    DEMSim_.InstructBoxDomainBoundingBC("top_open", mat_type_terrain_);

    // Add boundary planes
    float bottom = terrain_.world_bottom;
    DEMSim_.AddBCPlane(make_float3(0.0f, 0.0f, bottom), make_float3(0.0f, 0.0f, 1.0f), mat_type_terrain_);
    DEMSim_.AddBCPlane(make_float3(0.0f, static_cast<float>(world_size_y) / 2.0f, 0.0f), make_float3(0.0f, -1.0f, 0.0f), mat_type_terrain_);
    DEMSim_.AddBCPlane(make_float3(0.0f, -static_cast<float>(world_size_y) / 2.0f, 0.0f), make_float3(0.0f, 1.0f, 0.0f), mat_type_terrain_);
    DEMSim_.AddBCPlane(make_float3(-static_cast<float>(world_size_x) / 2.0f, 0.0f, 0.0f), make_float3(1.0f, 0.0f, 0.0f), mat_type_terrain_);
    DEMSim_.AddBCPlane(make_float3(static_cast<float>(world_size_x) / 2.0f, 0.0f, 0.0f), make_float3(-1.0f, 0.0f, 0.0f), mat_type_terrain_);

    // Define terrain particle templates
    float mass1 = terrain_.terrain_density * terrain_.volume1;
    float3 MOI1 = terrain_.MOI1 * terrain_.terrain_density;
    float mass2 = terrain_.terrain_density * terrain_.volume2;
    float3 MOI2 = terrain_.MOI2 * terrain_.terrain_density;

    std::cout << "Loading clump templates..." << std::endl;

    // Load clump templates
    std::shared_ptr<DEMClumpTemplate> my_template2 = DEMSim_.LoadClumpType(mass2, MOI2, GetDEMEDataFile("clumps/triangular_flat_6comp.csv"), mat_type_terrain_);
    std::shared_ptr<DEMClumpTemplate> my_template1 = DEMSim_.LoadClumpType(mass1, MOI1, GetDEMEDataFile("clumps/triangular_flat.csv"), mat_type_terrain_);
    std::vector<std::shared_ptr<DEMClumpTemplate>> ground_particle_templates = {
        my_template2,
        DEMSim_.Duplicate(my_template2),
        my_template1,
        DEMSim_.Duplicate(my_template1),
        DEMSim_.Duplicate(my_template1),
        DEMSim_.Duplicate(my_template1),
        DEMSim_.Duplicate(my_template1)
    };

    // Scale and name templates
    for (size_t i = 0; i < terrain_.scales.size(); ++i) {
        auto& tmpl = ground_particle_templates.at(i);
        tmpl->Scale(terrain_.scales.at(i));

        char t_name[20];
        std::sprintf(t_name, "%04zu", i);
        tmpl->AssignName(std::string(t_name));
    }

    // Load clump locations from file
    std::cout << "Making terrain..." << std::endl;
    std::unordered_map<std::string, std::vector<float3>> clump_xyz;
    std::unordered_map<std::string, std::vector<float4>> clump_quaternion;

    try {
        clump_xyz = DEMSim_.ReadClumpXyzFromCsv(terrain_.terrain_filepath);
        clump_quaternion = DEMSim_.ReadClumpQuatFromCsv(terrain_.terrain_filepath);
    } catch (...) {
        throw std::runtime_error("Failed to read clump checkpoint file. Ensure the file exists and is correctly formatted.");
    }

    std::vector<float3> in_xyz;
    std::vector<float4> in_quat;
    std::vector<std::shared_ptr<DEMClumpTemplate>> in_types;
    unsigned int t_num = 0;

    for (const auto& scale : terrain_.scales) {
        char t_name[20];
        std::sprintf(t_name, "%04u", t_num);

        auto this_type_xyz = clump_xyz[std::string(t_name)];
        auto this_type_quat = clump_quaternion[std::string(t_name)];

        size_t n_clump_this_type = this_type_xyz.size();
        std::cout << "Loading clump " << std::string(t_name) << " with " << n_clump_this_type << " particles." << std::endl;

        std::vector<std::shared_ptr<DEMClumpTemplate>> this_type(n_clump_this_type, ground_particle_templates.at(t_num));

        in_xyz.insert(in_xyz.end(), this_type_xyz.begin(), this_type_xyz.end());
        in_quat.insert(in_quat.end(), this_type_quat.begin(), this_type_quat.end());
        in_types.insert(in_types.end(), this_type.begin(), this_type.end());

        std::cout << "Added clump type " << t_num << std::endl;
        t_num++;
    }

    // Create and add clump batch
    DEMClumpBatch base_batch(in_xyz.size());
    base_batch.SetTypes(in_types);
    base_batch.SetPos(in_xyz);
    base_batch.SetOriQ(in_quat);

    DEMSim_.AddClumps(base_batch);
}

// Load in the wheel
void WheelSimulator::ConfigureWheel() {
    // Define simulation parameters
    total_pressure_ = wheel_.total_mass * Constants::GRAVITY_MAGNITUDE; // N
    added_pressure_ = (wheel_.total_mass - wheel_.mass) * Constants::GRAVITY_MAGNITUDE; // N

    std::cout << "Total Pressure: " << total_pressure_ << "N" << std::endl;
    std::cout << "Added Pressure: " << added_pressure_ << "N" << std::endl;

    // Add wheel object
    auto wheel = DEMSim_.AddWavefrontMeshObject(wheel_.mesh_file_path, DEMSim_.LoadMaterial(wheel_.material_properties));
    wheel->SetMass(wheel_.mass);
    wheel->SetMOI(make_float3(wheel_.IXX, wheel_.IYY, wheel_.IZZ));
    wheel->SetFamily(Family::ROTATING);
    wheel_tracker_ = DEMSim_.Track(wheel);
}

void WheelSimulator::SetupPrescribedMotions() {
    // Families' prescribed motions
    float v_ref = simparams_.rotational_velocity * wheel_.r_effective;

    //TODO: Turn family numbers into enums with descriptive names

    DEMSim_.SetFamilyPrescribedAngVel(Family::ROTATING, "0", Utils::toStringWithPrecision(simparams_.rotational_velocity), "0", false);
    DEMSim_.AddFamilyPrescribedAcc(Family::ROTATING, "none", "none", Utils::toStringWithPrecision(-added_pressure_ / wheel_.mass)); // TODO: What does this number mean?

    DEMSim_.SetFamilyPrescribedAngVel(Family::ROTATING_AND_TRANSLATING, "0", Utils::toStringWithPrecision(simparams_.rotational_velocity), "0", false);
    DEMSim_.SetFamilyPrescribedLinVel(Family::ROTATING_AND_TRANSLATING, Utils::toStringWithPrecision(v_ref * (1.0 - simparams_.slip)), "0", "none", false);
    DEMSim_.AddFamilyPrescribedAcc(Family::ROTATING_AND_TRANSLATING, "none", "none", Utils::toStringWithPrecision(-added_pressure_ / wheel_.mass)); // TODO: What does this number mean?
}

void WheelSimulator::SetupInspectors() {
    // Setup inspectors. These let us track and query data of objects.
    max_z_finder_ = DEMSim_.CreateInspector("clump_max_z");
    min_z_finder_ = DEMSim_.CreateInspector("clump_min_z");
    total_mass_finder_ = DEMSim_.CreateInspector("clump_mass");
    max_v_finder_ = DEMSim_.CreateInspector("clump_max_absv");
}

void WheelSimulator::WriteParticleCSV() {
    char filename[200];
    sprintf(filename, "%s/DEMdemo_output_%04d.csv", particles_dir_.c_str(), currframe_);
    DEMSim_.WriteSphereFile(std::string(filename));
}

void WheelSimulator::WriteWheelMesh() {
    char meshname[200];
    sprintf(meshname, "%s/DEMdemo_mesh_%04d.vtk", rover_dir_.c_str(), currframe_);
    DEMSim_.WriteMeshFile(std::string(meshname));
}

void WheelSimulator::PerformInitialSink() {
    // Put the wheel in place, then let the wheel sink in initially
    std::cout << "Setting max z value" << std::endl;
    float max_z = max_z_finder_->GetValue();
    std::cout << "Setting wheel position" << std::endl;
    if (wheel_tracker_) {
        wheel_tracker_->SetPos(make_float3(simparams_.offset_x, simparams_.offset_y, max_z + simparams_.offset_z + wheel_.r_outer));
        //offset wheel orientation a bit for steering test
        const float rad = simparams_.angle_deg * (float)M_PI / 180.0f;
        const float4 initQ = make_float4(0.0f, 0.0f, sinf(0.5f * rad), cosf(0.5f * rad));
        wheel_tracker_->SetOriQ(initQ);
    } else {
        std::cerr << "Error: wheel_tracker_ is null!" << std::endl;
    }
    std::cout << "Starting wheel settling loop" << std::endl;
    for (double t = 0; t < simparams_.settling_time; t += frame_time_) {
        std::cout << "Outputting frame: " << currframe_ << std::endl;
        std::cout << "Writing sphere" << std::endl;
        WriteParticleCSV();
        std::cout << "Writing mesh" << std::endl;
        WriteWheelMesh();
        std::cout << "Doing dynamics" << std::endl;
        DEMSim_.DoDynamicsThenSync(frame_time_);
        currframe_++;
    }
}

void WheelSimulator::ApplyWheelForwardMotion() {
    // Switch wheel from free fall into DP test. Tell it to start driving forward.
    DEMSim_.DoDynamicsThenSync(0);
    DEMSim_.ChangeFamily(Family::ROTATING, Family::ROTATING_AND_TRANSLATING);
}

void WheelSimulator::UpdateActiveBoxDomain(float box_halfsize_x, float box_halfsize_y) {
    // The active box domain is a trick to improve performance. 
    // Essentially, we only update the position of particles near the wheel, and freeze all particles outside the box
    // This likely causes some simualtion inaccuracies. We should perform a sensitivity study to see if this is the case.
    DEMSim_.DoDynamicsThenSync(0.0f);
    DEMSim_.ChangeClumpFamily(Family::FIXED);

    size_t num_changed = 0;
    // Retrieve wheel position
    float x = wheel_tracker_->Pos().x;
    float y = wheel_tracker_->Pos().y;

    std::pair<float, float> Xrange = {x - box_halfsize_x, x + box_halfsize_x};
    std::pair<float, float> Yrange = {y - box_halfsize_y, y + box_halfsize_y};
    num_changed += DEMSim_.ChangeClumpFamily(Family::FREE, Xrange, Yrange);

    std::cout << num_changed << " particles changed family number." << std::endl;
}

void WheelSimulator::WriteFrameData(double t, float3 forces) {
    // Write a new row of summary data to output.csv.
    try {
        output_datafile_<< t << "," 
                        << forces.x << "," 
                        << forces.y << "," 
                        << forces.z << ","
                        << forces.x/total_pressure_ << "," 
                        << max_v_finder_->GetValue() << ","
                        << wheel_tracker_->Pos().x << ","
                        << wheel_tracker_->Pos().y << ","
                        << wheel_tracker_->Pos().z << ","
                        << wheel_tracker_->OriQ().x << ","
                        << wheel_tracker_->OriQ().y << ","
                        << wheel_tracker_->OriQ().z << ","
                        << wheel_tracker_->OriQ().w << ","
                        << wheel_tracker_->Vel().x << ","
                        << wheel_tracker_->Vel().y << ","
                        << wheel_tracker_->Vel().z
                        << std::endl;
        output_datafile_.flush();
    } catch (...) {
        throw std::runtime_error("Unable to write content to output file.");
    }
}

void WheelSimulator::RunSimulationLoop() {
    std::cout << "Output at " << Constants::FPS << " FPS" << std::endl;
    // The main sim loop currently runs at twice the step size as the settling phase
    // TODO: This is likely confusing; we should split these into two parameters, e.g. step_size_settle, step_size_drive
    simparams_.step_size *= 2.;
    DEMSim_.UpdateStepSize(simparams_.step_size);

    // Start simulation timer
    auto start = std::chrono::high_resolution_clock::now();

    // Active box domain parameters
    float box_halfsize_x = wheel_.r_outer * 1.25f;
    float box_halfsize_y = wheel_.width * 2.0f;

    for (double t = 0.0; t < simparams_.sim_endtime; t += simparams_.step_size, curr_step_++) {
        if (curr_step_ % out_steps_ == 0) {
            UpdateActiveBoxDomain(box_halfsize_x, box_halfsize_y);

            // Write output files
            WriteWheelMesh();
            WriteParticleCSV();

            std::cout << "Outputting frame: " << currframe_ << std::endl;
            currframe_++;
            DEMSim_.ShowThreadCollaborationStats();
        }

        if (curr_step_ % report_steps_ == 0) {
            // Retrieve forces on the wheel
            float3 forces = wheel_tracker_->ContactAcc() * wheel_.mass;

            // Output to terminal
            std::cout << "Time: " << t << std::endl;
            std::cout << "Force on wheel: " << forces.x << ", " << forces.y << ", " << forces.z << std::endl;
            std::cout << "Drawbar pull coeff: " << (forces.x / total_pressure_) << std::endl;

            WriteFrameData(t, forces);

            // Termination condition
            // if (wheel_tracker_->Pos().x > (world_size_x / 2.0f - wheel_.r_outer * 1.2f)) {
            //     std::cout << "This is far enough, stopping the simulation..." << std::endl;
            //     DEMSim_.DoDynamicsThenSync(0.0f);
            //     break;
            // }
        }

        DEMSim_.DoDynamics(simparams_.step_size);
    }

    // End simulation timer
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> time_sec = std::chrono::duration_cast<std::chrono::duration<double>>(end - start);
    std::cout << time_sec.count() << " seconds (wall time) to finish the simulation." << std::endl;

    // Show stats
    DEMSim_.ShowTimingStats();
    DEMSim_.ShowAnomalies();

    std::cout << "WheelSimulator demo exiting..." << std::endl;
}