using EngeeDSP
using LinearAlgebra
using Statistics

function calculate_99_percent_bandwidth(frequencies, power; band_range=nothing)
    if band_range !== nothing
        min_freq, max_freq = band_range
        idx = findall(x -> min_freq <= x <= max_freq, frequencies)
        if isempty(idx)
            error("Нет данных в указанном диапазоне частот [$min_freq, $max_freq] Гц")
        end
        frequencies = frequencies[idx]
        power = power[idx]
    end
    total_power = sum(power)
    if total_power == 0
        return 0.0, 0.0, 0.0
    end
    normalized_power = power / total_power
    sorted_indices = sortperm(normalized_power, rev=true)
    cumulative = cumsum(normalized_power[sorted_indices])
    idx_99 = findfirst(x -> x >= 0.99, cumulative)
    if idx_99 === nothing
        idx_99 = length(cumulative)
    end
    freq_indices_99 = sorted_indices[1:idx_99]
    freq_values = frequencies[freq_indices_99]
    bandwidth = maximum(freq_values) - minimum(freq_values)
    freq_order = sortperm(frequencies)
    cumulative_by_freq = cumsum(normalized_power[freq_order])
    idx_low = findfirst(x -> x >= 0.005, cumulative_by_freq)
    idx_high = findfirst(x -> x >= 0.995, cumulative_by_freq)
    if idx_low !== nothing && idx_high !== nothing
        freq_sorted = frequencies[freq_order]
        bandwidth_symmetric = freq_sorted[idx_high] - freq_sorted[idx_low]
    else
        bandwidth_symmetric = bandwidth
    end
    centroid = sum(frequencies .* normalized_power) / sum(normalized_power)
    return bandwidth, bandwidth_symmetric, centroid
end

function bandwidth_99_direct(signal, sample_rate; band_range=nothing)
    nfft = min(1024, length(signal))
    segments = max(1, floor(Int, length(signal) / nfft))
    psd_total = zeros(nfft ÷ 2 + 1)
    for i in 1:segments
        start_idx = (i-1)*nfft + 1
        end_idx = min(i*nfft, length(signal))
        segment = signal[start_idx:end_idx]
        if length(segment) < 64
            continue
        end
        window = 0.5 .- 0.5*cos.(2π*(0:length(segment)-1)/(length(segment)-1))
        windowed = segment .* window
        fft_result = EngeeDSP.fft(windowed)
        psd = abs2.(fft_result) / (sample_rate * sum(window.^2))
        if length(psd) % 2 == 0
            idx = 1:length(psd)÷2 + 1
        else
            idx = 1:(length(psd)+1)÷2
        end
        psd_segment = psd[idx]
        if length(psd_segment) > 2
            psd_segment[2:end-1] .*= 2
        end
        if length(psd_segment) <= length(psd_total)
            psd_total[1:length(psd_segment)] .+= psd_segment
        end
    end
    psd_avg = psd_total ./ segments
    frequencies = range(0, sample_rate/2, length=length(psd_avg))
    if band_range !== nothing
        min_freq, max_freq = band_range
        idx = findall(x -> min_freq <= x <= max_freq, frequencies)
        if !isempty(idx)
            frequencies = frequencies[idx]
            psd_avg = psd_avg[idx]
        end
    end
    bandwidth, bandwidth_sym, centroid = calculate_99_percent_bandwidth(frequencies, psd_avg)
    return bandwidth, bandwidth_sym, centroid, frequencies, psd_avg
end

function bandwidth_99_spectrumanalyzer(signal, sample_rate; band_range=nothing, freq_span=nothing)
    if freq_span !== nothing
        span_start, span_end = freq_span
        scope = EngeeDSP.spectrumAnalyzer(
            SampleRate = sample_rate,
            SpectrumType = "power",
            SpectrumUnits = "Watts",
            Method = "welch",
            Window = "Hann",
            RBWSource = "auto",
            FrequencySpan = "span-and-center-frequency",
            Span = span_end - span_start,
            CenterFrequency = (span_start + span_end) / 2,
            PlotAsTwoSidedSpectrum = false,
            AveragingMethod = "exponential",
            ForgettingFactor = 1.0,
            VBWSource = "auto",
            OverlapPercent = 0
        )
    else
        scope = EngeeDSP.spectrumAnalyzer(
            SampleRate = sample_rate,
            SpectrumType = "power",
            SpectrumUnits = "Watts",
            Method = "welch",
            Window = "Hann",
            RBWSource = "auto",
            FrequencySpan = "full",
            PlotAsTwoSidedSpectrum = false,
            AveragingMethod = "exponential",
            ForgettingFactor = 1.0,
            VBWSource = "auto",
            OverlapPercent = 0
        )
    end
    scope(signal)
    spectrum_data = getSpectrumData(scope)
    if haskey(spectrum_data, "spectrum") && haskey(spectrum_data, "frequencies")
        spectrum = spectrum_data["spectrum"]
        frequencies = spectrum_data["frequencies"]
        if ndims(spectrum) > 1
            spectrum = vec(mean(spectrum, dims=2))
        end
        if ndims(frequencies) > 1
            frequencies = vec(frequencies)
        end
        if band_range !== nothing
            min_freq, max_freq = band_range
            idx = findall(x -> min_freq <= x <= max_freq, frequencies)
            if !isempty(idx)
                frequencies = frequencies[idx]
                spectrum = spectrum[idx]
            end
        end
        bandwidth, bandwidth_sym, centroid = calculate_99_percent_bandwidth(frequencies, spectrum)
        release!(scope)
        return bandwidth, bandwidth_sym, centroid, frequencies, spectrum
    else
        release!(scope)
        error("Не удалось получить данные спектра")
    end
end

function generate_test_signal(fs, duration)
    t = range(0, duration, step=1/fs)
    f1 = 1000
    f2 = 2500
    f3 = 4000
    signal = 1.0 * EngeeDSP.sin.(2π * f1 * t) +
             0.5 * EngeeDSP.sin.(2π * f2 * t) +
             0.2 * EngeeDSP.sin.(2π * f3 * t) +
             0.03 * randn(length(t))
    return signal
end

fs = 10000
duration = 0.1
signal = generate_test_signal(fs, duration)
println("\n1. Полный диапазон частот:")
bw1, bw1_sym, centroid1, freqs1, power1 = bandwidth_99_direct(signal, fs)
println("   Ширина полосы: $(round(bw1, digits=2)) Гц")
println("   Симметричная ширина: $(round(bw1_sym, digits=2)) Гц")
println("   Центральная частота: $(round(centroid1, digits=2)) Гц")
println("\n2. Ограниченный диапазон [500, 3000] Гц:")
bw2, bw2_sym, centroid2, freqs2, power2 = bandwidth_99_direct(signal, fs; band_range=(500.0, 3000.0))
println("   Ширина полосы: $(round(bw2, digits=2)) Гц")
println("   Симметричная ширина: $(round(bw2_sym, digits=2)) Гц")
println("   Центральная частота: $(round(centroid2, digits=2)) Гц")
println("\n3. SpectrumAnalyzer с диапазоном [1000, 3500] Гц:")
bw3, bw3_sym, centroid3, freqs3, power3 = bandwidth_99_spectrumanalyzer(signal, fs; band_range=(1000.0, 3500.0))
println("   Ширина полосы: $(round(bw3, digits=2)) Гц")
println("   Симметричная ширина: $(round(bw3_sym, digits=2)) Гц")
println("   Центральная частота: $(round(centroid3, digits=2)) Гц")
println("\n4. SpectrumAnalyzer с настройкой freq_span [500, 4500] Гц:")
bw4, bw4_sym, centroid4, freqs4, power4 = bandwidth_99_spectrumanalyzer(signal, fs; freq_span=(500.0, 4500.0))
println("   Ширина полосы: $(round(bw4, digits=2)) Гц")
println("   Симметричная ширина: $(round(bw4_sym, digits=2)) Гц")
println("   Центральная частота: $(round(centroid4, digits=2)) Гц")

p1 = plot(freqs1, 10*log10.(power1 .+ 1e-10), 
         label="Полный диапазон", 
         linewidth=2,
         xlabel="Частота (Гц)", 
         ylabel="Мощность (дБ)",
         title="Прямой БПФ - Полный диапазон",
         grid=true)
display(p1)
p2 = plot(freqs2, 10*log10.(power2 .+ 1e-10), 
         label="[500, 3000] Гц", 
         linewidth=2,
         xlabel="Частота (Гц)", 
         ylabel="Мощность (дБ)",
         title="Прямой БПФ - Ограниченный диапазон",
         grid=true)
display(p2)
p3 = plot(freqs3, 10*log10.(power3 .+ 1e-10), 
         label="SpectrumAnalyzer [1000, 3500] Гц", 
         linewidth=2,
         xlabel="Частота (Гц)", 
         ylabel="Мощность (дБ)",
         title="SpectrumAnalyzer с фильтрацией band_range",
         grid=true)
display(p3)
p4 = plot(freqs4, 10*log10.(power4 .+ 1e-10), 
         label="SpectrumAnalyzer с freq_span [500, 4500] Гц", 
         linewidth=2,
         xlabel="Частота (Гц)", 
         ylabel="Мощность (дБ)",
         title="SpectrumAnalyzer с настройкой freq_span",
         grid=true)
display(p4)

p5 = bar(["Полный", "[500,3000]", "SA [1000,3500]", "SA span [500,4500]"], 
         [bw1_sym, bw2_sym, bw3_sym, bw4_sym],
         xlabel="Метод",
         ylabel="Ширина полосы (Гц)",
         title="Сравнение симметричной ширины полосы 99%",
         color=[:blue :red :green :purple],
         legend=false,
         xtickfontsize=8)