#include "pch.h"
#include "XYImpedanceProcessor.h"

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

#include <algorithm>
#include <array>
#include <cmath>
#include <vector>

namespace {

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

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};

} // namespace

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

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

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

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

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

    enabled_ = enabled;
    Reset();
}

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

void XYImpedanceProcessor::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();
}

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

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

void XYImpedanceProcessor::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));
    }
}

XYParserImpedanceSummary XYImpedanceProcessor::BuildSummaryFromWindow() const
{
    XYParserImpedanceSummary 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 auto fill_channel = [&](std::size_t channel_index, XYParserLeadChannelNumber lead) {
        summary.impedance_values[static_cast<std::size_t>(lead)] = EstimateImpedanceValue(channel_index);
    };

    if (channel_count_ == 8) {
        for (std::size_t i = 0; i < k8ChLeadMap.size(); ++i) {
            fill_channel(i, k8ChLeadMap[i]);
        }
        return summary;
    }

    for (std::size_t i = 0; i < std::min<std::size_t>(channel_count_, k64ChLeadMap.size()); ++i) {
        fill_channel(i, k64ChLeadMap[i]);
    }
    return summary;
}

std::uint16_t XYImpedanceProcessor::EstimateImpedanceValue(std::size_t channel_index) const
{
    const int effective_sample_rate = sample_rate_ > 0 ? sample_rate_ : 1;
    const int cached_sample_count = static_cast<int>(cached_samples_.size());
    const int raw_sample_count = cached_sample_count < effective_sample_rate ? cached_sample_count : effective_sample_rate;
    if (raw_sample_count < 2 || channel_index >= channel_count_) {
        return 0;
    }

    const int fft_size = (raw_sample_count % 2 == 0) ? raw_sample_count : (raw_sample_count - 1);
    if (fft_size < 2) {
        return 0;
    }

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

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

    double window_sum = 0.0;
    for (int i = 0; i < fft_size; ++i) {
        time_data[static_cast<std::size_t>(i)] = cached_samples_[static_cast<std::size_t>(i)][channel_index] * window[static_cast<std::size_t>(i)];
        window_sum += window[static_cast<std::size_t>(i)];
    }
    if (window_sum <= 0.0) {
        kiss_fftr_free(cfg);
        return 0;
    }

    std::vector<kiss_fft_cpx> freq_data(static_cast<std::size_t>(fft_size / 2 + 1));
    kiss_fftr(cfg, time_data.data(), freq_data.data());
    kiss_fftr_free(cfg);

    const double sample_rate = sample_rate_ > 0 ? static_cast<double>(sample_rate_) : static_cast<double>(kDefaultSampleRate);
    const double freq_resolution = sample_rate / static_cast<double>(fft_size);
    constexpr double kLineNoiseFreq = 50.0;
    constexpr double kLineNoiseHarmonicFreq = 100.0;
    constexpr double kRejectHalfBand = 5.0;
    constexpr double kImpedanceScale = 0.0008;

    double max_magnitude = 0.0;
    for (std::size_t k = 1; k < freq_data.size(); ++k) {
        const double frequency = static_cast<double>(k) * freq_resolution;
        if (std::abs(frequency - kLineNoiseFreq) <= kRejectHalfBand ||
            std::abs(frequency - kLineNoiseHarmonicFreq) <= kRejectHalfBand) {
            continue;
        }

        const double real = freq_data[k].r;
        const double imag = freq_data[k].i;
        double magnitude = std::sqrt(real * real + imag * imag);
        if (k == freq_data.size() - 1 && (fft_size % 2 == 0)) {
            magnitude /= window_sum;
        } else {
            magnitude *= 2.0 / window_sum;
        }

        if (magnitude > max_magnitude) {
            max_magnitude = magnitude;
        }
    }

    const double scaled_value = std::round(max_magnitude) * kImpedanceScale;
    const long rounded_value = std::lround(scaled_value);
    const long clamped_value = rounded_value < 0 ? 0 : (rounded_value > 65535 ? 65535 : rounded_value);
    return static_cast<std::uint16_t>(clamped_value);
}
