Спектральный анализ аритмий сердца

struct WFDBHeader
    record::String
    nsig::Int
    fs::Float64
    nsamp::Int

    # first channel (for backward compatibility)
    sigfile::String
    fmt::Int
    gain::Union{Float64,Nothing}
    adc_res::Union{Int,Nothing}
    adc_zero::Union{Int,Nothing}
    baseline::Union{Int,Nothing}
    desc::String
    comments::Vector{String}

    # full per-channel info (new, for universal support)
    sigfiles::Vector{String}
    fmts::Vector{Int}
    gains::Vector{Union{Float64,Nothing}}
    baselines::Vector{Union{Int,Nothing}}
    units::Vector{String}
    sig_names::Vector{String}
end

struct WFDBAnnotation
    sample::Int
    code::Int
    subtyp::Int
    chan::Int
    num::Int
    aux::String
end

# ---------------------------
# Helpers
# ---------------------------

# Accept String/SubString/etc.
toS(x::AbstractString) = String(x)

# Parse fmt from token like "212", "16", "212x2"
parse_fmt(s::AbstractString) = parse(Int, match(r"^\d+", toS(s)).match)

# Parse gain token variants:
#   182.149/mV
#   200(0)/mV
#   400
#   400(0)
# Return (gain, baseline, units)


function parse_gain_units(tok::AbstractString)
    tok = toS(tok)
    m = match(r"^([0-9.]+)(?:\(([-]?\d+)\))?(?:/(.+))?$", tok)
    if m === nothing
        return (nothing, nothing, "")
    end
    g = tryparse(Float64, m.captures[1])
    bl = m.captures[2] === nothing ? nothing : tryparse(Int, m.captures[2])
    u = m.captures[3] === nothing ? "" : String(m.captures[3])
    return (g, bl, u)
end

# Heuristic: signal name is usually last token (ECG, ECG1, MLII, ...)
# But some headers have tail with multiple tokens. We'll take last token,
# unless it looks numeric-like.

function parse_sig_name(parts::Vector{<:AbstractString})
    p = toS.(parts)
    # try last token
    lasttok = p[end]
    if tryparse(Float64, lasttok) === nothing
        return lasttok
    end
    # fallback: find first non-number-ish token from the end
    for i in length(p):-1:1
        if tryparse(Float64, p[i]) === nothing
            return join(p[i:end], " ")
        end
    end
    return p[end]
end

# sign-extend 12-bit
@inline function signext12(x::Int)::Int
    x >= 0x800 ? x - 0x1000 : x
end

# Compute physical units matrix if possible
function compute_phys(raw::Matrix{Int}, gains::Vector{Union{Float64,Nothing}}, baselines::Vector{Union{Int,Nothing}})
    nsig = size(raw,2)
    ok = all(ch -> (gains[ch] !== nothing && gains[ch] != 0), 1:nsig)
    ok || return nothing
    phys = Matrix{Float64}(undef, size(raw,1), nsig)
    for ch in 1:nsig
        bl = baselines[ch] === nothing ? 0 : baselines[ch]
        g  = gains[ch]
        phys[:,ch] = (raw[:,ch] .- bl) ./ g
    end
    return phys
end

# ---------------------------
# 1) Read header (.hea)
# ---------------------------

function read_header(base::AbstractString)::WFDBHeader
    hea_path = toS(base) * ".hea"
    lines = readlines(hea_path)

    comments = String[]
    data_lines = String[]
    for ln in lines
        s = strip(toS(ln))
        isempty(s) && continue
        if startswith(s, "#")
            push!(comments, s)
        else
            push!(data_lines, s)
        end
    end
    isempty(data_lines) && error("Пустой header: $hea_path")

    # Record line: <record> <nsig> <fs> <nsamp> ...
    r = split(strip(data_lines[1]))
    record = toS(r[1])
    nsig   = parse(Int, r[2])
    fs     = length(r) >= 3 ? parse(Float64, r[3]) : NaN
    nsamp  = length(r) >= 4 ? parse(Int, r[4]) : -1
    nsig < 1 && error("nsig < 1 в $hea_path")

    sigfiles  = String[]
    fmts      = Int[]
    gains     = Vector{Union{Float64,Nothing}}()
    baselines = Vector{Union{Int,Nothing}}()
    units     = String[]
    sig_names = String[]

    # Per-channel lines (nsig lines)
    for ch in 1:nsig
        idx = 1 + ch
        idx > length(data_lines) && error("Не хватает строк каналов в .hea (ожидали $nsig)")

        parts = split(strip(data_lines[idx]))
        push!(sigfiles, toS(parts[1]))

        fmt = parse_fmt(parts[2])
        push!(fmts, fmt)

        # Common WFDB: parts[3] is gain or gain/baseline/units
        g = nothing
        bl = nothing
        u = ""

        if length(parts) >= 3
            g, bl, u = parse_gain_units(parts[3])
        end

        # adc_res and adc_zero might exist as numeric tokens after gain
        # We'll use them only for the FIRST channel in backwards-compatible fields.
        # Baseline: if not inside gain(...), often equals adc_zero (5th token).
        # This is dataset-dependent, but works for many PhysioNet headers.
        adc_res = nothing
        adc_zero = nothing
        if length(parts) >= 4
            adc_res = tryparse(Int, parts[4])
        end
        if length(parts) >= 5
            adc_zero = tryparse(Int, parts[5])
            if bl === nothing
                bl = adc_zero
            end
        end

        push!(gains, g)
        push!(baselines, bl)
        push!(units, u)
        push!(sig_names, parse_sig_name(parts))
    end

    # Backwards-compatible "single" fields = first channel
    sigfile  = sigfiles[1]
    fmt      = fmts[1]
    gain     = gains[1]
    baseline = baselines[1]
    desc     = sig_names[1]  # keep old "desc" meaning as signal name

    # adc_res/adc_zero old fields: try from first channel line if present
    parts1 = split(strip(data_lines[2]))
    adc_res = length(parts1) >= 4 ? tryparse(Int, parts1[4]) : nothing
    adc_zero = length(parts1) >= 5 ? tryparse(Int, parts1[5]) : nothing
    if baseline === nothing
        baseline = adc_zero
    end

    return WFDBHeader(
        record, nsig, fs, nsamp,
        sigfile, fmt, gain, adc_res, adc_zero, baseline, desc, comments,
        sigfiles, fmts, gains, baselines, units, sig_names
    )
end

# ---------------------------
# 2) Read .dat (fmt 212 and 16)
# Always returns raw as Matrix N×nsig
# ---------------------------

function read_dat_212(base::AbstractString; sampfrom::Int=0, sampto::Union{Int,Nothing}=nothing)
    hdr = read_header(base)
    all(hdr.fmts .== 212) || error("Этот ридер поддерживает только fmt=212 для всех каналов; fmts=$(hdr.fmts)")

    dat_path = joinpath(dirname(toS(base)), hdr.sigfiles[1])
    bytes = read(dat_path)

    # 1 frame = 3 bytes = 1 time sample for up to 2 channels
    nframes_total = length(bytes) ÷ 3
    if hdr.nsamp > 0
        nframes_total = min(nframes_total, hdr.nsamp)
    end

    f_from = clamp(sampfrom + 1, 1, nframes_total)
    f_to   = sampto === nothing ? nframes_total : clamp(sampto, 1, nframes_total)
    f_to < f_from && error("sampto < sampfrom")

    nsig = hdr.nsig
    nsig > 2 && error("fmt 212 кодирует максимум 2 канала; hdr.nsig=$(nsig)")

    N = f_to - f_from + 1
    raw = Matrix{Int}(undef, N, nsig)

    out_i = 1
    frame_i = 1

    # iterate frames
    for k in 1:3:(3*nframes_total)
        b0 = Int(bytes[k])
        b1 = Int(bytes[k+1])
        b2 = Int(bytes[k+2])

        c1 = b0 | ((b1 & 0x0F) << 8)
        c2 = ((b1 & 0xF0) >> 4) | (b2 << 4)

        v1 = signext12(c1)
        v2 = signext12(c2)

        if frame_i >= f_from && frame_i <= f_to
            raw[out_i, 1] = v1
            if nsig == 2
                raw[out_i, 2] = v2
            end
            out_i += 1
        end

        frame_i += 1
        frame_i > f_to && break
    end

    phys = compute_phys(raw, hdr.gains, hdr.baselines)
    return hdr, raw, phys
end

function read_dat_16(base::AbstractString; sampfrom::Int=0, sampto::Union{Int,Nothing}=nothing)
    hdr = read_header(base)
    all(hdr.fmts .== 16) || error("Этот ридер поддерживает только fmt=16 для всех каналов; fmts=$(hdr.fmts)")

    dat_path = joinpath(dirname(toS(base)), hdr.sigfiles[1])

    io = open(dat_path, "r")
    try
        nsig = hdr.nsig

        filesize_bytes = stat(dat_path).size
        total_samples_allch = filesize_bytes ÷ 2      # Int16 => 2 bytes
        nframes_total = total_samples_allch ÷ nsig    # frames over time

        if hdr.nsamp > 0
            nframes_total = min(nframes_total, hdr.nsamp)
        end

        f_from = clamp(sampfrom + 1, 1, nframes_total)
        f_to   = sampto === nothing ? nframes_total : clamp(sampto, 1, nframes_total)
        f_to < f_from && error("sampto < sampfrom")

        N = f_to - f_from + 1
        raw = Matrix{Int}(undef, N, nsig)

        # seek to starting frame (interleaved channels)
        seek(io, (f_from - 1) * nsig * 2)

        for i in 1:N
            for ch in 1:nsig
                v = read(io, Int16) # little-endian in typical WFDB .dat on x86
                raw[i, ch] = Int(v)
            end
        end

        phys = compute_phys(raw, hdr.gains, hdr.baselines)
        return hdr, raw, phys
    finally
        close(io)
    end
end

"""
Universal reader for fmt=16 and fmt=212.
"""
function read_dat(base::AbstractString; sampfrom::Int=0, sampto::Union{Int,Nothing}=nothing)
    hdr = read_header(base)
    fmts = unique(hdr.fmts)
    length(fmts) == 1 || error("Смешанные форматы по каналам пока не поддержаны: $(hdr.fmts)")
    fmt = fmts[1]
    if fmt == 212
        return read_dat_212(base; sampfrom=sampfrom, sampto=sampto)
    elseif fmt == 16
        return read_dat_16(base; sampfrom=sampfrom, sampto=sampto)
    else
        error("fmt=$fmt не поддержан. Поддерживаются только 16 и 212.")
    end
end

# ---------------------------
# 3) Read annotations (MIT format)
# ---------------------------

const SKIP = 59
const NUM  = 60
const SUB  = 61
const CHN  = 62
const AUX  = 63

function read_annotations(base::AbstractString, ext::AbstractString)
    ann_path = toS(base) * "." * toS(ext)
    buf = read(ann_path)

    anns = WFDBAnnotation[]
    i = 1

    t = 0
    cur_num = 0
    cur_chan = 0
    cur_subtyp = 0

    while i + 1 <= length(buf)
        b0 = Int(buf[i])
        b1 = Int(buf[i+1])
        i += 2

        A = (b1 >> 2) & 0x3F
        I = ((b1 & 0x03) << 8) | b0

        # EOF marker
        if A == 0 && I == 0
            break
        end

        if A == SKIP
            i + 3 <= length(buf) || break
            w0 = Int(buf[i]); w1 = Int(buf[i+1]); w2 = Int(buf[i+2]); w3 = Int(buf[i+3])
            i += 4
            t = w0 | (w1<<8) | (w2<<16) | (w3<<24)
            continue
        elseif A == NUM
            cur_num = I
            continue
        elseif A == CHN
            cur_chan = I
            continue
        elseif A == SUB
            cur_subtyp = I
            continue
        elseif A == AUX
            n = I
            n <= 0 && continue
            i + n - 1 <= length(buf) || break
            aux_bytes = buf[i:i+n-1]
            i += n
            # remove NUL
            aux = String(filter(b -> b != 0x00, aux_bytes))
            push!(anns, WFDBAnnotation(t, AUX, cur_subtyp, cur_chan, cur_num, aux))
            continue
        else
            t += I
            push!(anns, WFDBAnnotation(t, A, cur_subtyp, cur_chan, cur_num, ""))
        end
    end

    return anns
end;

safe_read_ann(base::AbstractString, ext::AbstractString) =
    isfile(toS(base) * "." * toS(ext)) ? read_annotations(base, ext) : WFDBAnnotation[];