import numpy as np
import matplotlib.pyplot as plt

def eliminate_repeats(arr):
    arr = np.asarray(arr, dtype=bool)
    repeats = np.concatenate(([False], arr[1:] == arr[:-1]))
    return np.where(repeats, False, arr)

def max_min(U):
    Uabs =np.abs(U)
    m = np.max(Uabs,axis=1)
    x_idx = np.argmax(Uabs,axis=1)
    eps = m[np.argmax(m)+1]-np.max(m)
    tx = m>np.max(m)-np.abs(eps)/2
    tx = eliminate_repeats(tx)
    #x & t indices of maxima/minima of U
    return x_idx[tx],tx
    
def freq(U,t):
    _, tx = max_min(U)
    w = 1/(np.diff(t[tx]))*np.pi
    return w

def vel(U,x,t):
    x_idx, tx = max_min(U)
    return np.diff(x[x_idx])/np.diff(t[tx])

def measure_all(U,x,t):
    v = np.average(vel(U,x,t))
    gamma = 1/np.sqrt(1-v**2)
    f = np.average(freq(U,t)*gamma)
    print(f"Frequency: {f:.2f}, velocity: {v:.2f}")
    return (f,v)

def u_t(U,dt):
    Ut = np.empty_like(U)
    Ut[1:-1,:] = (U[2:,:] - U[:-2,:]) / (2*dt)
    Ut[0,:]    = (U[1,:] - U[0,:]) / dt 
    Ut[-1,:]   = (U[-1,:] - U[-2,:]) / dt  
    return Ut

def u_x(U,dx):
    Ux = np.empty_like(U)
    Ux[:,1:-1] = (U[:,2:] - U[:,:-2]) / (2*dx)
    Ux[:,0]    = (U[:,1] - U[:,0]) / dx
    Ux[:,-1]   = (U[:,-1] - U[:,-2]) / dx  
    #return Ux 
    return (np.roll(U, -1, axis=1) - np.roll(U, 1, axis=1)) / (2*dx)

def energy_density(U,V,dx,dt):
    return  0.5*(u_x(U,dx)**2+u_t(U,dt)**2)+V