import numpy as np
import matplotlib.pyplot as plt
#%%
"""EDIT THESE"""
def dft(x):
    return np.fft.fft(x)

def d_idft(coef, x, slope=True):
    N = len(coef)
    coef = coef * (1 - np.exp(-1j * 2 * np.pi / N * np.arange(N)))
    slope = (x[-1] - x[0]) if slope else 0
    coef += slope
    return np.fft.ifft(coef)[1:], coef

#%%###########################################################################
N = 128
x = np.zeros(N)
x[0]  =  1
x[-2] = -1
#%%
coef = dft(x)
dx   = np.diff(x)
dx_est1, dcoef1 = d_idft(coef, x, slope=False)
dx_est2, dcoef2 = d_idft(coef, x, slope=True)
dx_est3, dcoef3 = np.fft.ifft(np.fft.fft(dx)), np.fft.fft(dx)
#%%
print("{:.64}".format(np.mean(np.abs(dx - dx_est1))))
print("{:.64}".format(np.mean(np.abs(dx - dx_est2))))
print("{:.64}".format(np.mean(np.abs(dx - dx_est3))))
#%%###########################################################################
theta,   r = np.angle(dcoef1), np.abs(dcoef1)
theta2, r2 = np.angle(dcoef2), np.abs(dcoef2)
theta3, r3 = np.angle(dcoef3), np.abs(dcoef3)
#%%
plt.plot(r); plt.plot(r2); plt.plot(r3); plt.show()  # can comment line below
theta = np.unwrap(theta); theta2 = np.unwrap(theta2); theta3 = np.unwrap(theta3)
plt.plot(theta); plt.plot(theta2); plt.plot(theta3); plt.show()
#%%###########################################################################
print("Y[0]:\nw/o slope = {}\nw/slope = {}\nground truth = {}\n".format(
    dcoef1[0], dcoef2[0], dcoef3[0]))
#%%
th = 1e-14
print("Slope term:",
      "PASSED" if abs(dcoef2[0] - dcoef3[0]) < th else "FAILED")
print("Reconstruction:",
      "PASSED" if sum(np.abs(dx_est2 - dx))  < th else "FAILED")
if len(dcoef2) == len(dx):
    print("len(Y): PASSED")
    print("Coefficients:",
          "PASSED" if sum(np.abs(dcoef2 - dcoef3)) < th else "FAILED")
else:
    print("len(Y): FAILED\nCoefficients: FAILED")