import numpy as np
import pylab as pl

n_period_points = 32  # Number of points in period
dx = 1.0 / n_period_points
x_vals = np.arange(0.0, 1.0, dx)  # This is the interval to be analyzed

waveName = "Step Sawtooth"

wave_y_vals = [0.0, 0.0, 0.5, 0.5, 0.5, 0.5, 1.0, 1.0,
               1.0, 1.0, 1.5, 1.5, 1.5, 1.5, 2.0, 2.0,
               -2.0, -2.0, -1.5, -1.5, -1.5, -1.5, -1.0, -1.0,
               -1.0, -1.0, -0.5, -0.5, -0.5, -0.5, 0.0, 0.0]

print("PHY372 Spring 2024 Problem Set 6")
print("Fourier analysis and re-synthesis of a", waveName)
print("P372-PS06-FourierSeriesTemplate")
print("Matthew Oros Feb 2024")

m_vals = np.arange(0, n_period_points // 2)

s_list: list[np.ndarray] = []
c_list: list[np.ndarray] = []
for m in m_vals:
    s_list.append(np.sin(m * 2 * np.pi * x_vals))
    c_list.append(np.cos(m * 2 * np.pi * x_vals))

a_coeffs: list[float] = []
for c in c_list:
    a_coeffs.append(2.0 * np.dot(wave_y_vals, c) * dx)

b_coeffs: list[float] = []
for s in s_list:
    b_coeffs.append(2.0 * np.dot(wave_y_vals, s) * dx)

y_vals = np.zeros(n_period_points)
for i in range(len(m_vals)):
    y_vals += a_coeffs[i] * c_list[i] + b_coeffs[i] * s_list[i]

pl.plot(x_vals, y_vals, 'r', linewidth=2.0)
pl.plot(x_vals, wave_y_vals, linewidth=1.0)
pl.axis((0.0, 1.0, -2.5, 2.5))
pl.xlabel('position x')
pl.ylabel('function')
pl.title('Fourier re-synthesis of ' + waveName)

pl.show()