/*
 * Copyright (c) 2025 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_EEG_HPP
#define EI_EEG_HPP

#include "edge-impulse-sdk/dsp/ei_dsp_handle.h"
#include "edge-impulse-enterprise/eeg/eeg.hpp"

// Include the generated header
#include "model-parameters/coeffs.hpp"

// Wrapper to use ei_malloc for allocation
class eeg_wrap : public ei::eeg_processor {
public:
    void* operator new(size_t size) {
        return ei_malloc(size);
    }
    void operator delete(void* ptr) {
        ei_free(ptr);
    }
    using ei::eeg_processor::eeg_processor;
};

class eeg_class : public DspHandle {
public:
    int print() override {
        ei_printf("EEG processor: channels=%zu\n", eeg.num_channels);
        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;
        // config_ptr is expected to be a struct with EEG config
        ei_dsp_config_eeg_t& config = *((ei_dsp_config_eeg_t*)config_ptr);
        size_t floats_per_inc = eeg.win_inc_samples * eeg.num_channels;
        assert(signal->total_length >= floats_per_inc);
        int out_idx = 0;
        for (size_t i = 0; i <= signal->total_length - floats_per_inc; i += floats_per_inc) {
            matrix_t temp(eeg.win_inc_samples, eeg.num_channels);
            signal->get_data(i, floats_per_inc, temp.buffer);
            // Create a temp psd output matrix structure, but just use the output matrix buffer
            matrix_t psd_mat(eeg.num_channels, eeg.epoch_samples / 2, output_matrix->buffer + out_idx);
            eeg.stream(&temp, psd_mat);
            out_idx += psd_mat.rows * psd_mat.cols;
        }
        return EIDSP_OK;
    }

    eeg_class(ei_dsp_config_eeg_t* config, float frequency)
        : eeg(frequency,
              config->axes,
              config->scale_axes,
              config->motion_sensitivity,
              (const float*)ei::eeg::filter_coeff,
              (const float*)ei::eeg::filter_zi,
              ei::eeg::num_sections,
              (const float*)ei::eeg::tapers,
              ei::eeg::num_tapers,
              (size_t)(frequency * config->epoch_length),
              true)
    {
        // Any additional config setup can go here
    }

    ~eeg_class() {}

    // Boilerplate below here
    static DspHandle* create(void* config, float frequency);

    void* operator new(size_t size) {
        return ei_malloc(size);
    }

    void operator delete(void* ptr) {
        ei_free(ptr);
    }
private:
    eeg_wrap eeg;
};

DspHandle* eeg_class::create(void* config_in, float frequency) {
    auto config = reinterpret_cast<ei_dsp_config_eeg_t*>(config_in);
    return new eeg_class(config, frequency);
}

#endif // EI_EEG_HPP