// CFG parity: f5_dit_forward_cfg (batched cond+uncond) must match the real
// f5_tts DiT run with cfg_infer=True — i.e. the null half sees ZEROED cond
// audio and FILLER text ids (python drop_audio_cond=True / drop_text=True).
// Goldens generated by /mnt/ai/f5-parity/gen_golden_cfg.py.
#include "engine/community_models/f5_tts/runtime.h"

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

namespace {

std::vector<float> load_bin(const std::string & path, size_t expected) {
    std::ifstream f(path, std::ios::binary);
    if (!f) {
        std::fprintf(stderr, "cannot open %s\n", path.c_str());
        std::exit(1);
    }
    std::vector<float> out(expected);
    f.read(reinterpret_cast<char *>(out.data()), static_cast<std::streamsize>(expected * sizeof(float)));
    if (static_cast<size_t>(f.gcount()) != expected * sizeof(float)) {
        std::fprintf(stderr, "short read on %s\n", path.c_str());
        std::exit(1);
    }
    return out;
}

std::vector<int32_t> load_ids(const std::string & path, size_t expected) {
    std::ifstream f(path, std::ios::binary);
    if (!f) {
        std::fprintf(stderr, "cannot open %s\n", path.c_str());
        std::exit(1);
    }
    std::vector<int32_t> out(expected);
    f.read(reinterpret_cast<char *>(out.data()), static_cast<std::streamsize>(expected * sizeof(int32_t)));
    if (static_cast<size_t>(f.gcount()) != expected * sizeof(int32_t)) {
        std::fprintf(stderr, "short read on %s\n", path.c_str());
        std::exit(1);
    }
    return out;
}

double cosine(const std::vector<float> & a, const std::vector<float> & b) {
    double dot = 0, na = 0, nb = 0;
    for (size_t i = 0; i < a.size(); ++i) {
        dot += static_cast<double>(a[i]) * b[i];
        na += static_cast<double>(a[i]) * a[i];
        nb += static_cast<double>(b[i]) * b[i];
    }
    return dot / (std::sqrt(na) * std::sqrt(nb) + 1e-12);
}

double max_abs(const std::vector<float> & a, const std::vector<float> & b) {
    double m = 0;
    for (size_t i = 0; i < a.size(); ++i) {
        m = std::max(m, std::abs(static_cast<double>(a[i]) - b[i]));
    }
    return m;
}

}  // namespace

int main(int argc, char ** argv) {
    const std::string gold = "/mnt/ai/f5-parity/golden";
    const std::string ckpt = argc > 1 ? argv[1] : "/mnt/ai/models/Habibi-TTS/Unified/model_200000.safetensors";
    constexpr int N = 64, MEL = 100, NT = 24;

    const auto x = load_bin(gold + "/cfg_input_x.bin", N * MEL);
    const auto cond = load_bin(gold + "/cfg_input_cond.bin", N * MEL);
    const auto ids = load_ids(gold + "/cfg_input_ids.bin", NT);
    const auto want_cond = load_bin(gold + "/cfg_out_cond.bin", N * MEL);
    const auto want_null = load_bin(gold + "/cfg_out_null.bin", N * MEL);

    engine::models::f5_tts::F5Architecture arch;
    engine::models::f5_tts::F5ComputeDevice dev;
    dev.use_cuda = std::getenv("F5_CUDA") != nullptr;
    dev.device = 1;  // GPU 1 is the TTS GPU on this rig

    const auto [got_cond, got_null] = engine::models::f5_tts::f5_dit_forward_cfg(
        ckpt, x, cond, ids, 0.42F, N, arch, &dev);

    int failures = 0;
    auto check = [&](const char * name, const std::vector<float> & got, const std::vector<float> & want) {
        const double c = cosine(got, want);
        const double m = max_abs(got, want);
        const bool ok = c >= 0.999 && std::isfinite(c);
        std::printf("%-14s cosine=%.6f maxabs=%.5f %s\n", name, c, m, ok ? "OK" : "FAIL");
        if (!ok) ++failures;
    };
    check("cfg_cond", got_cond, want_cond);
    check("cfg_null", got_null, want_null);  // fails if the null half sees real cond/text

    if (failures > 0) {
        std::printf("CFG PARITY FAILED (%d halves)\n", failures);
        return 1;
    }
    std::printf("CFG PARITY PASSED\n");
    return 0;
}