/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 "regionPlugin.h"
#include <memory>
#include <string_view>

namespace nvinfer1::plugin
{
namespace
{
char const* const kREGION_PLUGIN_VERSION{"1"};
char const* const kREGION_PLUGIN_NAME{"Region_TRT"};

template <typename T>
void safeFree(T* ptr)
{
    if (ptr)
    {
        free(ptr);
        ptr = nullptr;
    }
}

template <typename T>
void allocateChunk(T*& ptr, int32_t count)
{
    ptr = static_cast<T*>(malloc(count * sizeof(T)));
}

struct SoftmaxTreeDeleter
{
    void operator()(softmaxTree* smTree) const
    {
        if (smTree)
        {
            // free individual elements first
            safeFree(smTree->leaf);
            safeFree(smTree->parent);
            safeFree(smTree->child);
            safeFree(smTree->group);
            if (smTree->name)
            {
                for (int32_t i = 0; i < smTree->n; i++)
                {
                    safeFree(smTree->name[i]);
                }
                safeFree(smTree->name);
            }
            safeFree(smTree->groupSize);
            safeFree(smTree->groupOffset);

            // free softmax tree
            safeFree(smTree);
        }
    }
};

} // namespace

Region::Region(RegionParameters params)
    : num(params.num)
    , coords(params.coords)
    , classes(params.classes)
    , smTree(params.smTree, SoftmaxTreeDeleter())
{
}

Region::Region(RegionParameters params, int32_t C, int32_t H, int32_t W)
    : num(params.num)
    , coords(params.coords)
    , classes(params.classes)
    , smTree(params.smTree, SoftmaxTreeDeleter())
    , C(C)
    , H(H)
    , W(W)
{
}

// NOLINTNEXTLINE(readability-function-cognitive-complexity)
Region::Region(void const* buffer, size_t length)
{
    char const *d = reinterpret_cast<char const*>(buffer), *a = d;
    C = read<int32_t>(d);
    H = read<int32_t>(d);
    W = read<int32_t>(d);
    num = read<int32_t>(d);
    classes = read<int32_t>(d);
    coords = read<int32_t>(d);
    bool softmaxTreePresent = read<bool>(d);
    bool leafPresent = read<bool>(d);
    bool parentPresent = read<bool>(d);
    bool childPresent = read<bool>(d);
    bool groupPresent = read<bool>(d);
    bool namePresent = read<bool>(d);
    bool groupSizePresent = read<bool>(d);
    bool groupOffsetPresent = read<bool>(d);
    if (softmaxTreePresent)
    {
        softmaxTree* smTreeTemp;
        // need to read each element individually
        allocateChunk(smTreeTemp, 1);

        smTreeTemp->n = read<int32_t>(d);

        if (leafPresent)
        {
            allocateChunk(smTreeTemp->leaf, smTreeTemp->n);
        }
        else
        {
            smTreeTemp->leaf = nullptr;
        }
        if (parentPresent)
        {
            allocateChunk(smTreeTemp->parent, smTreeTemp->n);
        }
        else
        {
            smTreeTemp->parent = nullptr;
        }
        if (childPresent)
        {
            allocateChunk(smTreeTemp->child, smTreeTemp->n);
        }
        else
        {
            smTreeTemp->child = nullptr;
        }
        if (groupPresent)
        {
            allocateChunk(smTreeTemp->group, smTreeTemp->n);
        }
        else
        {
            smTreeTemp->group = nullptr;
        }

        for (int32_t i = 0; i < smTreeTemp->n; i++)
        {
            if (leafPresent)
            {
                smTreeTemp->leaf[i] = read<int32_t>(d);
            }
            if (parentPresent)
            {
                smTreeTemp->parent[i] = read<int32_t>(d);
            }
            if (childPresent)
            {
                smTreeTemp->child[i] = read<int32_t>(d);
            }
            if (groupPresent)
            {
                smTreeTemp->group[i] = read<int32_t>(d);
            }
        }

        if (namePresent)
        {
            allocateChunk(smTreeTemp->name, smTreeTemp->n);
        }
        else
        {
            smTreeTemp->name = nullptr;
        }

        if (namePresent)
        {
            for (int32_t i = 0; i < smTreeTemp->n; i++)
            {
                allocateChunk(smTreeTemp->name[i], 256);
                for (int32_t j = 0; j < 256; j++)
                {
                    smTreeTemp->name[i][j] = read<char>(d);
                }
            }
        }

        smTreeTemp->groups = read<int32_t>(d);
        if (groupSizePresent)
        {
            allocateChunk(smTreeTemp->groupSize, smTreeTemp->groups);
        }
        else
        {
            smTreeTemp->groupSize = nullptr;
        }
        if (groupOffsetPresent)
        {
            allocateChunk(smTreeTemp->groupOffset, smTreeTemp->groups);
        }
        else
        {
            smTreeTemp->groupOffset = nullptr;
        }
        for (int32_t i = 0; i < smTreeTemp->groups; i++)
        {
            if (groupSizePresent)
            {
                smTreeTemp->groupSize[i] = read<int32_t>(d);
            }
            if (groupOffsetPresent)
            {
                smTreeTemp->groupOffset[i] = read<int32_t>(d);
            }
        }
        smTree = std::shared_ptr<softmaxTree>(smTreeTemp, SoftmaxTreeDeleter());
    }
    else
    {
        smTree.reset();
    }
    PLUGIN_VALIDATE(d == a + length);
}

int32_t Region::getNbOutputs() const noexcept
{
    return 1;
}

Dims Region::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
{
    PLUGIN_ASSERT(nbInputDims == 1);
    PLUGIN_ASSERT(index == 0);
    return inputs[0];
}

int32_t Region::enqueue(
    int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
    void const* inputData = inputs[0];
    void* outputData = outputs[0];
    if (smTree.get())
    {
        hasSoftmaxTree = true;
    }
    else
    {
        hasSoftmaxTree = false;
    }
    pluginStatus_t status = regionInference(
        stream, batchSize, C, H, W, num, coords, classes, hasSoftmaxTree, smTree.get(), inputData, outputData);
    return status;
}

size_t Region::getSerializationSize() const noexcept
{
    // C, H, W, num, classes, coords, smTree !nullptr and other array members !nullptr, softmaxTree members
    size_t count = 6 * sizeof(int32_t) + 8 * sizeof(bool);
    if (smTree.get())
    {
        count += 2 * sizeof(int32_t);

        if (smTree->leaf)
        {
            count += smTree->n * sizeof(int32_t);
        }
        if (smTree->parent)
        {
            count += smTree->n * sizeof(int32_t);
        }
        if (smTree->child)
        {
            count += smTree->n * sizeof(int32_t);
        }
        if (smTree->group)
        {
            count += smTree->n * sizeof(int32_t);
        }
        if (smTree->name)
        {
            count += smTree->n * 256 * sizeof(char);
        }
        if (smTree->groupSize)
        {
            count += smTree->groups * sizeof(int32_t);
        }
        if (smTree->groupOffset)
        {
            count += smTree->groups * sizeof(int32_t);
        }
    }
    return count;
}

// NOLINTNEXTLINE(readability-function-cognitive-complexity)
void Region::serialize(void* buffer) const noexcept
{
    char *d = reinterpret_cast<char*>(buffer), *a = d;
    write(d, C);
    write(d, H);
    write(d, W);
    write(d, num);
    write(d, classes);
    write(d, coords);
    write(d, smTree != nullptr);
    write(d, smTree != nullptr && smTree->leaf != nullptr);
    write(d, smTree != nullptr && smTree->parent != nullptr);
    write(d, smTree != nullptr && smTree->child != nullptr);
    write(d, smTree != nullptr && smTree->group != nullptr);
    write(d, smTree != nullptr && smTree->name != nullptr);
    write(d, smTree != nullptr && smTree->groupSize != nullptr);
    write(d, smTree != nullptr && smTree->groupOffset != nullptr);
    // need to do a deep copy
    if (smTree)
    {
        write(d, smTree->n);
        for (int32_t i = 0; i < smTree->n; i++)
        {
            if (smTree->leaf)
            {
                write(d, smTree->leaf[i]);
            }
            if (smTree->parent)
            {
                write(d, smTree->parent[i]);
            }
            if (smTree->child)
            {
                write(d, smTree->child[i]);
            }
            if (smTree->group)
            {
                write(d, smTree->group[i]);
            }
        }
        if (smTree->name)
        {
            for (int32_t i = 0; i < smTree->n; i++)
            {
                char const* str = smTree->name[i];
                for (int32_t j = 0; j < 256; j++)
                {
                    write(d, str[j]);
                }
            }
        }
        write(d, smTree->groups);
        for (int32_t i = 0; i < smTree->groups; i++)
        {
            if (smTree->groupSize)
            {
                write(d, smTree->groupSize[i]);
            }
            if (smTree->groupOffset)
            {
                write(d, smTree->groupOffset[i]);
            }
        }
    }
    PLUGIN_ASSERT(d == a + getSerializationSize());
}

bool Region::supportsFormat(DataType type, PluginFormat format) const noexcept
{
    return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
}

int32_t Region::initialize() noexcept
{
    return STATUS_SUCCESS;
}

void Region::terminate() noexcept {}

char const* Region::getPluginType() const noexcept
{
    return kREGION_PLUGIN_NAME;
}

char const* Region::getPluginVersion() const noexcept
{
    return kREGION_PLUGIN_VERSION;
}

size_t Region::getWorkspaceSize(int32_t maxBatchSize) const noexcept
{
    return 0;
}

void Region::destroy() noexcept
{
    delete this;
}

IPluginV2Ext* Region::clone() const noexcept
{
    try
    {
        RegionParameters params{num, coords, classes, nullptr};
        auto plugin = std::make_unique<Region>(params, C, H, W);
        plugin->setPluginNamespace(mPluginNamespace.c_str());
        plugin->setSoftmaxTree(smTree);

        return plugin.release();
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}

// Set plugin namespace
void Region::setPluginNamespace(char const* pluginNamespace) noexcept
{
    mPluginNamespace = pluginNamespace;
}

char const* Region::getPluginNamespace() const noexcept
{
    return mPluginNamespace.c_str();
}

// Return the DataType of the plugin output at the requested index
DataType Region::getOutputDataType(int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
{
    PLUGIN_ASSERT(index == 0);
    return DataType::kFLOAT;
}

// Configure the layer with input and output data types.
void Region::configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims, int32_t nbOutputs,
    DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
    bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept
{
    PLUGIN_ASSERT(*inputTypes == DataType::kFLOAT && floatFormat == PluginFormat::kLINEAR);
    PLUGIN_ASSERT(nbInputs == 1);
    PLUGIN_ASSERT(nbOutputs == 1);
    C = inputDims[0].d[0];
    H = inputDims[0].d[1];
    W = inputDims[0].d[2];
    /*
     * In the below assertion, 1 stands for the objectness of the bounding box
     * We should also
     * PLUGIN_ASSERT(coords == 4);
     */
    PLUGIN_ASSERT(C == num * (coords + 1 + classes));
}

// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
void Region::attachToContext(
    cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
{
}

// Detach the plugin object from its execution context.
void Region::detachFromContext() noexcept {}

RegionPluginCreator::RegionPluginCreator()
{
    mPluginAttributes.clear();
    mPluginAttributes.emplace_back(PluginField("num", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("coords", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("classes", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("smTree", nullptr, PluginFieldType::kINT32, 1));

    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

char const* RegionPluginCreator::getPluginName() const noexcept
{
    return kREGION_PLUGIN_NAME;
}

char const* RegionPluginCreator::getPluginVersion() const noexcept
{
    return kREGION_PLUGIN_VERSION;
}

PluginFieldCollection const* RegionPluginCreator::getFieldNames() noexcept
{
    return &mFC;
}

IPluginV2Ext* RegionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
{
    try
    {
        using namespace std::string_view_literals;
        PluginField const* fields = fc->fields;
        for (int32_t i = 0; i < fc->nbFields; ++i)
        {
            std::string_view const attrName = fields[i].name;
            if (attrName == "num"sv)
            {
                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
                params.num = *(static_cast<int32_t const*>(fields[i].data));
            }
            if (attrName == "coords"sv)
            {
                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
                params.coords = *(static_cast<int32_t const*>(fields[i].data));
            }
            if (attrName == "classes"sv)
            {
                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
                params.classes = *(static_cast<int32_t const*>(fields[i].data));
            }
            if (attrName == "smTree"sv)
            {
                // TODO not sure if this will work
                void* tmpData = const_cast<void*>(fields[i].data);
                params.smTree = static_cast<nvinfer1::plugin::softmaxTree*>(tmpData);
            }
        }

        auto obj = std::make_unique<Region>(params);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj.release();
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}

IPluginV2Ext* RegionPluginCreator::deserializePlugin(
    char const* name, void const* serialData, size_t serialLength) noexcept
{
    try
    {
        // This object will be deleted when the network is destroyed, which will
        // call Region::destroy()
        auto obj = std::make_unique<Region>(serialData, serialLength);
        obj->setPluginNamespace(mNamespace.c_str());
        return obj.release();
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}
} // namespace nvinfer1::plugin