/*
 * Copyright (c) 2022 Arm Limited. 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
 *
 * 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.
 */

#pragma once

#include "tensorflow/lite/schema/schema_generated.h"

#include <stdlib.h>
#include <string>

namespace InferenceProcess {

template <typename T, typename U>
class Array {
public:
    Array() = delete;
    Array(T *const data, U &size, size_t capacity) : _data{data}, _size{size}, _capacity{capacity} {}

    auto size() const {
        return _size;
    }

    auto capacity() const {
        return _capacity;
    }

    void push_back(const T &data) {
        _data[_size++] = data;
    }

private:
    T *const _data;
    U &_size;
    const size_t _capacity{};
};

template <typename T, typename U>
Array<T, U> makeArray(T *const data, U &size, size_t capacity) {
    return Array<T, U>{data, size, capacity};
}

class InferenceParser {
public:
    const tflite::Model *getModel(const void *buffer, size_t size) {
        // Verify buffer
        flatbuffers::Verifier base_verifier(reinterpret_cast<const uint8_t *>(buffer), size);
        if (!tflite::VerifyModelBuffer(base_verifier)) {
            printf("Warning: the model is not valid\n");
            return nullptr;
        }

        // Create model handle
        const tflite::Model *model = tflite::GetModel(buffer);
        if (model->subgraphs() == nullptr) {
            printf("Warning: nullptr subgraph\n");
            return nullptr;
        }

        return model;
    }

    template <typename T, typename U, size_t S>
    bool parseModel(const void *buffer, size_t size, char (&description)[S], T &&ifmDims, U &&ofmDims) {
        const tflite::Model *model = getModel(buffer, size);
        if (model == nullptr) {
            return true;
        }
        strncpy(description, model->description()->c_str(), sizeof(description));

        // Get input dimensions for first subgraph
        auto *subgraph = *model->subgraphs()->begin();
        bool failed    = getSubGraphDims(subgraph, subgraph->inputs(), ifmDims);
        if (failed) {
            return true;
        }

        // Get output dimensions for last subgraph
        subgraph = *model->subgraphs()->rbegin();
        failed   = getSubGraphDims(subgraph, subgraph->outputs(), ofmDims);
        if (failed) {
            return true;
        }

        return false;
    }

private:
    bool getShapeSize(const flatbuffers::Vector<int32_t> *shape, size_t &size) {
        size = 1;

        if (shape == nullptr) {
            printf("Warning: nullptr shape size.\n");
            return true;
        }

        if (shape->Length() == 0) {
            printf("Warning: shape zero length.\n");
            return true;
        }

        for (auto it = shape->begin(); it != shape->end(); ++it) {
            size *= *it;
        }

        return false;
    }

    bool getTensorTypeSize(const enum tflite::TensorType type, size_t &size) {
        switch (type) {
        case tflite::TensorType::TensorType_UINT8:
        case tflite::TensorType::TensorType_INT8:
            size = 1;
            break;
        case tflite::TensorType::TensorType_INT16:
            size = 2;
            break;
        case tflite::TensorType::TensorType_INT32:
        case tflite::TensorType::TensorType_FLOAT32:
            size = 4;
            break;
        default:
            printf("Warning: Unsupported tensor type\n");
            return true;
        }

        return false;
    }

    template <typename T>
    bool getSubGraphDims(const tflite::SubGraph *subgraph, const flatbuffers::Vector<int32_t> *tensorMap, T &dims) {
        if (subgraph == nullptr || tensorMap == nullptr) {
            printf("Warning: nullptr subgraph or tensormap.\n");
            return true;
        }

        if ((dims.capacity() - dims.size()) < tensorMap->size()) {
            printf("Warning: tensormap size is larger than dimension capacity.\n");
            return true;
        }

        for (auto index = tensorMap->begin(); index != tensorMap->end(); ++index) {
            auto tensor = subgraph->tensors()->Get(*index);
            size_t size;
            size_t tensorSize;

            bool failed = getShapeSize(tensor->shape(), size);
            if (failed) {
                return true;
            }

            failed = getTensorTypeSize(tensor->type(), tensorSize);
            if (failed) {
                return true;
            }

            size *= tensorSize;

            if (size > 0) {
                dims.push_back(size);
            }
        }

        return false;
    }
};

} // namespace InferenceProcess