import numpy as np
import sounddevice as sd
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt
import time

import taichi as ti
ti.init(arch=ti.cpu)

dx = 0.01 #space step
dt = 0.001 #time step

xmin = 0
xmax = 2*np.pi 

N = int(xmax/dx)

x_range = np.linspace(xmin,xmax,N)
substeps = int(1 / 60 // dt)

v = 1 #speed

gamma = 0.001
gamma0 = -0.05

u0 = ti.field(float, shape=(N,))#np.zeros((N))

# IMPULSE
i_index = int(N/2) #where the impulse is

#u0[i_index] = 0.5

# #not impulse
# i_index = int(N/2) #where the wave center is
# h = 0.5

# for i in range((N/2-5),(N/2+5)):
#     u[i] = 



#compute initial values
u_prev = u0
u_prevprev = ti.field(float, shape=(N,))#np.zeros(N)
for i in range(1,(N-2)):
    u_prevprev[i] = u_prev[i] + 1/2*((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])

u_next = ti.field(float, shape=(N,))#np.zeros(N)
# print(u_prev)
# print(u_prevprev)

#@ti.kernel
#def nextStep(u_prev: ti.types.ndarray(),u_prevprev:ti.types.ndarray()):
# def nextStep(u_prev,u_prevprev):
#     u_next = np.zeros(N)#ti.field(float, shape=(N))#
#     #ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
#     for i in range(1,(N-2)):
#         u_next[i] = -u_prevprev[i] + 2*u_prev[i] +((v**2)*(dx**2)/(dt**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
#     #boundary conditions
#     u_next[0] = 0
#     u_next[N-1] = 0
#     #update priors
#     return u_next

@ti.kernel
def nextStep():
    ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
    for i in range(1,(N-2)):
        #u_next[i] = -1*u_prevprev[i] + 2*u_prev[i] +((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
        #u_next[i] = -1*u_prevprev[i] + 2*u_prev[i] +((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1] + gamma0*((u_prev[i]-u_prevprev[i])/dt))
        u_next[i] = -1*u_prevprev[i] + 2*u_prev[i] +((v**2)*(dt**2)/(dx**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1] + gamma*((u_prev[i+1]-2*u_prev[i]+u_prev[i-1] - (u_prevprev[i+1]-2*u_prevprev[i]+u_prevprev[i-1]))/(dt)))
    #boundary conditions
    u_next[0] = 0
    u_next[N-1] = 0


@ti.kernel
def updatePrev():
    for i in range(N-1):
        u_prevprev[i] = u_prev[i]
        u_prev[i] = u_next[i]

# def update(u0,xmax,xmin,dx,dt,N):
#     u_next = np.zeros((N))
#     u_prev = u0
#     u_prevprev = np.zeros(N)
#     for i in range(1,(N-2)):
#         u_prevprev[i] = u_prev[i] + 1/2*((v**2)*(dx**2)/(dt**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
#     while True:
#         for i in range(1,(N-2)):
#             u_next[i] = -u_prevprev[i] + 2*u_prev[i] +((v**2)*(dx**2)/(dt**2))*(u_prev[i+1]-2*u_prev[i]+u_prev[i-1])
#         u_next[0] = 0
#         u_next[N-1] = 0
#         u_prevprev = u_prev
#         u_prev = u_next
#         yield u_next
# my_iter = update(u0,xmax,xmin,dx,dt,N)


x_range_resize = np.linspace(0,1,N)
offset = ti.field(float,shape=(N,))
offset.fill(0.5)

gui = ti.GUI('click to add impulse', res=(640, 400))
while gui.running:
    if gui.get_event(ti.GUI.PRESS,ti.GUI.RMB):
        mouse_x, mouse_y = gui.get_cursor_pos()
        u_prev[int(mouse_x*N)] = mouse_y*1.5 - 0.5
        u_prevprev[int(mouse_x*N)] = mouse_y*1.5 - 0.5
    for i in range(substeps):
        nextStep()
        updatePrev()
    pos = np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
    
    gui.lines(begin=pos[:(N-2)],end=pos[1:(N-1)],radius=2,color=0x068587)
    gui.show()

#sd.play(pos[:,1], 44100)

tune = pos[:,1] #just the y values.
tunef = fft(tune)
xf = fftfreq(N,dx)
plt.plot(xf,tunef)
plt.title('FFT of final wave')
plt.show()

tunemax = np.abs(tune[0:N//2].copy())
ind = []
for i in range(5):
     ind.append(np.argmax(tunemax))
     tunemax[np.argmax(tunemax)] = 0

print(xf[ind])

#temp_x = np.linspace()
f_s = 44100 #sampling rate
s = 3 #duration in seconds

def resynth(f):
    f_s = 44100 #sampling rate
    s = 3 #duration in seconds
    samples = np.zeros(round(f_s*s))
    W = 2*np.pi*f/f_s #frequency

    for n in range(round(f_s*s)):
            samples[n] = np.sin(W*n)
    return samples

samples = np.zeros(round(f_s*s))
for i in range(5):
    samples = np.add(samples, 10*resynth(30*xf[ind[i]]))


# f = 200 #frequency of signal
# f_s = 44100 #sampling rate
# s = 5 #duration in seconds
# samples = np.zeros(round(f_s*s))
# W = 2*np.pi*f/f_s #frequency

# for n in range(round(f_s*s)):
# 		samples[n] = np.sin(W*n)
		
plt.plot(samples)
plt.title('reconstituted wave, close window to play sound')
plt.show()

sd.play(samples,f_s)

time.sleep(3.0)

sd.stop()