from bokeh.plotting import figure, output_file, vplot, show
import numpy as np
import random

n = 1024
a = 1.0
sample_period = 0.001
t = np.linspace(0, n * sample_period, n)
dt = t[1] - t[0]
data = np.abs(np.sin(t * 2 * np.pi) * a)

def fft(data, sample_period, power=False, use_db=True):
    dt = sample_period
    sp = np.fft.rfft(data)

    if power:
        spectrum = (np.abs(sp) * 2 * dt) ** 2
    else:
        spectrum = np.abs(sp) * 2 * dt

    if use_db:
        max_input = np.max(data)
        if power:
            spectrum = 20 * np.log10(spectrum / max_input)
        else:
            spectrum = 10 * np.log10(spectrum / max_input)
    n = len(data)
    freqs = np.fft.fftfreq(n, sample_period)

    # Ignore the negative part of frequency. It's because of symmetry of FFT.
    idx = np.argsort(freqs)
    idx = filter(lambda i: freqs[i] > 0, idx)

    return freqs[idx], spectrum[idx].real

freqs, spectrum = fft(data, dt, use_db=False)
freqs_db, spectrum_db = fft(data, dt, use_db=True)
output_file("lines.html")

p1 = figure(width=800, height=300, title="data")
p1.line(t, data, legend="data", line_width=1, color="red")

p2 = figure(width=800, height=300, title="FFT", x_axis_label="Frequency(Hz)", y_axis_label="Amplitude")
p2.line(freqs, spectrum, legend="data", line_width=2, color="blue")

p3 = figure(width=800, height=300, title="FFT(dB)", x_axis_label="Frequency(Hz)", y_axis_label="Amplitude(dB)")
p3.line(freqs_db, spectrum_db, legend="data", line_width=2, color="blue")

p = vplot(p1, p2, p3)

# Show the results.
show(p)
print(spectrum.real[10])