/*
 * Copyright (C) 2024–2025 Amlogic, Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#include "postprocess.h"
#include <iostream>
#include <fstream>
#include <numeric>
#include <algorithm>

void preprocess(const cv::Mat& src, float* dst) {
    cv::Mat rgb, resized;
    cv::cvtColor(src, rgb, cv::COLOR_BGR2RGB);
    cv::resize(rgb, resized, cv::Size(kInputW, kInputH));

    for (int i = 0; i < kInputH; ++i) {
        for (int j = 0; j < kInputW; ++j) {
            cv::Vec3f pixel = resized.at<cv::Vec3b>(i, j);
            int idx = (i * kInputW + j) * 3;
            dst[idx + 0] = (pixel[0] - MEAN[0]) / STD[0];
            dst[idx + 1] = (pixel[1] - MEAN[1]) / STD[1];
            dst[idx + 2] = (pixel[2] - MEAN[2]) / STD[2];
        }
    }
}

void postprocess_topk(float* logits, int size, const std::vector<std::string>& labels, int k) {
    std::vector<int> indices(size);
    std::iota(indices.begin(), indices.end(), 0);

    std::partial_sort(indices.begin(), indices.begin() + k, indices.end(),
                      [&](int a, int b) { return logits[a] > logits[b]; });

    std::cout << "\nTop-" << k << " Results:" << std::endl;
    for (int i = 0; i < k; ++i) {
        int idx = indices[i];
        std::string name = (idx < (int)labels.size()) ? labels[idx] : "N/A";
        printf("%d: %-20s  score=%.6f\n", i + 1, name.c_str(), logits[idx]);
    }
}

std::vector<std::string> load_labels(const std::string& path) {
    std::vector<std::string> labels;
    std::ifstream f(path);
    std::string line;
    while (std::getline(f, line)) {
        if(!line.empty()) labels.push_back(line);
    }
    return labels;
}