import math
import matplotlib.pyplot as plt
import numpy as np

def goertzel_sinusoid(freq, duration, sample_rate, amplitude):
    n = int(duration * sample_rate)
    omega = (2.0 * math.pi * freq) / sample_rate
    coeff = 2.0 * math.cos(omega)
    
    # Initialize state variables for cos wave
    # You can use 1.0 for q1 and q2 for a cos wave with slightly larger amplitude and a phase shift of 1/2 a sample
    # q1 previous sample in sinusoid
    q1 = math.cos(omega * (n - 1)) # use math.sin(omega * (n-1)) for sin wave
    # q2 previous previous sample
    q2 = math.cos(omega * (n - 2)) # use math.sin(omega * (n-2)) for sin wave
    result = np.zeros(n)
    
    for i in range(n):
        sample = coeff * q1 - q2
        q2 = q1
        q1 = sample
        result[i] = sample * amplitude
    
    return result

def sinusoid(n, cycles=1, phase=0):
    result = []
    for i in range(phase, n+phase):
        sample = math.cos(2.0 * cycles * math.pi * (i/n))
        result.append(sample)
    return result

def main():
    samp_per_cycle = 512
    cycles = 3
    len = cycles * samp_per_cycle
    signal = goertzel_sinusoid(1, cycles, samp_per_cycle, 1)
    reference = np.array(sinusoid(len, cycles))
    error = reference - signal
    print("max amplitude:", signal.max())
    print("max error:", error.max())
    x = np.array(range(0, len))
    plt.plot(x, signal)
    plt.plot(x, reference)
    plt.plot(x, error)
    #plt.plot(x, error*100000000000)
    plt.show()

if __name__ == "__main__":
    main()