/*
* Prior boxes for RetinaNet
* Tomas Goldmann,2023
*/

#include <iostream>
#include <vector>
#include <cmath>
#include <map>
#include <algorithm>
#include "prior_boxes.hpp"

PriorBox::PriorBox(std::vector<int> image_size, std::string phase)
{

    std::map<std::string, std::vector<std::vector<float>>> cfg;
    this->min_sizes = {{16.0f, 32.0f}, {64.0f, 128.0f}, {256.0f, 512.0f}};
    this->image_size = {image_size[1], image_size[0]};
    this->steps = {8, 16, 32};

    for (const float step : steps)
    {
        this->feature_maps.push_back({(int)std::ceil(image_size[0] / step), (int)std::ceil(image_size[1] / step)});
    }
}

std::vector<std::vector<float>> PriorBox::forward()
{

    std::vector<float> anchors;
    for (size_t k = 0; k < this->feature_maps.size(); k++)
    {
        const auto &f = this->feature_maps[k];
        const auto &min_sizes = this->min_sizes[k];

        for (int i = 0; i < f[1]; i++)
        {
            for (int j = 0; j < f[0]; j++)
            {
                for (const auto &min_size : min_sizes)
                {
                    float s_kx = min_size / this->image_size[1];
                    float s_ky = min_size / this->image_size[0];
                    std::vector<float> dense_cx = {static_cast<float>(j + 0.5) * this->steps[k] / this->image_size[1]};
                    std::vector<float> dense_cy = {static_cast<float>(i + 0.5) * this->steps[k] / this->image_size[0]};

                    for (const auto &cy : dense_cy)
                    {
                        for (const auto &cx : dense_cx)
                        {

                            anchors.push_back(cx);
                            anchors.push_back(cy);
                            anchors.push_back(s_kx);
                            anchors.push_back(s_ky);
                        }
                    }
                }
            }
        }
    }

    std::vector<std::vector<float>> output;

    size_t num_anchors = anchors.size() / 4;
    for (size_t i = 0; i < num_anchors; i++)
    {
        for (size_t j = 0; j < 4; j++)
        {
            anchors[i * 4 + j] = std::min(std::max(anchors[i * 4 + j], 0.0f), 1.0f);
        }
        output.push_back({anchors[i * 4], anchors[i * 4 + 1], anchors[i * 4 + 2], anchors[i * 4 + 3]});
    }
    //std::cout << "len  " << output.size() << std::endl;

    return output;
}

std::vector<std::vector<float>> decode(const std::vector<std::vector<float>> &loc, const std::vector<std::vector<float>> &priors, const std::vector<float> &variances)
{
    std::vector<std::vector<float>> boxes;
    for (size_t i = 0; i < loc.size(); i++)
    {
        const auto &prior = priors[i];
        const auto &loc_pred = loc[i];
        float prior_x = prior[0];
        float prior_y = prior[1];
        float prior_w = prior[2];
        float prior_h = prior[3];
        float var_x = variances[0];
        float var_y = variances[1];
        float var_w = variances[2];
        float var_h = variances[3];

        float decoded_x = prior_x + loc_pred[0] * var_x * prior_w;
        float decoded_y = prior_y + loc_pred[1] * var_y * prior_h;
        float decoded_w = prior_w * std::exp(loc_pred[2] * var_w);
        float decoded_h = prior_h * std::exp(loc_pred[3] * var_h);

        float decoded_xmin = decoded_x - decoded_w / 2;
        float decoded_ymin = decoded_y - decoded_h / 2;
        float decoded_xmax = decoded_x + decoded_w / 2;
        float decoded_ymax = decoded_y + decoded_h / 2;

        boxes.push_back({decoded_xmin, decoded_ymin, decoded_xmax, decoded_ymax});
    }

    return boxes;
}
