#include <iostream>
#include <iomanip>
#include <vector>
#include <string>
#include <filesystem>
#include <algorithm>
#include <cmath>
#include <numeric>
#include <map>
#include <chrono>
#include <cstring>

#include "acoustic_analyzer/core/pipeline.h"
#include "acoustic_analyzer/io/wav_file_source.h"

namespace fs = std::filesystem;
using namespace acoustic;

struct Prediction {
    std::string file_path;
    std::string true_label;
    std::string pred_label;
    float confidence = 0.0f;
    float azimuth = 0.0f;
    float elevation = 0.0f;
    float distance = -1.0f;
    bool detected = false;
};

void print_usage(const char* prog) {
    std::cerr << "Usage: " << prog << " <file_or_dir> [options]" << std::endl;
    std::cerr << "Options:" << std::endl;
    std::cerr << "  --model <path>           ONNX model path (default: models/gunshot_classifier.onnx)" << std::endl;
    std::cerr << "  --label_map <path>       Label map file (default: models/label_map.json)" << std::endl;
    std::cerr << "  --threshold <float>      Detection threshold (default: 0.5)" << std::endl;
    std::cerr << "  --num_mics <int>         Number of channels in WAV (default: 4)" << std::endl;
    std::cerr << "  --spacing <float>        Mic spacing in meters (default: 0.15)" << std::endl;
    std::cerr << "  --layout <str>           Array layout: cross/linear/circular (default: cross)" << std::endl;
    std::cerr << "  --ref_spl <float>        Reference SPL for distance estimation (default: 150)" << std::endl;
    std::cerr << "  --ground_azimuth <float> Ground-truth azimuth for error calc (optional)" << std::endl;
    std::cerr << "  --ground_distance <float> Ground-truth distance for error calc (optional)" << std::endl;
}

bool ends_with(const std::string& s, const std::string& suffix) {
    if (s.size() < suffix.size()) return false;
    return s.compare(s.size() - suffix.size(), suffix.size(), suffix) == 0;
}

std::string get_parent_folder_name(const std::string& path) {
    fs::path p(path);
    if (p.has_parent_path()) {
        return p.parent_path().filename().string();
    }
    return "";
}

// Convert vector-of-vectors [channels][samples] to flat interleaved [ch0_s0, ch1_s0, ...]
std::vector<float> flatten_audio(const std::vector<std::vector<float>>& audio, int channels) {
    if (audio.empty() || channels == 0) return {};
    size_t samples = audio[0].size();
    std::vector<float> flat(samples * channels);
    for (size_t s = 0; s < samples; ++s) {
        for (int ch = 0; ch < channels; ++ch) {
            flat[s * channels + ch] = (ch < static_cast<int>(audio.size()) && s < audio[ch].size())
                                          ? audio[ch][s] : 0.0f;
        }
    }
    return flat;
}

Prediction process_file(const std::string& path,
                        Pipeline& pipeline,
                        int num_mics,
                        float ground_azimuth,
                        float ground_distance) {
    Prediction result;
    result.file_path = path;
    result.true_label = get_parent_folder_name(path);

    WavFileSource wav(path);
    if (!wav.open()) {
        std::cerr << "[SKIP] Cannot open: " << path << std::endl;
        result.pred_label = "error";
        return result;
    }

    int sr = wav.sample_rate();
    int file_ch = wav.num_channels();
    if (file_ch != num_mics) {
        std::cerr << "[WARN] " << path << " has " << file_ch
                  << " channels, expected " << num_mics << std::endl;
    }
    // Use actual file channels if mismatch, but pipeline config should match
    int effective_channels = (file_ch < num_mics) ? file_ch : num_mics;

    size_t chunk_samples = static_cast<size_t>(sr * pipeline.Config().chunk_duration);
    std::vector<std::vector<float>> audio;
    size_t got = wav.read(audio, chunk_samples);
    if (got == 0 || audio.empty()) {
        result.pred_label = "empty";
        return result;
    }

    // Process only the first chunk (like demo_offline)
    // For sliding-window analysis, call Process() on hop-sized chunks
    auto flat = flatten_audio(audio, effective_channels);
    auto frame = pipeline.Process(flat);

    if (!frame.is_clear && !frame.threats.empty()) {
        const auto& t = frame.threats[0];
        result.detected = true;
        result.pred_label = t.sound_type;
        result.confidence = t.confidence;
        result.azimuth = t.azimuth;
        result.elevation = t.elevation;
        result.distance = t.distance;
    } else {
        result.pred_label = "ambient";
        result.confidence = 0.0f;
    }

    std::cout << "File: " << fs::path(path).filename().string()
              << " | True: " << result.true_label
              << " | Pred: " << result.pred_label
              << " | Conf: " << std::fixed << std::setprecision(4) << result.confidence;
    if (result.detected) {
        std::cout << " | Az: " << std::setprecision(2) << result.azimuth << "°"
                  << " | El: " << std::setprecision(2) << result.elevation << "°"
                  << " | Dist: " << std::setprecision(2) << result.distance << "m";
        if (ground_azimuth >= 0.0f) {
            float az_err = std::fabs(result.azimuth - ground_azimuth);
            if (az_err > 180.0f) az_err = 360.0f - az_err;
            std::cout << " | AzErr: " << az_err << "°";
        }
        if (ground_distance >= 0.0f) {
            std::cout << " | DistErr: " << std::fabs(result.distance - ground_distance) << "m";
        }
    }
    std::cout << std::endl;
    return result;
}

void collect_wav_files(const std::string& target, std::vector<std::string>& out) {
    if (fs::is_regular_file(target) && ends_with(target, ".wav")) {
        out.push_back(target);
        return;
    }
    if (!fs::is_directory(target)) return;

    for (const auto& entry : fs::recursive_directory_iterator(target)) {
        if (entry.is_regular_file() && ends_with(entry.path().string(), ".wav")) {
            out.push_back(entry.path().string());
        }
    }
    std::sort(out.begin(), out.end());
}

void print_report(const std::vector<Prediction>& results,
                  float ground_azimuth,
                  float ground_distance) {
    std::map<std::string, int> total_by_true;
    std::map<std::string, int> correct_by_true;
    std::map<std::string, float> conf_sum_by_true;
    std::map<std::string, std::map<std::string, int>> confusion;

    int total = 0, correct = 0;
    int detected_count = 0;
    float az_err_sum = 0.0f, dist_err_sum = 0.0f;
    int az_err_count = 0, dist_err_count = 0;

    for (const auto& r : results) {
        if (r.pred_label == "error" || r.pred_label == "empty") continue;
        total++;
        total_by_true[r.true_label]++;
        conf_sum_by_true[r.true_label] += r.confidence;
        confusion[r.true_label][r.pred_label]++;
        if (r.true_label == r.pred_label) {
            correct++;
            correct_by_true[r.true_label]++;
        }
        if (r.detected) {
            detected_count++;
            if (ground_azimuth >= 0.0f) {
                float az_err = std::fabs(r.azimuth - ground_azimuth);
                if (az_err > 180.0f) az_err = 360.0f - az_err;
                az_err_sum += az_err;
                az_err_count++;
            }
            if (ground_distance >= 0.0f) {
                dist_err_sum += std::fabs(r.distance - ground_distance);
                dist_err_count++;
            }
        }
    }

    std::cout << "\n==========================================" << std::endl;
    std::cout << "      MULTICHANNEL VALIDATION REPORT" << std::endl;
    std::cout << "==========================================" << std::endl;
    std::cout << "Total samples:    " << total << std::endl;
    std::cout << "Correct:          " << correct << std::endl;
    std::cout << "Accuracy:         " << std::fixed << std::setprecision(2)
              << (total > 0 ? 100.0f * correct / total : 0.0f) << "%" << std::endl;
    std::cout << "Detected frames:  " << detected_count << std::endl;

    if (az_err_count > 0) {
        std::cout << "Azimuth RMSE:     " << std::setprecision(2)
                  << std::sqrt(az_err_sum / az_err_count) << "°" << std::endl;
    }
    if (dist_err_count > 0) {
        std::cout << "Distance MAE:     " << std::setprecision(2)
                  << (dist_err_sum / dist_err_count) << "m" << std::endl;
    }

    std::cout << "\nPer-class breakdown:" << std::endl;
    for (const auto& kv : total_by_true) {
        const std::string& cls = kv.first;
        int cls_total = kv.second;
        int cls_correct = correct_by_true[cls];
        float avg_conf = conf_sum_by_true[cls] / cls_total;
        std::cout << "  " << std::setw(10) << std::left << cls
                  << "  Count: " << std::setw(3) << cls_total
                  << "  Correct: " << std::setw(3) << cls_correct
                  << "  Acc: " << std::setw(6) << std::fixed << std::setprecision(2)
                  << (100.0f * cls_correct / cls_total) << "%"
                  << "  AvgConf: " << std::setprecision(4) << avg_conf << std::endl;
    }

    std::cout << "\nConfusion matrix (rows=true, cols=pred):" << std::endl;
    std::vector<std::string> labels;
    for (const auto& row : confusion) labels.push_back(row.first);
    for (const auto& row : confusion) {
        for (const auto& col : row.second) {
            if (std::find(labels.begin(), labels.end(), col.first) == labels.end()) {
                labels.push_back(col.first);
            }
        }
    }
    std::sort(labels.begin(), labels.end());

    std::cout << std::setw(12) << " ";
    for (const auto& l : labels) std::cout << std::setw(10) << l;
    std::cout << std::endl;
    for (const auto& true_l : labels) {
        std::cout << std::setw(10) << std::left << true_l << "  ";
        for (const auto& pred_l : labels) {
            int count = confusion.count(true_l) ? confusion[true_l].count(pred_l) ? confusion[true_l].at(pred_l) : 0 : 0;
            std::cout << std::setw(10) << count;
        }
        std::cout << std::endl;
    }
    std::cout << "==========================================" << std::endl;
}

int main(int argc, char** argv) {
    if (argc < 2 || std::strcmp(argv[1], "--help") == 0 || std::strcmp(argv[1], "-h") == 0) {
        print_usage(argv[0]);
        return argc < 2 ? 1 : 0;
    }

    std::string target = argv[1];
    std::string model_path = "models/gunshot_classifier.onnx";
    std::string label_map_path = "models/label_map.json";
    float threshold = 0.5f;
    int num_mics = 4;
    float spacing = 0.15f;
    std::string layout = "cross";
    float ref_spl = 150.0f;
    float ground_azimuth = -1.0f;
    float ground_distance = -1.0f;

    for (int i = 2; i < argc; ++i) {
        if (std::strcmp(argv[i], "--model") == 0 && i + 1 < argc) model_path = argv[++i];
        else if (std::strcmp(argv[i], "--label_map") == 0 && i + 1 < argc) label_map_path = argv[++i];
        else if (std::strcmp(argv[i], "--threshold") == 0 && i + 1 < argc) threshold = std::stof(argv[++i]);
        else if (std::strcmp(argv[i], "--num_mics") == 0 && i + 1 < argc) num_mics = std::stoi(argv[++i]);
        else if (std::strcmp(argv[i], "--spacing") == 0 && i + 1 < argc) spacing = std::stof(argv[++i]);
        else if (std::strcmp(argv[i], "--layout") == 0 && i + 1 < argc) layout = argv[++i];
        else if (std::strcmp(argv[i], "--ref_spl") == 0 && i + 1 < argc) ref_spl = std::stof(argv[++i]);
        else if (std::strcmp(argv[i], "--ground_azimuth") == 0 && i + 1 < argc) ground_azimuth = std::stof(argv[++i]);
        else if (std::strcmp(argv[i], "--ground_distance") == 0 && i + 1 < argc) ground_distance = std::stof(argv[++i]);
    }

    // Build PipelineConfig directly (no yaml-cpp needed)
    PipelineConfig config;
    config.sample_rate = 16000;
    config.chunk_duration = 2.0f;
    config.hop_duration = 0.5f;
    config.n_mels = 64;
    config.confidence_threshold = threshold;
    config.classifier.model_path = model_path;
    config.classifier.label_map_path = label_map_path;
    config.classifier.threshold = threshold;
    config.classifier.smoothing_window = 1; // offline: no temporal smoothing
    config.mic_array.num_mics = static_cast<uint32_t>(num_mics);
    config.mic_array.layout = layout;
    config.mic_array.spacing = spacing;
    config.distance.ref_spl_gunshot = ref_spl;
    config.distance.ref_spl_artillery = ref_spl;
    config.distance.ref_spl_explosion = ref_spl;

    Pipeline pipeline(config);

    std::vector<std::string> files;
    collect_wav_files(target, files);
    if (files.empty()) {
        std::cerr << "No .wav files found in: " << target << std::endl;
        return 1;
    }
    std::cout << "Found " << files.size() << " WAV file(s)." << std::endl;
    std::cout << "Channels: " << num_mics << ", Layout: " << layout
              << ", Spacing: " << spacing << "m" << std::endl;
    std::cout << std::endl;

    std::vector<Prediction> results;
    results.reserve(files.size());
    double total_ms = 0.0;
    for (const auto& f : files) {
        auto t0 = std::chrono::steady_clock::now();
        results.push_back(process_file(f, pipeline, num_mics, ground_azimuth, ground_distance));
        auto t1 = std::chrono::steady_clock::now();
        total_ms += std::chrono::duration<double, std::milli>(t1 - t0).count();
    }

    std::cout << "\nTotal inference time: " << std::fixed << std::setprecision(2)
              << total_ms << " ms"
              << " | Avg per file: " << (files.empty() ? 0.0 : total_ms / files.size())
              << " ms" << std::endl;

    print_report(results, ground_azimuth, ground_distance);
    return 0;
}
