/*
 * Copyright (c) 2024 EdgeImpulse Inc.
 *
 * Generated by Edge Impulse and licensed under the applicable Edge Impulse
 * Terms of Service. Community and Professional Terms of Service
 * (https://edgeimpulse.com/legal/terms-of-service) or Enterprise Terms of
 * Service (https://edgeimpulse.com/legal/enterprise-terms-of-service),
 * according to your product plan subscription (the “License”).
 *
 * This software, documentation and other associated files (collectively referred
 * to as the “Software”) is a single SDK variation generated by the Edge Impulse
 * platform and requires an active paid Edge Impulse subscription to use this
 * Software for any purpose.
 *
 * You may NOT use this Software unless you have an active Edge Impulse subscription
 * that meets the eligibility requirements for the applicable License, subject to
 * your full and continued compliance with the terms and conditions of the License,
 * including without limitation any usage restrictions under the applicable License.
 *
 * If you do not have an active Edge Impulse product plan subscription, or if use
 * of this Software exceeds the usage limitations of your Edge Impulse product plan
 * subscription, you are not permitted to use this Software and must immediately
 * delete and erase all copies of this Software within your control or possession.
 * Edge Impulse reserves all rights and remedies available to enforce its rights.
 *
 * Unless required by applicable law or agreed to in writing, the Software 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, disclaimers and limitations under the License.
 */
#ifndef __EI_FLATTEN__H__
#define __EI_FLATTEN__H__

#include "edge-impulse-sdk/dsp/ei_vector.h"
#include "edge-impulse-sdk/dsp/returntypes.hpp"
#include "edge-impulse-sdk/dsp/ei_dsp_handle.h"
#include "model-parameters/model_metadata.h"
#include "edge-impulse-sdk/dsp/numpy.hpp"
#include "edge-impulse-sdk/dsp/config.hpp"

class flatten_class : public DspHandle {
public:
    int print() override {
        ei_printf("means: ");
        for(int axis = 0; (size_t)axis < this->means.size(); axis++) {
        ei_printf("axis: %i\n", axis);
            for (size_t i = 0; i < this->means.size(); i++) {
                ei_printf("%f ", this->means[axis][i]);
            }
        }
        ei_printf("\n");
        return ei::EIDSP_OK;
    }

    int extract(
        ei::signal_t *signal,
        ei::matrix_t *output_matrix,
        void *config_ptr,
        const float frequency,
        ei_impulse_result_t *result) override
    {
        using namespace ei;

        ei_dsp_config_flatten_t config = *((ei_dsp_config_flatten_t*)config_ptr);

        uint32_t expected_matrix_size = 0;
        if (config.average) expected_matrix_size += config.axes;
        if (config.minimum) expected_matrix_size += config.axes;
        if (config.maximum) expected_matrix_size += config.axes;
        if (config.rms) expected_matrix_size += config.axes;
        if (config.stdev) expected_matrix_size += config.axes;
        if (config.skewness) expected_matrix_size += config.axes;
        if (config.kurtosis) expected_matrix_size += config.axes;
        if (config.moving_avg_num_windows) expected_matrix_size += config.axes;

        if (output_matrix->rows * output_matrix->cols != expected_matrix_size) {
            EIDSP_ERR(EIDSP_MATRIX_SIZE_MISMATCH);
        }

        int ret;

        // input matrix from the raw signal
        matrix_t input_matrix(signal->total_length / config.axes, config.axes);
        if (!input_matrix.buffer) {
            EIDSP_ERR(EIDSP_OUT_OF_MEM);
        }
        signal->get_data(0, signal->total_length, input_matrix.buffer);

        // scale the signal
        ret = numpy::scale(&input_matrix, config.scale_axes);
        if (ret != EIDSP_OK) {
            ei_printf("ERR: Failed to scale signal (%d)\n", ret);
            EIDSP_ERR(ret);
        }

        // transpose the matrix so we have one row per axis
        numpy::transpose_in_place(&input_matrix);

        size_t out_matrix_ix = 0;

        for (size_t row = 0; row < input_matrix.rows; row++) {
            matrix_t row_matrix(1, input_matrix.cols, input_matrix.buffer + (row * input_matrix.cols));

            float mean; // to use with moving average

            if (config.average || config.moving_avg_num_windows) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::mean(&row_matrix, &out_matrix);
                mean = out_matrix.buffer[0];
                if (config.average) {
                    output_matrix->buffer[out_matrix_ix++] = mean;
                }
            }

            if (config.minimum) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::min(&row_matrix, &out_matrix);
                output_matrix->buffer[out_matrix_ix++] = out_matrix.buffer[0];
            }

            if (config.maximum) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::max(&row_matrix, &out_matrix);
                output_matrix->buffer[out_matrix_ix++] = out_matrix.buffer[0];
            }

            if (config.rms) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::rms(&row_matrix, &out_matrix);
                output_matrix->buffer[out_matrix_ix++] = out_matrix.buffer[0];
            }

            if (config.stdev) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::stdev(&row_matrix, &out_matrix);
                output_matrix->buffer[out_matrix_ix++] = out_matrix.buffer[0];
            }

            if (config.skewness) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::skew(&row_matrix, &out_matrix);
                output_matrix->buffer[out_matrix_ix++] = out_matrix.buffer[0];
            }

            if (config.kurtosis) {
                float fbuffer;
                matrix_t out_matrix(1, 1, &fbuffer);
                numpy::kurtosis(&row_matrix, &out_matrix);
                output_matrix->buffer[out_matrix_ix++] = out_matrix.buffer[0];
            }

            if (config.moving_avg_num_windows) {
                push_mean(row, mean);
                output_matrix->buffer[out_matrix_ix++] = numpy::mean(means[row].data(), means[row].size());
            }
        }

        // flatten again
        output_matrix->cols = output_matrix->rows * output_matrix->cols;
        output_matrix->rows = 1;

        return EIDSP_OK;
    }

    static DspHandle* create(void* config, float _sampling_frequency);

    void* operator new(size_t size) {
        // Custom memory allocation logic here
        return ei_malloc(size);
    }

    void operator delete(void* ptr) {
        // Custom memory deallocation logic here
        ei_free(ptr);
    }

private:
    ei_vector<ei_vector<float>> means;
    ei_vector<size_t> head_indexes;
    size_t moving_avg_num_windows;

    flatten_class(int moving_avg_num_windows, int axes_count) : means(axes_count), head_indexes(axes_count, 0) {
        this->moving_avg_num_windows = moving_avg_num_windows;
    }

    void push_mean(int axis, float mean) {
        auto& head = head_indexes[axis];
        if (head_indexes[axis] >= means[axis].size()) {
            means[axis].push_back(mean);
        } else {
            means[axis][head] = mean;
        }
        head = head + 1;
        // This is a lot cheaper than mod (%)
        if (head >= moving_avg_num_windows) {
            head = 0;
        }
    }
};

DspHandle* flatten_class::create(void* config_in, float _sampling_frequency) { // NOLINT def in header is OK at EI
    auto config = reinterpret_cast<ei_dsp_config_flatten_t*>(config_in);
    return new flatten_class(config->moving_avg_num_windows, config->axes);
};

#endif  //!__EI_FLATTEN__H__