// ss_mesh — run the stage-1 SS decoder and export the occupancy isosurface as
// a Wavefront OBJ via (tetrahedral) marching cubes.
//
//   usage: ss_mesh <ss_dec.gguf> <z_s.latent> [out.obj] [--iso V] [--normalize]
//
// Decodes z_s -> a 64^3 occupancy logit grid, extracts the {logit = iso} surface
// (default iso 0, i.e. the occupancy boundary), and writes it as out.obj
// (default ss_mesh.obj). With --normalize, vertices are mapped from grid-index
// units into the centered unit cube [-0.5, 0.5]^3.

#include "trellis2.h"
#include "marching_cubes.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <string>
#include <vector>

int main(int argc, char ** argv) {
    if (argc < 3) {
        std::fprintf(stderr, "usage: %s <ss_dec.gguf> <z_s.latent> [out.obj] [--iso V] [--normalize]\n", argv[0]);
        return 2;
    }
    const std::string gguf_path = argv[1];
    const std::string lat_path  = argv[2];
    std::string out_path = "ss_mesh.obj";
    float iso = 0.0f;
    bool normalize = false;
    bool out_set = false;
    for (int i = 3; i < argc; ++i) {
        if (std::strcmp(argv[i], "--iso") == 0 && i + 1 < argc) iso = (float) std::atof(argv[++i]);
        else if (std::strcmp(argv[i], "--normalize") == 0) normalize = true;
        else if (!out_set) { out_path = argv[i]; out_set = true; }
    }

    std::printf("trellis2.cpp %s\n", trellis2_version());

    std::string err;
    trellis2_ss_dec_model * m = trellis2_ss_dec_load(gguf_path, true, &err);
    if (!m) { std::fprintf(stderr, "model load error: %s\n", err.c_str()); return 1; }
    std::printf("backend: %s\n", trellis2_ss_dec_backend_name(m));

    const trellis2_ss_dec_hparams hp = trellis2_ss_dec_hparams_of(m);
    const int    Rin  = hp.res_in();
    const int    Rout = hp.res_out();
    const size_t n_in  = (size_t) hp.latent_channels * Rin * Rin * Rin;
    const size_t n_out = (size_t) hp.out_channels * Rout * Rout * Rout;

    std::ifstream f(lat_path, std::ios::binary);
    if (!f) { std::fprintf(stderr, "cannot open latent %s\n", lat_path.c_str()); trellis2_ss_dec_free(m); return 1; }
    std::vector<float> latent(n_in);
    f.read(reinterpret_cast<char *>(latent.data()), (std::streamsize) (n_in * sizeof(float)));
    if ((size_t) f.gcount() != n_in * sizeof(float)) {
        std::fprintf(stderr, "latent size mismatch: got %zu floats, want %zu\n",
                     (size_t) f.gcount() / sizeof(float), n_in);
        trellis2_ss_dec_free(m); return 1;
    }

    std::vector<float> logits(n_out, 0.0f);
    if (!trellis2_ss_dec_decode(m, latent.data(), logits.data(), &err)) {
        std::fprintf(stderr, "decode error: %s\n", err.c_str());
        trellis2_ss_dec_free(m); return 1;
    }
    trellis2_ss_dec_free(m);

    size_t occ = 0;
    for (float v : logits) if (v > iso) ++occ;
    std::printf("logits : [%d^3]  occupied(>%.2f)=%zu/%zu (%.2f%%)\n",
                Rout, iso, occ, n_out, 100.0 * (double) occ / (double) n_out);

    // Decoder output is channel-major [1, R, R, R] with linear index
    // i*R^2 + j*R + k (k fastest) -> matches marching_cubes' x + y*R + z*R^2 with
    // (x,y,z) = (k, j, i).
    mc::Mesh mesh = mc::extract(logits.data(), Rout, Rout, Rout, iso);
    std::printf("mesh   : %zu verts, %zu tris\n", mesh.n_verts(), mesh.n_tris());

    if (mesh.n_tris() == 0) {
        std::fprintf(stderr, "warning: empty surface at iso=%.3f (nothing to write)\n", iso);
        return 1;
    }

    if (normalize) {
        for (size_t i = 0; i < mesh.verts.size(); ++i)
            mesh.verts[i] = mesh.verts[i] / (float) Rout - 0.5f;
    }

    if (!mc::write_obj(mesh, out_path.c_str())) {
        std::fprintf(stderr, "failed to write %s\n", out_path.c_str());
        return 1;
    }
    std::printf("wrote %s\n", out_path.c_str());
    return 0;
}