Обработка данных электроэнцефалографии (ЭЭГ)

# Чтение EDF (European Data Format) файла

# Обращение к парамтерам стуктуры:
#   hdr.ver, hdr.patientID, hdr.recordID, hdr.startdate, hdr.starttime,
#   hdr.ns, hdr.records, hdr.duration, hdr.samples,
#   hdr.labels, hdr.physicalDims, hdr.digitalMins, hdr.digitalMaxs ...
#
# Возвращает:
#   hdr    :: EDFHeader (структура)
#   record :: Matrix{Float64} размера (hdr.ns × N)
#
# Особенности реаоизации:
#   EDF допускает разные частоты дискретизации у разных каналов
#   следовательно разные количество отсчетов для каждого из каналов.
#   Чтобы сохранить record[ch, :], недостающие значения заполняем NaN.


# ---------------------------
# Структура заголовка EDF
# ---------------------------
struct EDFHeader
    # Основные поля стуктуры
    ver::Float64        # Версия файла
    patientID::String   # ID пациента
    recordID::String    # Описание записи
    startdate::String   # Дата записи
    starttime::String   # Время записи
    bytes::Int          # Размер заголовка в байтах
    records::Int        # Количество блокоов записи
    duration::Float64   # Длительность блока записи
    ns::Int             # Количество каналов данных

    # Поля по каналам
    labels::Vector{String}          # Имена каналов
    transducers::Vector{String}     # Тип датчика
    physicalDims::Vector{String}    # Единицы измерения (символы  mV, uV, %  и т.д.)
    physicalMins::Vector{Float64}   # Минимальное значение физической велечины
    physicalMaxs::Vector{Float64}   # Максимаьное значение физической велечины
    digitalMins::Vector{Float64}    # Минимальное цифровое значение
    digitalMaxs::Vector{Float64}    # Максмальное цифровое значение
    prefilters::Vector{String}      # Тип фильтра при записи
    samples::Vector{Int}            # Число отсчётов на одну запись для каждого канала
end

# ---------------------------
# Вспомогательные функции
# ---------------------------

# Чтение фиксированного количества байт и преобразование в строку
# EDF поля — ASCII, фиксированная длина, часто дополнены пробелами
read_fixed(io::IO, n::Integer) = String(read(io, n))

# Удаляем '\0' и обрезаем пробелы по краям
strip0(s::AbstractString) = strip(replace(s, '\0' => ' '))

# Очистка имени сигнала
sanitize_label(s::AbstractString) = replace(strip0(s), r"\W+" => "")

# Парсинг числа из ASCII поля
parse_num(T, s::AbstractString) = parse(T, strip0(s))

# Проверка little-endian
function is_little_endian()
    if isdefined(Base, :islittleendian)
        return Base.islittleendian()
    elseif isdefined(Sys, :islittleendian)
        return Sys.islittleendian()
    else
        return true
    end
end

# Чтение Int16 массива
function read_int16_le!(io::IO, buf::Vector{Int16})
    read!(io, buf)
    if !is_little_endian()
        @inbounds for i in eachindex(buf)
            u = reinterpret(UInt16, buf[i])
            u = (u >> 8) | (u << 8)
            buf[i] = reinterpret(Int16, u)
        end
    end
    return buf
end

# Нормализация целевого сигнала
# Возвращаем индексы сигналов, которые нужно читать
function normalize_target_signals(targetSignals, all_labels_sanitized::Vector{String})
    if targetSignals === nothing
        return collect(1:length(all_labels_sanitized))
    end

    # Если передали индексы
    if targetSignals isa AbstractVector{<:Integer}
        return collect(Int.(targetSignals))
    end

    # Если передали имя или массив имён
    wanted = String[]
    if targetSignals isa AbstractString
        push!(wanted, sanitize_label(targetSignals))
    elseif targetSignals isa AbstractVector{<:AbstractString}
        for s in targetSignals
            push!(wanted, sanitize_label(s))
        end
    else
        error("EDFREAD: Unsupported type for targetSignals: $(typeof(targetSignals)).")
    end

    idx = findall(lbl -> lbl in wanted, all_labels_sanitized)
    return collect(idx)
end


# ---------------------------
# Основная функция edfread
# ---------------------------
function edfread(fname::AbstractString;
                 assignToVariables::Bool = false,
                 targetSignals = nothing,
                 assignModule::Module = Main)

    io = open(fname, "r")

    try
        
        # 1) Чтение фиксированных заголовков

        ver       = parse_num(Float64, read_fixed(io, 8))
        patientID = strip0(read_fixed(io, 80))
        recordID  = strip0(read_fixed(io, 80))
        startdate = strip0(read_fixed(io, 8))   # дд.мм.гг
        starttime = strip0(read_fixed(io, 8))   # чч.мм.сс
        bytes     = parse_num(Int, read_fixed(io, 8))
        _reserved = read(io, 44)                   # зарезервировано, не используем
        records   = parse_num(Int, read_fixed(io, 8))
        duration  = parse_num(Float64, read_fixed(io, 8))
        ns        = parse_num(Int, read_fixed(io, 4))

        if ns <= 0
            error("EDFREAD: Некорректное число сигналов (ns=$ns).")
        end

        # 2) Чтение полей по каналам (ns * 256 байт)

        # labels: ns * 16
        all_labels = Vector{String}(undef, ns)
        for i in 1:ns
            all_labels[i] = sanitize_label(read_fixed(io, 16))
        end

        # Определяем, какие каналы читать
        target_idx = normalize_target_signals(targetSignals, all_labels)

        if isempty(target_idx)
            error("EDFREAD: Запрошенные сигналы не найдены в заголовке.")
        end
        for i in target_idx
            if i < 1 || i > ns
                error("EDFREAD: Индекс целевого канала вне диапазона: $i (ns=$ns).")
            end
        end

        # transducer: ns * 80
        all_transducers = Vector{String}(undef, ns)
        for i in 1:ns
            all_transducers[i] = strip0(read_fixed(io, 80))
        end

        # physical dimension (units): ns * 8
        all_units = Vector{String}(undef, ns)
        for i in 1:ns
            all_units[i] = strip0(read_fixed(io, 8))
        end

        # physical min/max: ns * 8
        all_pmin = Vector{Float64}(undef, ns)
        all_pmax = Vector{Float64}(undef, ns)
        # max
        for i in 1:ns
            all_pmin[i] = parse_num(Float64, read_fixed(io, 8))
        end
        # min
        for i in 1:ns
            all_pmax[i] = parse_num(Float64, read_fixed(io, 8))
        end

        # digital min/max: ns * 8
        all_dmin = Vector{Float64}(undef, ns)
        all_dmax = Vector{Float64}(undef, ns)
        for i in 1:ns
            all_dmin[i] = parse_num(Float64, read_fixed(io, 8))
        end
        for i in 1:ns
            all_dmax[i] = parse_num(Float64, read_fixed(io, 8))
        end

        # prefilter: ns * 80
        all_prefilter = Vector{String}(undef, ns)
        for i in 1:ns
            all_prefilter[i] = strip0(read_fixed(io, 80))
        end

        # samples per record: ns * 8
        all_samples = Vector{Int}(undef, ns)
        for i in 1:ns
            all_samples[i] = parse_num(Int, read_fixed(io, 8))
        end

        # reserved: ns * 32
        for _ in 1:ns
            _ = read_fixed(io, 32)
        end

        # 3) Формирование структуры hdr

        sel_labels      = all_labels[target_idx]
        sel_transducers = all_transducers[target_idx]
        sel_units       = all_units[target_idx]         
        sel_pmin        = all_pmin[target_idx]
        sel_pmax        = all_pmax[target_idx]
        sel_dmin        = all_dmin[target_idx]
        sel_dmax        = all_dmax[target_idx]
        sel_pref        = all_prefilter[target_idx]
        sel_samples     = all_samples[target_idx]
        sel_ns          = length(target_idx)

        hdr = EDFHeader(
            ver, patientID, recordID, startdate, starttime,
            bytes, records, duration, sel_ns,
            sel_labels, sel_transducers, sel_units,
            sel_pmin, sel_pmax, sel_dmin, sel_dmax, sel_pref, sel_samples
        )

        # 4) Чтение данных 

        function read_edf_data()
            if records == -1
                error("EDFREAD: records=-1 (неизвестно число записей). Требуется конечное число records.")
            end
            if records < 0
                error("EDFREAD: Некорректное число записей: $records.")
            end

            # Масштабирование (линейное):
            # scalefac = (physMax-physMin)/(digMax-digMin)
            # dc = physMax - scalefac*digMax
            scalefac_full = (all_pmax .- all_pmin) ./ (all_dmax .- all_dmin)
            dc_full       = all_pmax .- scalefac_full .* all_dmax

            # Матрица: nsel × maxLen (остаток заполняется - NaN)
            lengths = [all_samples[i] * records for i in target_idx]
            maxLen  = maximum(lengths)
            record  = fill(NaN, sel_ns, maxLen)

            # Позиции записи для каждого канала
            writepos = ones(Int, sel_ns)

            # Карта: индекс сигнала -> позиция в выбранном наборе (или 0)
            pos_in_selected = zeros(Int, ns)
            for (j, sig_i) in enumerate(target_idx)
                pos_in_selected[sig_i] = j
            end

            # Читаем все записи (records)
            for _recnum in 1:records
                for sig_i in 1:ns
                    nsamp = all_samples[sig_i]
                    if nsamp < 0
                        error("EDFREAD: Некорректное значение 'samples' для сигнала $sig_i: $nsamp.")
                    end

                    j = pos_in_selected[sig_i]
                    if j != 0
                        buf = Vector{Int16}(undef, nsamp)
                        read_int16_le!(io, buf)

                        sf = scalefac_full[sig_i]
                        dc = dc_full[sig_i]

                        p = writepos[j]
                        @inbounds for k in 1:nsamp
                            record[j, p] = Float64(buf[k]) * sf + dc
                            p += 1
                        end
                        writepos[j] = p
                    else
                        # пропускаем невыбранный сигнал
                        seek(io, position(io) + nsamp * 2)
                    end
                end
            end

            return record
        end

        record = read_edf_data()

        # 5) Присваем поля 

        if assignToVariables
            for ch in 1:hdr.ns
                name = hdr.labels[ch]
                sym  = Symbol(name)
                vecsig = vec(record[ch, :])
                @eval assignModule begin
                    $(sym) = $(vecsig)
                end
            end
            return hdr, Matrix{Float64}(undef, 0, 0)
        end

        #
        return hdr, record

    finally
        close(io)
    end
end