import numpy as np
from scipy.fft import fft
import soundfile as sf
from numpy import log10
from matplotlib import pyplot as plt
import math

# Average Noise Spectral Level
def ANSL(u, fs, M, A=0, N=None, D=None, w=None):
    """
    :param u: Voltage data (Unit: V)
    :param fs: Sampling rate (Unit: Hz)
    :param M: Sensitivity (Unit: dBV re 1uPa)
    :param A: Preamplifier (Unit: dB)
    :param N: Segment length
    :param D: Overlap length
    :param w: Window
    :return: Spectral level
    """
    L = len(u)
    if N == None:
        if L >= fs:
            N = round(fs)
        else:
            N = L
    elif N > L:
        raise Exception("Segment length is larger than data length.")

    if D == None:
        D = 0

    if w is None:
        w = np.hamming(N)
    G = np.mean(np.square(abs(w)))

    K = 1 + math.floor((L-N) / (N-D))

    T = N / fs
    df = 1 / T
    X2 = np.zeros(N)
    for k in range(K):
        X = fft(w * u[round(k*(N-D)):round(k*(N-D)+N)])
        X2 += np.square(abs(X))

    spl = 10 * log10(X2) - 10 * log10(N*fs*G*K/2) - M - A
    freq = np.arange(0, fs/2, df)
    spl = spl[0:len(freq)]

    return freq, spl


# 1/3 Octave Spectral Level
def OneThirdOct_SPL(freq, spl, f_start, Npoints):
    """
    :param freq: Frequency
    :param spl: 1 Hz bandwidth spectral level
    :param f_start: Starting frequency of the 1/3 octave
    :param Npoints: The number of frequency points
    :return: 1/3 Octave Spectral Level
    """
    fC = f_start * np.power(2, np.arange(Npoints)/3)
    fL = f_start * np.power(2, np.arange(Npoints)/3 - 1.0/6)
    fH = f_start * np.power(2, np.arange(Npoints)/3 + 1.0/6)

    spl_oct = np.zeros(Npoints)
    for it in range(Npoints):
        # idx = (freq >= fL[it]) & (freq <= fH[it])
        idx = np.where((freq >= fL[it]) & (freq <= fH[it]))
        # print(idx)
        # print(spl[idx])
        spl_oct[it] = 10 * log10(np.mean(np.power(10, spl[idx]/10)))
        # print(spl_oct[it])

    return fC, spl_oct

def call_oanp(soundFile, M, A=0, isPlot=False):
    [x, fs] = sf.read(soundFile)
    u = 3 * x           # Voltage data
    # M = -164.0          # Sensitivity
    # A = 0               # Preamplifier
    N = fs              # Segment length
    D = N/2             # Overlap length
    w = np.hanning(N)   # Window

    # Average Noise Spectral Level
    [freq, spl] = ANSL(u, fs, M, A, N, D, w)
    # 1/3 Octave Spectral Level
    [fC, spl_oct] = OneThirdOct_SPL(freq, spl, f_start=20, Npoints=28)

    with open('oanp_data.txt', 'w') as f:
        f.write('frequency and noise level \r\n')
        for i in range(len(fC)):
            f.write(f'{round(fC[i])} \t {round(spl_oct[i],2)}\r\n')

    if isPlot:
        plt.semilogx(freq[20:10000], spl[20:10000])
        plt.semilogx(fC, spl_oct)
        plt.xlabel("Frequency / Hz")
        plt.ylabel("Noise Level / dB")
        plt.show()


def oancor(x,fs,sl=None,ol=None,fmin=None,fmax=None):
    """
    :param x: M*N array (Unit: none)
    :param fs: Sampling rate (Unit: Hz)
    :param sl: Segment length (Unit: none)
    :param ol: Overlap length (Unit: none)
    :return: [frequency list, Correlation array]
    """
    M=x.shape[0]
    N=x.shape[1]
    if sl==None:
        sl=round(fs)
    if ol==None:
        ol=0
    if fmin==None:
        fmin=0
    if fmax==None:
        fmax=fs


    if sl>=N :
        L=1
    else:
        L=1+math.floor((N-sl)/(sl-ol))

    f=np.arange(0,fs,fs/sl)
    f_ind=np.where((f>=fmin)&(f<=fmax))[0]
    f=f[f_ind]
    C=np.zeros((sl,M,M),np.complex64)
    for it in range(L):
        X=fft(x[:,it*(sl-ol):it*(sl-ol)+sl])
        X=X[:,f_ind[0]:(f_ind[-1]+1)]
        for jt in range(len(f)):
            cw=np.outer(X[:,jt],X[:,jt].conjugate())
            C[jt,:,:]=C[jt,:,:]+cw/np.sqrt(np.real(np.outer(np.diagonal(cw),np.diagonal(cw))))
    C=C/L
    #freq=np.arange(0,sl,1)*fs/sl
    return f,C

def oandi(x,fs,d,dtheta=None,sl=None,ol=None,c=None,fmin=None,fmax=None):
    """
    :param x: M*N array (Unit: none)
    :param fs: Sampling rate (Unit: Hz)
    :param d: M element depth in ascending order (Unit: m)
    :param dtheta: angle interval (Unit: degree)
    :param sl: Segment length (Unit: none)
    :param ol: Overlap length (Unit: none)
    :param c: sound velocity (Unit: m/s)
    :return: [theta, frequency  list, directionality array]
    """
    M=x.shape[0]
    N=x.shape[1]
    if c==None:
        c=1500
    if dtheta==None:
        dtheta=1
    if sl==None:
        sl=round(fs)
    if ol==None:
        ol=0
    if fmin==None:
        fmin=0
    if fmax==None:
        fmax=fs

    if sl>=N :
        L=1
    else:
        L=1+math.floor((N-sl)/(sl-ol))

    D=d-(d[-1]-d[0])/2
    theta=np.arange(-90,90,dtheta)
    f=np.arange(0,fs,fs/sl)
    f_ind=np.where((f>=fmin)&(f<=fmax))[0]
    f=f[f_ind]
    ome=2*np.pi*f
    B=np.zeros((len(theta),len(f)),np.float64)
    Bs=np.zeros((len(theta),len(f)),np.complex64)
    for jt in range(L):

        X=fft(x[:,jt*(sl-ol):jt*(sl-ol)+sl])
        X=X[:,f_ind[0]:(f_ind[-1]+1)]
        #X=X.astype(np.complex64)

        Bs[:]=0
        for it in range(M):
            phi=np.outer(-1j*D[it]*np.sin(theta/180*np.pi),ome/c)
            Bi=np.exp(phi)*np.tile(X[it,:],(len(theta),1))
            Bs=Bs+Bi
        B=B+np.abs(Bs)**2
    B=B/L
    return theta,f,B


def plot_oancor(f,C,d,fp,fs,c=1500):
    """
    :param freq: frequency list
    :param C: M*M Correlation array (Unit: none)
    :param d: M element depth in ascending order (Unit: m)
    :param fp: appointed frequency (Unit: Hz)
    :param fs: Sampling rate (Unit: Hz)
    :param c: sound velocity (Unit: m/s)
    :return: d(lambda)--correl
    """
    x=(d-d[0])/(c/fp)
    ind=round(fp*len(f)/fs)
    ind=np.where(f>=fp)[0][0]
    y=np.real(C[ind,:,0])
    plt.plot(x,y)
    plt.xlabel(r"d($\lambda$)")
    plt.ylabel("Vertical correlation coefficient")
    plt.title('%dHz' %f[ind])
    plt.show()

def plot_oandi(theta,f,B,fp):
    """
    :param B: directionality array (Unit: none)
    :param fs: Sampling rate (Unit: Hz)
    :param fp: appointed frequency (Unit: Hz)
    :param theta: angle list (Unit: degree)
    :return: degree--direction in frequency fp
    """
    
    x=theta
    ind=np.where(f>=fp)[0][0]
    y=B[:,ind]
    y=10*log10(y)
    plt.plot(x,y)
    plt.xlabel(r"$\theta$ / degree")
    plt.ylabel("10lg(B^2)/dB")
    plt.title('%dHz' %fp)
    plt.show()