////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Copyright (c) 2019-2021 Qualcomm Technologies, Inc.
// All Rights Reserved.
// Confidential and Proprietary - Qualcomm Technologies, Inc.
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

#pragma once

#include "is_dbg_util.h"
#include "eis_util.h"

class ConvolutionFilter
{
private:
    const uint32_t              _max_size;
    uint32_t                    _size;

protected:
    double*                     _coefs;

public: // members
    const uint32_t              _max_radius;
    uint32_t                    _radius;

public: // definitions
    enum class Mode { valid = 0, same, full };

public: // methods
    ConvolutionFilter(uint32_t radius, int32_t& status)
    :
        _max_size(2 * radius + 1),
        _size(_max_size),
        _coefs(NULL),
        _max_radius(radius),
        _radius(_max_radius)
    {
        IS_ASSERT(IS_RET_SUCCESS == status);

        _coefs = (double *)malloc((2 * _max_radius + 1) * sizeof(_coefs[0]));
        if (NULL == _coefs)
        {
            IS_ERR(IS_ERR_OUT_OF_MEMORY, "ConvolutionFilter failed");
            status = IS_RET_OUT_OF_MEMORY;
        }
    };

    virtual ~ConvolutionFilter()
    {
        FREE(_coefs);
    };

    inline uint32_t size() const { return _size; };
    inline uint32_t max_size() const { return _max_size; };

    void set_radius(uint32_t radius)
    {
        IS_ASSERT(radius <= _max_radius);
        if (radius <= _max_radius)
        {
            _radius = radius;
            _size = 2 * _radius + 1;
        }
        else
        {
            _radius = _max_radius;
            _size = 2 * _radius + 1;
        }
    }

    inline double& operator[](const uint32_t n) const
    {
        IS_ASSERT(n >= 0);
        IS_ASSERT(n < size());

        return _coefs[n];
    }

    template<class InputContainer, class OutputContainer>
    void filter(const InputContainer &input, OutputContainer &output, Mode mode = Mode::same)
    {
        filter_n(input, input.size(), output, mode);
    }

    template<class InputContainer, class OutputContainer>
    void filter_n(
        const InputContainer&   input,
        size_t                  input_size,
        OutputContainer&        output,
        Mode                    mode = Mode::same)
    {
        int nf = (int)size();
        int ni = (int)input_size;
        int n, first;

        switch (mode) {
        case Mode::valid:
            n = MAX(ni - MAX(0, nf - 1), 0);
            first = MIN(nf, ni) - 1;
            break;

        case Mode::full:
            n = nf + ni - 1;
            first = 0;
            break;

        case Mode::same:
        default:
            n = ni;
            first = _radius;
            break;
        }

        IS_ASSERT(n <= (int)output.maxSize());
        //        output.resize(n);

        for (int i = first; i < n + first; i++)
        {
            const int jmin = (i >= ni - 1) ? i - (ni - 1) : 0;
            const int jmax = (i < nf - 1) ? i : nf - 1;
            // need to come with something better like (T)0
            output[i - first] = 0;

            // missing values from start

            for (int j = jmin; j <= jmax; j++)
            {
                output[i - first] += (*this)[j] * input[i - j];
            }

            // in case of missing from the end
        }
    }

    uint32_t filter_by_idx(
        const RingBuffer<sample_data_t>&    input,
        uint32_t                            theta_idx,
        uint32_t                            start_idx,
        uint32_t                            end_idx,
        uint32_t                            output_buffer_max_sz,
        double*                             output,
        uint64_t*                           output_ts)
    {
        IS_ASSERT(end_idx >= start_idx); //TODO (etsairi): check if working when this is not true

        const uint32_t input_size = (uint32_t)input.size();

        // Number of elements to filter, filter the elements between [start_idx, end_idx], obviously buffer is cyclic
        const uint32_t tmp_cnt = (end_idx >= start_idx) ? (end_idx - start_idx + 1) :
                                                          (uint32_t)((int32_t)end_idx - (int32_t)start_idx + 1 + (int32_t)input_size);
        IS_ASSERT(tmp_cnt <= output_buffer_max_sz);

        const uint32_t cnt = MIN(output_buffer_max_sz, tmp_cnt);

        uint32_t input_idx = start_idx;
        uint32_t output_idx;

        for (output_idx = 0; output_idx < cnt; output_idx++)
        {
            uint32_t start_padding_len = 0;
            uint32_t main_len = _size;
            uint32_t end_padding_len = 0;

            if (input_idx < _radius)
            {
                start_padding_len = (_radius - input_idx);
                main_len -= start_padding_len;
            }

            if ((input_idx + _radius + 1) > input_size)
            {
                end_padding_len = (input_idx + _radius + 1) - input_size;
                main_len -= end_padding_len;
            }

            // This assert could happen and filter is ok with it but shouldn't happen in real cases too much
            IS_ASSERT((0 == start_padding_len) || (0 == end_padding_len)); // make sure that input_size is large enough

            IS_ASSERT(_size == (start_padding_len + main_len + end_padding_len)); // make sure that the entire filter is covered

            output[output_idx] = 0;

            // padding from the beginning
            for (uint32_t filter_idx = 0; filter_idx < start_padding_len; filter_idx++)
            {
                output[output_idx] += (*this)[filter_idx] * input[0].data[theta_idx];
            }

            // main filter
            for (uint32_t filter_idx = start_padding_len; filter_idx < start_padding_len + main_len; filter_idx++)
            {
                output[output_idx] += (*this)[filter_idx] * input[input_idx + filter_idx - _radius].data[theta_idx];
            }

            // padding from the end
            for (uint32_t filter_idx = start_padding_len + main_len ; filter_idx < _size; filter_idx++)
            {
                output[output_idx] += (*this)[filter_idx] * input[input_size - 1].data[theta_idx];
            }

            output_ts[output_idx] = input[input_idx].ts;

            input_idx = (input_idx + 1) % input_size;
        }

        IS_ASSERT(output_idx == cnt);
        return cnt;
    }
};

class SavitzkyGolay : public ConvolutionFilter
{
    static double gram_polynomial(const int i, const int m, const int k, const int s) {
        if (k > 0) {
            return (4. * k - 2.) / (k * (2. * m - k + 1.)) * (i * gram_polynomial(i, m, k - 1, s) + s * gram_polynomial(i, m, k - 1, s - 1)) - ((k - 1.) * (2. * m + k)) / (k * (2. * m - k + 1.)) * gram_polynomial(i, m, k - 2, s);
        }
        else {
            if (k == 0 && s == 0)
                return 1.;
            else
                return 0.;
        }
    }

    static double gram_polynomial_optimized(const int i, const int m, const int k, const int s) {
        if (k > 0) {
            double temp1 = (i == 0) ? 0 : i * gram_polynomial(i, m, k - 1, s);
            double temp2 = (s == 0) ? 0 : s * gram_polynomial(i, m, k - 1, s - 1);
            return (4. * k - 2.) / (k * (2. * m - k + 1.)) * (temp1 + temp2) - ((k - 1.) * (2. * m + k)) / (k * (2. * m - k + 1.)) * gram_polynomial(i, m, k - 2, s);
        }
        else {
            if (k == 0 && s == 0)
                return 1.;
            else
                return 0.;
        }
    }

    static double generalized_factorial(const int a, const int b) {
        double gf = 1.;

        for (int j = (a - b) + 1; j < a + 1; j++) {
            gf *= j;
        }
        return gf;
    }

    static double savitzky_golay_weights_optimized(const int i, const int m, const int n)
    {
        double w = 0;
        for (int k = 0; k < n + 1; k++)
        {
            w = w + (2 * k + 1) * (generalized_factorial(2 * m, k) / generalized_factorial(2 * m + k + 1, k + 1)) * gram_polynomial_optimized(i, m, k, 0) * gram_polynomial_optimized(0, m, k, 0);
        }
        return w;
    }

public:
    SavitzkyGolay(uint32_t radius, uint32_t polynomial_order, int32_t& status) : ConvolutionFilter(radius, status)
    {
        IS_ASSERT(IS_RET_SUCCESS == status);

        if (NULL != _coefs)
        {
            for (int i = 0; i < (int)radius + 1; i++)
            {
                double value = savitzky_golay_weights_optimized(i - (int)radius, radius, polynomial_order);
                _coefs[i] = value;
                _coefs[2 * radius - i] = value;
            }
        }
    }

    virtual ~SavitzkyGolay() {};
};

class Blackman : public ConvolutionFilter
{
private:
    double calculate_blackman_n_weight(uint32_t n) const
    {
        double w;
        if (n <= _radius)
        {
            const double length = size();
            w = 0.42 - 0.5 * std::cos((2 * IS_PI * n) / (length - 1)) + 0.08 * std::cos((4 * IS_PI * n) / (length - 1));
        }
        else
        {
            w = 0;
        }
        return w;
    }

    // TODO (asimchi): optimize
    void calculate_blackman_weights()
    {
        if (NULL != _coefs)
        {
            for (uint32_t i = 0; i < max_size(); i++)
            {
                _coefs[i] = 0;
            }

            if (1 == _radius)
            {
                _coefs[0] = 0;
                _coefs[1] = 1;
                _coefs[2] = 0;
            }
            else
            {
                double weight_sum = 0;

                // calculate weights
                for (uint32_t i = 0; i < _radius; i++)
                {
                    _coefs[i] = calculate_blackman_n_weight(i);
                    weight_sum += _coefs[i] * 2;
                }
                _coefs[_radius] = calculate_blackman_n_weight(_radius);
                weight_sum += _coefs[_radius];

                // normalize
                for (uint32_t i = 0; i < _radius; i++)
                {
                    _coefs[i] /= weight_sum;
                    _coefs[2 * _radius - i] = _coefs[i];
                }
                _coefs[_radius] /= weight_sum;
            }
        }
    }

public:
    Blackman(uint32_t radius, int32_t& status) : ConvolutionFilter(radius, status)
    {
        calculate_blackman_weights();
    }

    void update_blackman_weights(uint32_t radius)
    {

        set_radius(radius);
        calculate_blackman_weights();
    }

    virtual ~Blackman() {};
};
