XYParser/XYParser/XYEegParserCommon.h

#pragma once

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>

namespace xyparser {

template <std::size_t ChannelCount>
struct XYEegSample {
    std::array<double, ChannelCount> channel_values_uv{};
    std::uint8_t trigger_type = 0;
    std::uint8_t trigger_index = 0;
};

template <std::size_t ChannelCount>
struct XYEegFrame {
    static constexpr std::size_t kSamplesPerFrame = 5;

    std::uint8_t header = 0;
    std::uint32_t index = 0;
    std::uint16_t payload_length = 0;
    std::uint8_t battery = 0;
    std::uint8_t channel_count = 0;
    std::int16_t elevation_angle = 0;
    std::int16_t rolling_angle = 0;
    std::int16_t yaw_angle = 0;
    std::int16_t ecg = 0;
    std::int16_t blood_oxygen = 0;
    std::array<std::uint8_t, 6> reserved{};
    std::uint8_t impedance_enabled = 0;
    std::uint8_t current_gain = 0;
    std::uint16_t current_sample_rate_hz = 0;
    std::uint8_t cap_type = 0;
    std::uint8_t gnd_detached = 0;
    std::array<XYEegSample<ChannelCount>, kSamplesPerFrame> samples{};
    std::uint8_t crc = 0;
    std::array<std::uint8_t, 2> tails{};
};

template <std::size_t ChannelCount>
class XYEegTcpParserCommon {
public:
    using Frame = XYEegFrame<ChannelCount>;

    static constexpr std::uint8_t kFrameHeader = 0xAA;
    static constexpr std::uint8_t kFrameTail = 0x55;
    static constexpr std::size_t kSamplesPerFrame = 5;
    static constexpr std::size_t kFrameHeaderLen = 1;
    static constexpr std::size_t kFrameTailLen = 2;
    static constexpr std::size_t kFrameTagLen = 25;

    void SetAdcParams(double vref, double gain) noexcept
    {
        if (vref > 0.0) {
            vref_ = vref;
        }
        if (gain > 0.0) {
            gain_ = gain;
        }
    }

    void SetBypassChecksum(bool bypass) noexcept
    {
        bypass_checksum_ = bypass;
    }

    bool BypassChecksum() const noexcept
    {
        return bypass_checksum_;
    }

    void Reset()
    {
        buffer_.clear();
        last_error_.clear();
    }

    const std::string &LastError() const noexcept
    {
        return last_error_;
    }

    std::vector<Frame> Feed(const std::vector<std::uint8_t> &bytes)
    {
        return Feed(bytes.data(), bytes.size());
    }

    std::vector<Frame> Feed(const std::uint8_t *data, std::size_t size)
    {
        std::vector<Frame> frames;
        if (data == nullptr || size == 0) {
            return frames;
        }

        buffer_.insert(buffer_.end(), data, data + size);

        while (true) {
            const auto header_it = std::find(buffer_.begin(), buffer_.end(), kFrameHeader);
            if (header_it == buffer_.end()) {
                buffer_.clear();
                break;
            }

            const std::size_t header_index = static_cast<std::size_t>(std::distance(buffer_.begin(), header_it));
            if (buffer_.size() < header_index + 7) {
                if (header_index > 0) {
                    buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(header_index));
                }
                break;
            }

            const std::uint16_t payload_length = ReadBE16(buffer_, header_index + 5);
            const std::size_t total_frame_len = payload_length + kFrameTagLen + kFrameTailLen + kFrameHeaderLen;
            if (buffer_.size() - header_index < total_frame_len) {
                if (header_index > 0) {
                    buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(header_index));
                }
                break;
            }

            const std::size_t tail_index = header_index + payload_length + kFrameTagLen + kFrameHeaderLen;
            if (tail_index + 1 >= buffer_.size() || buffer_[tail_index] != kFrameTail || buffer_[tail_index + 1] != kFrameTail) {
                buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(header_index + 1));
                continue;
            }

            std::vector<std::uint8_t> frame_bytes(
                buffer_.begin() + static_cast<std::ptrdiff_t>(header_index),
                buffer_.begin() + static_cast<std::ptrdiff_t>(tail_index + 2));
            buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<std::ptrdiff_t>(tail_index + 2));

            Frame frame;
            if (TryParseFrame(frame_bytes, frame)) {
                frames.push_back(frame);
            }
        }

        return frames;
    }

private:
    static std::uint16_t ReadBE16(const std::vector<std::uint8_t> &bytes, std::size_t offset) noexcept
    {
        return static_cast<std::uint16_t>((static_cast<std::uint16_t>(bytes[offset]) << 8) |
                                          static_cast<std::uint16_t>(bytes[offset + 1]));
    }

    static std::uint16_t ReadLE16(const std::vector<std::uint8_t> &bytes, std::size_t offset) noexcept
    {
        return static_cast<std::uint16_t>(static_cast<std::uint16_t>(bytes[offset]) |
                                          (static_cast<std::uint16_t>(bytes[offset + 1]) << 8));
    }

    static std::uint32_t ReadLE32(const std::vector<std::uint8_t> &bytes, std::size_t offset) noexcept
    {
        return static_cast<std::uint32_t>(bytes[offset]) |
               (static_cast<std::uint32_t>(bytes[offset + 1]) << 8) |
               (static_cast<std::uint32_t>(bytes[offset + 2]) << 16) |
               (static_cast<std::uint32_t>(bytes[offset + 3]) << 24);
    }

    static std::uint16_t DecodeSampleRateHz(std::uint8_t sample_rate_code) noexcept
    {
        switch (sample_rate_code) {
        case 0:
            return 250;
        case 1:
            return 500;
        case 2:
            return 1000;
        case 3:
            return 2000;
        default:
            return 0;
        }
    }

    double ConvertAdcToUv(const std::uint8_t raw0, const std::uint8_t raw1, const std::uint8_t raw2) const noexcept
    {
        double value = static_cast<double>(raw0) * 65536.0 +
                       static_cast<double>(raw1) * 256.0 +
                       static_cast<double>(raw2);
        const double factor = vref_ / gain_ / 8388608.0 * 1000000.0;
        if (value < 8388608.0) {
            return value * factor;
        }
        return -(16777216.0 - value) * factor;
    }

    bool CheckChecksum(const std::vector<std::uint8_t> &frame_bytes) const noexcept
    {
        if (frame_bytes.size() < 4) {
            return false;
        }

        const std::size_t crc_index = frame_bytes.size() - 3;
        std::uint8_t checksum = 0;
        for (std::size_t i = 0; i < crc_index; ++i) {
            checksum = static_cast<std::uint8_t>(checksum + frame_bytes[i]);
        }
        return checksum == frame_bytes[crc_index];
    }

    bool TryParseFrame(const std::vector<std::uint8_t> &frame_bytes, Frame &frame)
    {
        if (!bypass_checksum_ && !CheckChecksum(frame_bytes)) {
            last_error_ = "checksum failed";
            return false;
        }

        const std::size_t min_frame_len = kFrameHeaderLen + kFrameTagLen + kFrameTailLen;
        if (frame_bytes.size() < min_frame_len) {
            last_error_ = "frame too short";
            return false;
        }

        const std::size_t channel_count_offset = kFrameHeaderLen + 4 + 2 + 1;
        const std::uint8_t channel_count = frame_bytes[channel_count_offset];
        if (channel_count != ChannelCount) {
            last_error_ = "unexpected channel count";
            return false;
        }

        const std::size_t sample_bytes = ChannelCount * 3 + 2;
        const std::size_t expected_frame_len = kFrameHeaderLen + kFrameTagLen +
                                               kSamplesPerFrame * sample_bytes + kFrameTailLen;
        if (frame_bytes.size() != expected_frame_len) {
            last_error_ = "frame length mismatch";
            return false;
        }

        std::size_t offset = 0;
        frame.header = frame_bytes[offset++];
        frame.index = ReadLE32(frame_bytes, offset);
        offset += 4;
        frame.payload_length = ReadBE16(frame_bytes, offset);
        offset += 2;
        frame.battery = frame_bytes[offset++];
        frame.channel_count = frame_bytes[offset++];
        frame.elevation_angle = static_cast<std::int16_t>(ReadLE16(frame_bytes, offset));
        offset += 2;
        frame.rolling_angle = static_cast<std::int16_t>(ReadLE16(frame_bytes, offset));
        offset += 2;
        frame.yaw_angle = static_cast<std::int16_t>(ReadLE16(frame_bytes, offset));
        offset += 2;
        frame.ecg = static_cast<std::int16_t>(ReadLE16(frame_bytes, offset));
        offset += 2;
        frame.blood_oxygen = static_cast<std::int16_t>(ReadLE16(frame_bytes, offset));
        offset += 2;

        for (std::size_t i = 0; i < frame.reserved.size(); ++i) {
            frame.reserved[i] = frame_bytes[offset++];
        }
        frame.impedance_enabled = static_cast<std::uint8_t>(frame.reserved[0] & 0x01U);
        frame.current_gain = frame.reserved[1];
        frame.current_sample_rate_hz = DecodeSampleRateHz(frame.reserved[2]);
        frame.cap_type = frame.reserved[3];
        frame.gnd_detached = static_cast<std::uint8_t>(frame.reserved[4] != 0 ? 1 : 0);

        for (auto &sample : frame.samples) {
            for (std::size_t channel = 0; channel < ChannelCount; ++channel) {
                sample.channel_values_uv[channel] = ConvertAdcToUv(
                    frame_bytes[offset],
                    frame_bytes[offset + 1],
                    frame_bytes[offset + 2]);
                offset += 3;
            }
            sample.trigger_type = frame_bytes[offset++];
            sample.trigger_index = frame_bytes[offset++];
        }

        frame.crc = frame_bytes[offset++];
        frame.tails[0] = frame_bytes[offset++];
        frame.tails[1] = frame_bytes[offset++];

        last_error_.clear();
        return true;
    }

private:
    std::vector<std::uint8_t> buffer_;
    std::string last_error_;
    double vref_ = 4.5;
    double gain_ = 6.0;
    bool bypass_checksum_ = true;
};

} // namespace xyparser