/*
 * Copyright (c) 2021 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.
 */

/****************************************************************************
 * Includes
 ****************************************************************************/

#include "command_stream.hpp"

#include <inttypes.h>
#include <stdio.h>

using namespace std;

namespace EthosU {
namespace CommandStream {

/****************************************************************************
 * DataPointer
 ****************************************************************************/

DataPointer::DataPointer() : data(nullptr), size(0) {}

DataPointer::DataPointer(const char *_data, size_t _size) : data(_data), size(_size) {}

bool DataPointer::operator!=(const DataPointer &other) {
    if (size != other.size) {
        return true;
    }

    for (size_t i = 0; i < size; i++) {
        if (data[i] != other.data[i]) {
            return true;
        }
    }

    return false;
}

/****************************************************************************
 * PmuConfig
 ****************************************************************************/

Pmu::Pmu(ethosu_driver *_drv, const PmuEvents &_config) :
    drv(_drv),
    config(_config)
{
    // Enable PMU block
    ETHOSU_PMU_Enable_v2(drv);

    // Enable cycle counter
    ETHOSU_PMU_CNTR_Enable_v2(drv, ETHOSU_PMU_CCNT_Msk);

    // Configure event types
    for (size_t i = 0; i < config.size(); i++) {
        ETHOSU_PMU_Set_EVTYPER_v2(drv, i, config[i]);
        ETHOSU_PMU_CNTR_Enable_v2(drv, 1 << i);
    }
}

void Pmu::clear() {
    ETHOSU_PMU_CYCCNT_Reset_v2(drv);
    ETHOSU_PMU_EVCNTR_ALL_Reset_v2(drv);
}

void Pmu::print() {
    printf("PMU={cycleCount=%llu, events=[%" PRIu32 ", %" PRIu32 ", %" PRIu32 ", %" PRIu32 "]}\n",
            ETHOSU_PMU_Get_CCNTR_v2(drv),
            ETHOSU_PMU_Get_EVCNTR_v2(drv, 0),
            ETHOSU_PMU_Get_EVCNTR_v2(drv, 1),
            ETHOSU_PMU_Get_EVCNTR_v2(drv, 2),
            ETHOSU_PMU_Get_EVCNTR_v2(drv, 3));
}

uint64_t Pmu::getCycleCount() const {
    return ETHOSU_PMU_Get_CCNTR_v2(drv);
}

uint32_t Pmu::getEventCount(size_t index) const {
    return ETHOSU_PMU_Get_EVCNTR_v2(drv, index);
}

/****************************************************************************
 * CommandStream
 ****************************************************************************/

CommandStream::CommandStream(const DataPointer &_commandStream,
                             const BasePointers &_basePointers,
                             const PmuEvents &_pmuEvents) :
    drv(ethosu_reserve_driver()),
    commandStream(_commandStream),
    basePointers(_basePointers),
    pmu(drv, _pmuEvents)
{
    // Disable clock gating, else the NPU PMU will be clock gated and report too few cycles
    ethosu_set_clock_and_power(&drv->dev, ETHOSU_CLOCK_Q_DISABLE, ETHOSU_POWER_Q_DISABLE);

    // Use simplified driver setup
    ethosu_set_power_mode_v2(drv, true);
}

CommandStream::~CommandStream() {
    ethosu_set_power_mode_v2(drv, false);
    ethosu_release_driver(drv);
}

int CommandStream::run(size_t repeat) {
    // Base pointer array
    uint64_t baseAddress[ETHOSU_DRIVER_BASEP_INDEXES];
    size_t baseAddressSize[ETHOSU_DRIVER_BASEP_INDEXES];
    for (size_t i = 0; i < ETHOSU_DRIVER_BASEP_INDEXES; i++) {
        baseAddress[i] = reinterpret_cast<uint64_t>(basePointers[i].data);
        baseAddressSize[i] = reinterpret_cast<size_t>(basePointers[i].size);
    }

    while (repeat-- > 0) {
        int error = ethosu_invoke_v3(drv, commandStream.data, commandStream.size, baseAddress, baseAddressSize, ETHOSU_DRIVER_BASEP_INDEXES);

        if (error != 0) {
            printf("Inference failed. error=%d\n", error);
            return 1;
        }

        // Wait for interrupt
        while (true) {
            uint16_t status;
            ethosu_get_status_mask(&drv->dev, &status);

            // Return if NPU raise error status
            if (status & 0xcc) {
                printf("Job failed with error. status=0x%08x\n", status);
                return 1;
            }

            // Break loop if job is no longer running
            if ((status & 1) == 0) {
                break;
            }

            // Sleep waiting on interrupt
            __WFI();
        }
    }

    return 0;
}

DataPointer &CommandStream::getCommandStream()
{
    return commandStream;
}

BasePointers &CommandStream::getBasePointers()
{
    return basePointers;
}

Pmu &CommandStream::getPmu()
{
    return pmu;
}

}; // namespace CommandStream
}; // namespace EthosU