#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <cstring>
#include "acoustic_analyzer/core/gunshot_classifier.h"
#include "acoustic_analyzer/core/feature_extractor.h"
#include "acoustic_analyzer/io/wav_file_source.h"

using namespace acoustic;

void print_usage(const char* prog) {
    std::cerr << "Usage: " << prog << " --model <onnx> --wav <wav_file> [--label_map <json>]" << std::endl;
}

int main(int argc, char** argv) {
    std::string model_path;
    std::string wav_path;
    std::string label_map_path;

    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) {
            model_path = argv[++i];
        } else if (std::strcmp(argv[i], "--wav") == 0 && i + 1 < argc) {
            wav_path = argv[++i];
        } else if (std::strcmp(argv[i], "--label_map") == 0 && i + 1 < argc) {
            label_map_path = argv[++i];
        }
    }

    if (model_path.empty() || wav_path.empty()) {
        print_usage(argv[0]);
        return 1;
    }

    WavFileSource wav(wav_path);
    if (!wav.open()) {
        std::cerr << "Failed to open WAV: " << wav_path << std::endl;
        return 1;
    }

    std::vector<std::vector<float>> audio;
    size_t n = wav.read(audio, wav.total_samples());
    if (n == 0 || audio.empty()) {
        std::cerr << "Empty WAV file" << std::endl;
        return 1;
    }

    FeatureExtractor extractor(wav.sample_rate(), 2048, 512, 64);
    Eigen::MatrixXf mel = extractor.MelSpectrogram(audio[0]);

    ClassifierConfig cc;
    cc.model_path = model_path;
    cc.label_map_path = label_map_path;
    cc.threshold = 0.7f;
    cc.smoothing_window = 3;
    GunshotClassifier classifier(cc);
    if (!classifier.IsLoaded()) {
        std::cerr << "Classifier init failed" << std::endl;
        return 1;
    }

    auto [label, confidence] = classifier.Predict(mel);
    std::cout << "Label: " << label << std::endl;
    std::cout << "Confidence: " << confidence << std::endl;
    return 0;
}
