//@version=6
indicator("Close FFT Spectrum Pane", overlay=false, max_lines_count=500, max_labels_count=500, max_boxes_count=500)

int windowLen = input.int(64, "FFT窗口长度", minval=16, maxval=256, step=8)
int freqIndexInput = input.int(1, "观察的频率序号 (0=直流)", minval=0, maxval=256)
bool normalizeSpectrum = input.bool(true, "幅值归一化（除以窗口长度）")
int smoothLen = input.int(1, "幅值平滑长度 (1 = 不平滑)", minval=1, maxval=20)
bool showTable = input.bool(true, "显示主频表格")
int topCount = input.int(5, "主频数量", minval=1, maxval=12)

calcSpectrum(float[] samples, bool normalize) =>
    int size = array.size(samples)
    float[] mags = array.new_float()
    if size == 0
        mags
    float norm = normalize ? float(size) : 1.0
    int freqLimit = size / 2
    if freqLimit < 1
        freqLimit := 1
    float twoPi = 2.0 * math.pi
    for freq = 0 to freqLimit
        float sumReal = 0.0
        float sumImag = 0.0
        for n = 0 to size - 1
            float sample = array.get(samples, size - 1 - n)
            float angle = twoPi * float(freq) * float(n) / float(size)
            sumReal += sample * math.cos(angle)
            sumImag -= sample * math.sin(angle)
        float magnitude = math.sqrt(sumReal * sumReal + sumImag * sumImag) / norm
        array.push(mags, magnitude)
    mags

getTopIndexes(float[] mags, int count) =>
    int available = array.size(mags)
    int limit = count < available ? count : available
    int[] results = array.new_int()
    if limit == 0
        results
    float[] scratch = array.copy(mags)
    for rank = 0 to limit - 1
        float bestVal = na
        int bestIdx = -1
        for i = 0 to array.size(scratch) - 1
            float candidate = array.get(scratch, i)
            if na(candidate)
                continue
            if na(bestVal) or candidate > bestVal
                bestVal := candidate
                bestIdx := i
        if bestIdx == -1
            break
        array.push(results, bestIdx)
        array.set(scratch, bestIdx, na)
    results

var float[] priceBuffer = array.new_float()
var int lastWindowSetting = na
if na(lastWindowSetting) or lastWindowSetting != windowLen
    lastWindowSetting := windowLen
    array.clear(priceBuffer)

if not na(close)
    if array.size(priceBuffer) >= windowLen
        array.shift(priceBuffer)
    array.push(priceBuffer, close)

bool ready = array.size(priceBuffer) == windowLen
float[] spectrum = ready ? calcSpectrum(priceBuffer, normalizeSpectrum) : array.new_float()
int freqCount = array.size(spectrum)

int clampedFreqIdx = freqCount > 0 ? math.min(freqIndexInput, freqCount - 1) : 0
float paneMagnitude = ready and freqCount > 0 ? array.get(spectrum, clampedFreqIdx) : na
float panePlotted = smoothLen > 1 ? ta.sma(paneMagnitude, smoothLen) : paneMagnitude
plot(panePlotted, title="频率幅值", color=color.blue, linewidth=2)

var table spectrumTable = na
var int lastTableRows = na
if showTable
    if na(spectrumTable) or na(lastTableRows) or lastTableRows != topCount + 1
        if not na(spectrumTable)
            table.delete(spectrumTable)
        spectrumTable := table.new(position.top_right, 3, topCount + 1, bgcolor=color.new(color.black, 80), frame_color=color.new(color.gray, 60))
        lastTableRows := topCount + 1
else if not na(spectrumTable)
    table.delete(spectrumTable)
    spectrumTable := na
    lastTableRows := na

if showTable and not na(spectrumTable)
    color headerColor = color.new(color.white, 0)
    color valueColor = color.new(color.aqua, 0)
    color periodColor = color.new(color.silver, 0)
    table.cell(spectrumTable, 0, 0, "频率k", text_color=headerColor, text_halign=text.align_center)
    table.cell(spectrumTable, 1, 0, "幅值", text_color=headerColor, text_halign=text.align_center)
    table.cell(spectrumTable, 2, 0, "周期(Bar)", text_color=headerColor, text_halign=text.align_center)
    if ready and freqCount > 0
        int[] leaders = getTopIndexes(spectrum, topCount)
        int leaderCount = array.size(leaders)
        for row = 0 to topCount - 1
            int tableRow = row + 1
            if row < leaderCount
                int freqIdx = array.get(leaders, row)
                float amp = array.get(spectrum, freqIdx)
                float period = freqIdx == 0 ? na : float(windowLen) / float(freqIdx)
                float rounded = na(period) ? na : math.round(period * 100.0) / 100.0
                string freqText = str.tostring(freqIdx)
                string ampText = str.tostring(amp, format.mintick)
                string periodText = na(rounded) ? "∞" : str.tostring(rounded, format.mintick)
                table.cell(spectrumTable, 0, tableRow, freqText, text_color=headerColor)
                table.cell(spectrumTable, 1, tableRow, ampText, text_color=valueColor)
                table.cell(spectrumTable, 2, tableRow, periodText, text_color=periodColor)
            else
                table.cell(spectrumTable, 0, tableRow, "-", text_color=headerColor)
                table.cell(spectrumTable, 1, tableRow, "-", text_color=valueColor)
                table.cell(spectrumTable, 2, tableRow, "-", text_color=periodColor)
    else
        for row = 1 to topCount
            table.cell(spectrumTable, 0, row, "-", text_color=headerColor)
            table.cell(spectrumTable, 1, row, "-", text_color=valueColor)
            table.cell(spectrumTable, 2, row, "-", text_color=periodColor)