#include "diarizen_segmenter.h"

#include <cmath>
#include <cstring>
#include <algorithm>
#include <numeric>

// AX Engine runtime API placeholder.
// When AX_RUNTIME_ROOT is provided, include the real headers:
// #include "ax_engine.h"

namespace diarizen {

struct DiarizenSegmenter::Impl {
    std::string cnn_path;
    std::string backend_path;
    // void* cnn_handle = nullptr;  // AX engine handle
};

DiarizenSegmenter::DiarizenSegmenter(
    const std::string& cnn_model_path,
    const std::string& backend_onnx_path)
    : impl_(std::make_unique<Impl>())
{
    impl_->cnn_path = cnn_model_path;
    impl_->backend_path = backend_onnx_path;
    // TODO: Load CNN axmodel via AX Engine API
    // TODO: Load backend ONNX via ONNX Runtime C++ API
}

DiarizenSegmenter::~DiarizenSegmenter() = default;

SegmentResult DiarizenSegmenter::run(const float* audio, int num_samples) {
    SegmentResult result;
    result.log_probs.resize(result.num_frames * result.num_classes, 0.0f);

    // Preprocessing: LayerNorm on CPU
    if (num_samples != 64000) {
        // Input must be exactly 64000 samples (4s @ 16kHz)
        return result;
    }

    float mean = 0.0f;
    for (int i = 0; i < num_samples; ++i) mean += audio[i];
    mean /= num_samples;

    float var = 0.0f;
    for (int i = 0; i < num_samples; ++i) {
        float d = audio[i] - mean;
        var += d * d;
    }
    var = var / num_samples + 1e-5f;
    float inv_std = 1.0f / std::sqrt(var);

    std::vector<float> normalized(num_samples);
    for (int i = 0; i < num_samples; ++i) {
        normalized[i] = (audio[i] - mean) * inv_std;
    }

    // TODO: Run CNN NPU inference
    // TODO: Run backend ONNX inference
    // Placeholder: fill with uniform log(1/11) = -2.398
    float uniform_log_prob = std::log(1.0f / result.num_classes);
    std::fill(result.log_probs.begin(), result.log_probs.end(), uniform_log_prob);

    return result;
}

}  // namespace diarizen