engee.clear()

# gr()
plotlyjs()

include("$(@__DIR__)/read.jl")

let
    installed_packages = collect(x.name for (_, x) in Pkg.dependencies() if x.is_direct_dep)
    list_packages = ["Statistics"]
    for pack in list_packages
        pack in installed_packages || Pkg.add(pack)
    end
end
using Statistics

import EngeeDSP.Functions: fft, butter, filter, findpeaks, conv, freqz

function calc_fft(x, fs)
    len = length(x)
    S = fft(x)  # DFT
    
    # 单向频谱
    K = len ÷ 2
    S = S[1:K]      
    freq = (0:K-1) .* fs ./ len     # 频率轴从0到fs/2

    # 单侧振幅谱
    S_amp = 2 .* abs.(S) ./ len
    S_amp[1] = 0.0

    # 单侧振幅谱，单位为dB
    S_dB = 20 .* log10.(S_amp .+ 1e-12)

    return freq, S_amp, S_dB
end

data = "$(@__DIR__)/mitdb/100"

# 读取头文件并获取采样率
hdr = read_header(data)
fs = Float64(hdr.fs)

# 选择用于分析的信号片段
t_start = 5.0 * 60.0        # 片段从5分钟开始。
t_dur   = 5.0  * 60.0      # 片段的持续时间为5分钟。

# 将时间从秒转换为倒计时数字
sampfrom = Int(round(t_start * fs))
sampto   = sampfrom + Int(round(t_dur * fs))

# 读取ECG信号的选定片段
_, raw, phys = read_dat_212(data; sampfrom=sampfrom, sampto=sampto)

# 选择用于分析的引导通道
ecg = phys === nothing ? Float64.(raw[:,1]) : Float64.(phys[:,1])

# 计算计数数并形成时间轴
N = length(ecg)
t = collect(0:N-1) ./ fs

p1 = plot(t, ecg,
    xlabel="时间，从",
    ylabel="振幅，mV",
    label="原始信号"
)

freq0, S0_amp, _ = calc_fft(ecg, fs)

p2 = plot(freq0, S0_amp,
    xlabel="频率，赫兹",
    ylabel="振幅，mV",
    label="原始信号",
    xlim=(0, fs/2)
)

f_center = 60.0     # 谐振频率，Hz
bw = 6.0        # 滤波器带宽，Hz
f1 = f_center - bw/2      # 下限频率，Hz
f2 = f_center + bw/2      # 上限频率，Hz

# 归一化频率相对于采样率
wn1 = f1 / (fs / 2)
wn2 = f2 / (fs / 2)

order = 6      # 投影滤波器的顺序

b, a = butter(order, [wn1, wn2], "stop")    # 巴特沃斯陷波滤波器的合成

# 滤波器幅频响应(AFC)的计算
h, w = freqz(b, a, 2048)
f_resp = w .* fs ./ (2π)    # 将角频率转换为Hz

# 频率响应的计算（以dB为单位）
amp_resp = 20 .* log10.(abs.(h) .+ 1e-12)

plot(f_resp, amp_resp,
    xlabel="频率，赫兹",
    ylabel="频率响应，dB",
    legend=false,
    xlim=(0, fs/2)
)

ecg_filt = filter(b, a, ecg)        # 心电信号滤波

freq1, S1_amp, _ = calc_fft(ecg_filt, fs)

plot!(p2, freq1, S1_amp, label="过滤后")

plot!(p1, t, ecg_filt, label="过滤后")

d_ecg = vcat(0.0, diff(ecg_filt))   # 心电信号导数的计算

plot(t, d_ecg,
    xlabel="时间，从",
    ylabel="导数的幅度，rel。 单位",
    title="信号的导数",
    legend=false
)

sq_ecg = d_ecg .^ 2     # 平方信号的导数

plot(t, sq_ecg,
    xlabel="时间，从",
    ylabel="振幅，相对单位",
    title="导数的平方",
    legend=false
)

# 使用移动平均方法平滑信号
win = max(3, Int(round(0.010 * fs)))
kernel = ones(win) ./ win

det_sig = conv(sq_ecg, kernel)  # 使用卷积的信号平滑
det_sig = det_sig[1:N]  # 将信号长度减小到原始大小

plot(t, det_sig,
    xlabel="时间，从",
    ylabel="振幅，相对单位",
    title="用于检测R峰的信号",
    legend=false
)

thr = mean(det_sig) + 1.5 * std(det_sig)    # 自适应检测阈值
min_dist = Int(round(0.5 * fs))   # R峰之间的最小距离

# 搜索R峰
pks, locs = findpeaks(
    det_sig, out=:data,
    MinPeakDistance = min_dist,
    MinPeakHeight  = thr
)

plot(t, det_sig,
    xlabel="时间，从",
    ylabel="振幅，相对单位",
    title="发现R峰",
    legend=false
)
r_times = (locs .- 1) ./ fs   # 将找到的峰值的指数转换为秒

scatter!(r_times, pks, markersize=3)  # 在检测信号上显示发现的峰值

# 澄清R峰的位置
halfwin = Int(round(0.05 * fs))   # 搜索框
r_locs = Int[]

for idx in locs
    # 找到的峰周围区域的边界
    l = max(1, idx - halfwin)
    r = min(N, idx + halfwin)

     # 搜索选定区域内的最大值
    seg = ecg_filt[l:r]
    _, imax = findmax(seg)

    push!(r_locs, l + imax - 1)
end

r_locs = unique(sort(r_locs))   # 删除重复索引

r_times = (r_locs .- 1) ./ fs   # 发现的R峰的时间
r_vals = ecg_filt[r_locs]       # 发现的R峰的振幅

plot(t, ecg_filt,
    xlabel="时间，从",
    ylabel="振幅，mV",
    title="发现R峰",
    legend=false
)

scatter!(r_times, r_vals, markersize=3)

rr = diff(r_times)  # 相邻R峰之间的差异，具有

rr_ms = rr.* 1e3  # 将RR间隔转换为ms

rr_time = r_times[2:end]    # RR间隔的时间轴

heart_rate = 60.0 ./ rr # 计算心率、心跳/分钟

plot(rr_time, rr_ms,
    xlabel="时间，从",
    ylabel="RR-间隔，ms",
    title="RR间隔",
    legend=false
)

plot(rr_time, heart_rate,
    xlabel="时间，从",
    ylabel="心率，心跳/分钟",
    title="心率",
    legend=false
)