#include "pch.h"
#include "XYWelchProcessor.h"

#include "third_party/kissfft/kiss_fftr.h"
#include "third_party/kissfft/window_functions.h"

#include <cmath>
#include <cstring>
#include <limits>
#include <vector>

namespace {

constexpr int kDefaultSampleRate = 250;
constexpr int k8ChLeadCount = 8;
constexpr int k64ChLeadCount = 64;

struct PsdBandDefinition {
    double low_hz;
    double high_hz;
};

constexpr std::array<XYParserLeadChannelNumber, k8ChLeadCount> k8ChLeadMap = {
    LeadChannel_PO5, LeadChannel_POZ, LeadChannel_PO6, LeadChannel_PO7,
    LeadChannel_O1, LeadChannel_OZ, LeadChannel_O2, LeadChannel_PO8};

constexpr std::array<XYParserLeadChannelNumber, k64ChLeadCount> k64ChLeadMap = {
    LeadChannel_FP1, LeadChannel_FP2, LeadChannel_PO6, LeadChannel_POZ,
    LeadChannel_F3, LeadChannel_F4, LeadChannel_FPZ, LeadChannel_AF4,
    LeadChannel_FC3, LeadChannel_PO8, LeadChannel_CP2, LeadChannel_CP1,
    LeadChannel_FCZ, LeadChannel_PO5, LeadChannel_FC2, LeadChannel_FC1,
    LeadChannel_C3, LeadChannel_C4, LeadChannel_FC4, LeadChannel_CP4,
    LeadChannel_P3, LeadChannel_P4, LeadChannel_F5, LeadChannel_C5,
    LeadChannel_F6, LeadChannel_PO4, LeadChannel_CP6, LeadChannel_CP5,
    LeadChannel_PO3, LeadChannel_CP3, LeadChannel_FC6, LeadChannel_FC5,
    LeadChannel_CB1, LeadChannel_CB2, LeadChannel_P5, LeadChannel_AF7,
    LeadChannel_A1, LeadChannel_T7, LeadChannel_FT7, LeadChannel_TP7,
    LeadChannel_FT8, LeadChannel_AF8, LeadChannel_F8, LeadChannel_F7,
    LeadChannel_P6, LeadChannel_C6, LeadChannel_O2, LeadChannel_O1,
    LeadChannel_T8, LeadChannel_P7, LeadChannel_CZ, LeadChannel_PZ,
    LeadChannel_P8, LeadChannel_FZ, LeadChannel_OZ, LeadChannel_PO7,
    LeadChannel_TP8, LeadChannel_AF3, LeadChannel_C2, LeadChannel_C1,
    LeadChannel_P2, LeadChannel_P1, LeadChannel_F2, LeadChannel_F1};

constexpr std::array<PsdBandDefinition, XYPARSER_WELCH_BAND_COUNT> kWelchBands = {{
    {0.5, 4.0},
    {4.0, 8.0},
    {8.0, 13.0},
    {13.0, 30.0},
    {30.0, 50.0},
}};

} // namespace

XYWelchProcessor::XYWelchProcessor(std::uint8_t channel_count)
    : channel_count_(channel_count)
{
}

void XYWelchProcessor::SetChannelCount(std::uint8_t channel_count)
{
    channel_count_ = channel_count;
    Reset();
}

void XYWelchProcessor::SetSampleRate(int sample_rate)
{
    if (sample_rate <= 0) {
        return;
    }

    if (sample_rate_ != sample_rate) {
        sample_rate_ = sample_rate;
        Reset();
    }
}

void XYWelchProcessor::SetEnabled(bool enabled)
{
    if (enabled_ == enabled) {
        return;
    }

    enabled_ = enabled;
    Reset();
}

void XYWelchProcessor::Reset()
{
    cached_samples_.clear();
    pending_results_.clear();
}

void XYWelchProcessor::PushFrame(const XYParserFrameSummary& frame)
{
    if (!enabled_ || frame.channel_count != channel_count_) {
        return;
    }

    for (std::size_t sample_index = 0; sample_index < frame.sample_count; ++sample_index) {
        Sample sample{};
        for (std::size_t channel_index = 0; channel_index < frame.channel_count; ++channel_index) {
            sample[channel_index] = frame.channel_values_uv[sample_index][channel_index];
        }
        cached_samples_.push_back(sample);
    }

    ProcessPendingWindows();
}

void XYWelchProcessor::PushSample(const std::array<double, XYPARSER_MAX_CHANNELS>& sample)
{
    if (!enabled_ || channel_count_ == 0) {
        return;
    }

    cached_samples_.push_back(sample);
    ProcessPendingWindows();
}

int XYWelchProcessor::Read(XYParserWelchSummary* out_summaries, int max_summaries)
{
    if (out_summaries == nullptr || max_summaries <= 0) {
        return 0;
    }

    const int available_count = static_cast<int>(pending_results_.size());
    const int read_count = available_count < max_summaries ? available_count : max_summaries;
    for (int i = 0; i < read_count; ++i) {
        out_summaries[i] = pending_results_.front();
        pending_results_.pop_front();
    }
    return read_count;
}

void XYWelchProcessor::ProcessPendingWindows()
{
    const int effective_sample_rate = sample_rate_ > 0 ? sample_rate_ : 1;
    const std::size_t window_sample_count = static_cast<std::size_t>(effective_sample_rate);
    while (cached_samples_.size() >= window_sample_count) {
        pending_results_.push_back(BuildSummaryFromWindow());
        cached_samples_.erase(cached_samples_.begin(), cached_samples_.begin() + static_cast<std::ptrdiff_t>(window_sample_count));
    }
}

XYParserWelchSummary XYWelchProcessor::BuildSummaryFromWindow() const
{
    XYParserWelchSummary summary{};
    summary.channel_count = channel_count_;
    summary.sample_rate = static_cast<std::uint32_t>(sample_rate_ > 0 ? sample_rate_ : kDefaultSampleRate);
    summary.window_sample_count = summary.sample_rate;

    const int sample_count = static_cast<int>(summary.window_sample_count);
    if (channel_count_ == 0 || sample_count < 2) {
        return summary;
    }

    int n_per_seg = sample_count;
    if (n_per_seg % 2 != 0) {
        --n_per_seg;
    }
    if (n_per_seg < 2) {
        return summary;
    }

    const int n_overlap = n_per_seg / 2;
    const int n_step = (n_per_seg - n_overlap) > 0 ? (n_per_seg - n_overlap) : 1;
    const int n_freq_bin_count = n_per_seg / 2 + 1;

    std::vector<double> window(static_cast<std::size_t>(n_per_seg), 0.0);
    hanning_function(n_per_seg, window.data());

    double window_power = 0.0;
    for (int i = 0; i < n_per_seg; ++i) {
        window_power += window[static_cast<std::size_t>(i)] * window[static_cast<std::size_t>(i)];
    }
    if (window_power <= std::numeric_limits<double>::epsilon()) {
        return summary;
    }

    std::vector<std::vector<double>> psd_accum(static_cast<std::size_t>(channel_count_),
                                               std::vector<double>(static_cast<std::size_t>(n_freq_bin_count), 0.0));
    std::vector<kiss_fft_scalar> segment(static_cast<std::size_t>(n_per_seg), 0.0);
    std::vector<kiss_fft_cpx> fft_output(static_cast<std::size_t>(n_freq_bin_count));

    kiss_fftr_cfg cfg = kiss_fftr_alloc(n_per_seg, 0, nullptr, nullptr);
    if (cfg == nullptr) {
        return summary;
    }

    int segment_count = 0;
    for (int start = 0; start + n_per_seg <= sample_count; start += n_step) {
        ++segment_count;
        for (std::size_t channel_index = 0; channel_index < channel_count_; ++channel_index) {
            double mean = 0.0;
            for (int i = 0; i < n_per_seg; ++i) {
                mean += cached_samples_[static_cast<std::size_t>(start + i)][channel_index];
            }
            mean /= static_cast<double>(n_per_seg);

            for (int i = 0; i < n_per_seg; ++i) {
                segment[static_cast<std::size_t>(i)] =
                    (cached_samples_[static_cast<std::size_t>(start + i)][channel_index] - mean) *
                    window[static_cast<std::size_t>(i)];
            }

            kiss_fftr(cfg, segment.data(), fft_output.data());
            for (int bin = 0; bin < n_freq_bin_count; ++bin) {
                const double real = fft_output[static_cast<std::size_t>(bin)].r;
                const double imag = fft_output[static_cast<std::size_t>(bin)].i;
                double psd = (real * real + imag * imag) /
                             (static_cast<double>(summary.sample_rate) * window_power);
                if (bin != 0 && bin != (n_per_seg / 2)) {
                    psd *= 2.0;
                }
                psd_accum[channel_index][static_cast<std::size_t>(bin)] += psd;
            }
        }
    }

    kiss_fftr_free(cfg);
    if (segment_count <= 0) {
        return summary;
    }

    std::vector<double> all_frequencies(static_cast<std::size_t>(n_freq_bin_count), 0.0);
    const double freq_resolution = static_cast<double>(summary.sample_rate) / static_cast<double>(n_per_seg);
    for (int bin = 0; bin < n_freq_bin_count; ++bin) {
        all_frequencies[static_cast<std::size_t>(bin)] = freq_resolution * static_cast<double>(bin);
    }

    for (int bin = 0; bin < n_freq_bin_count; ++bin) {
        const double frequency = all_frequencies[static_cast<std::size_t>(bin)];
        if (frequency >= 0.5 && frequency <= 50.0 &&
            summary.frequency_count < XYPARSER_WELCH_MAX_FREQUENCY_COUNT) {
            summary.frequencies[summary.frequency_count] = frequency;
            ++summary.frequency_count;
        }
    }

    const auto fill_channel_result = [&](std::size_t raw_channel_index, XYParserLeadChannelNumber lead) {
        const std::size_t lead_index = static_cast<std::size_t>(lead);

        std::uint32_t selected_index = 0;
        for (int bin = 0; bin < n_freq_bin_count; ++bin) {
            const double frequency = all_frequencies[static_cast<std::size_t>(bin)];
            if (frequency >= 0.5 && frequency <= 50.0 &&
                selected_index < summary.frequency_count) {
                summary.psd_values[lead_index][selected_index] =
                    psd_accum[raw_channel_index][static_cast<std::size_t>(bin)] / static_cast<double>(segment_count);
                ++selected_index;
            }
        }

        for (std::size_t band_index = 0; band_index < kWelchBands.size(); ++band_index) {
            int band_bin_count = 0;
            double band_sum = 0.0;
            for (int bin = 0; bin < n_freq_bin_count; ++bin) {
                const double frequency = all_frequencies[static_cast<std::size_t>(bin)];
                if (frequency < kWelchBands[band_index].low_hz || frequency > kWelchBands[band_index].high_hz) {
                    continue;
                }
                ++band_bin_count;
                band_sum += psd_accum[raw_channel_index][static_cast<std::size_t>(bin)] / static_cast<double>(segment_count);
            }
            if (band_bin_count > 0) {
                summary.band_values[band_index][lead_index] = band_sum / static_cast<double>(band_bin_count);
            }
        }
    };

    if (channel_count_ == 8) {
        for (std::size_t i = 0; i < k8ChLeadMap.size(); ++i) {
            fill_channel_result(i, k8ChLeadMap[i]);
        }
    } else {
        const std::size_t fill_count = channel_count_ < k64ChLeadMap.size() ? channel_count_ : k64ChLeadMap.size();
        for (std::size_t i = 0; i < fill_count; ++i) {
            fill_channel_result(i, k64ChLeadMap[i]);
        }
    }

    summary.ok = 1;
    return summary;
}
