##############################################################################
#
# Copyright (C) 2024 XLIM-Laboratory UMR-CNRS, Rapid.Space International
#
# This program is free software: you can Use, Study, Modify and Redistribute
# it under the terms of the GNU General Public License version 3, or (at your
# option) any later version, as published by the Free Software Foundation.
# 
# You can also Link and Combine this program with other software covered by
# the terms of any of the Free Software licenses or any of the Open Source
# Initiative approved licenses and Convey the resulting work. Corresponding
# source of such a combination shall include the source code for all other
# software used.
# 
# This program is distributed WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# 
# See COPYING file for full licensing terms.
#
##############################################################################

import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import get_window
from scipy.fft import fft, fftshift
from scipy import signal


def nmse(xMes, xEst):
    """
    Function to compute the NMSE in time domain
    """
    error = xMes - xEst
    nmse_dB = 10 * np.log10(np.sum(np.abs(error)**2) / np.sum(np.abs(xMes)**2))
    return nmse_dB

######################################


def algoLMS(PPHI, out, start, nSample):
    """
    Classical LMS Algorithm
    gmpModel['coefGMP'] = np.linalg.inv((np.dot(np.conj(PPHI.T), PPHI))) @ np.conj(PPHI.T) @ out
    """

    """
    LMS Algorithm with regularization
    """
    PHI = PPHI[start:nSample, :]
    U, S, V = np.linalg.svd(np.conj(PHI.T) @ PHI)
    Splus = np.diag(1 / np.where(S > 1e-8, S, np.inf))
    theta = (U @ Splus @ np.conj(U.T)) @ np.conj(PHI.T) @ out[start:nSample]
    H = (U @ Splus @ np.conj(U.T)) @ np.conj(PHI.T)

    return theta, H

########################################


def computeFFT(x, fs, win):
    """
    Perform FFT on x 
    x    : signal
    fx   : sampling rate
    win  : fft windows

    returns
    fftAbs : fft absolute value
    freq : frequency array
    """

    N = len(x)
    fft_block_size = N
    w = get_window(win, fft_block_size)
    fftSig = fft(x * w) / fft_block_size
    fftAbs = np.abs(fftSig)
    freq = (np.arange(0, N) - N/2) * fs / fft_block_size
    # Shift zero frequency component to center of spectrum
    freq = fftshift(freq)
    return fftAbs, freq

#########################################


def toPlotSpectrum(x, fs, win, Nblocks, interval):
    """
    To plot the spectrum
    x    : signal
    fx   : sampling rate
    win  : fft windows
    Nbblocks : number of blocks to perform average spectrum
    interval = "half" to spectrum in [0 fs/2] else spectrum in [-fs/2 fs/2]

    returns
    pxx : average fft absolute value
    freq : frequency array
    """

    # Spectrum computation (one or two sided)
    N = len(x)
    fft_block_size = int(2**(np.ceil(np.log2(N/Nblocks))-1))
    window = get_window(win, fft_block_size)
    fft_input = np.zeros(fft_block_size, dtype='complex128')

    for n in range(1, Nblocks+1):
        idx = slice((n-1)*fft_block_size, n*fft_block_size)
        fft_block = fft(x[idx] * window) / fft_block_size
        fft_input += (fft_block * np.conj(fft_block))
    fft_input /= Nblocks
    if interval == "half":  # spectre entre [0 fs/2]
        fft_input = fft_input[:fft_block_size//2]
        pxx = np.abs(fft_input)
        freq = np.arange(0, fs/2, fs/(2*fft_block_size))
    else:  # spectre entre [-fs/2 fs/2]
        pxx = np.abs(fft_input)
        pxx = fftshift(pxx)
        freq = (np.arange(0, fft_block_size) -
                fft_block_size/2) * fs / fft_block_size
    return pxx, freq


#################################################

def corelation(tx, rx, osf):
    """
    Performs tx/rx synchronization and signal alignement
    tx    : transmitted signal
    rx    : received signal
    osf   : over-sampling factor

    returns
    txSig, rxSig  : synchonized tx and rx signals
    """

    preambule = 10*80  # 2 STF + 2 LTF + 1 SIG + 5 another
    n = int(osf*preambule)

    # Resample transmit waveform
    rap = 16  # Another upSampling for best synchronization
    tx = signal.resample(tx, len(tx) * rap)
    rx = signal.resample(rx, len(rx) * rap)
    nSamples = len(tx)
    lag = xcorr(np.abs(rx[:nSamples]), np.abs(tx[:n*rap]))

    # Data alignement and downSampling
    start = 0
    txSig = tx[start:len(tx)]
    rxSig = rx[lag+start:lag + start + len(txSig)]
    txSig = signal.resample(txSig, len(txSig) // rap)
    rxSig = signal.resample(rxSig, len(rxSig) // rap)

    return txSig, rxSig

#################################################


def xcorr(x, y):
    """
    Perform Cross-Correlation on x and y
    x    : 1st signal
    y    : 2nd signal

    returns
    lags : lags of correlation
    corr : coefficients of correlation
    """
    corr = signal.correlate(x, y, mode="full")
    lags = signal.correlation_lags(len(x), len(y), mode="full")
    lag = lags[np.argmax(corr)]

    return lag

#################################################


def coarsePhaseCorrection(x, y, osf):

    # Coarse phase correction using the preambule (zeroing AM/PM)

    preambule = 5*80  # 2 STF + 2 LTF + 1 SIG
    n = int(osf*preambule)
    nSample = len(x)

    dephase = np.angle(np.conj(x[:n].T) @ y[:n])
    y = y * np.exp(-1j * dephase)

    return y

##################################################


def amam_ampm_plot(tx, ty):
    """
    Perform AM/AM and AM/PM curves 
    ttx    : transmitted signal
    rx    : received signal

    """

    plt.figure()
    plt.subplot(211)
    plt.plot(np.abs(tx), np.abs(ty), '.')
    plt.title('Dynamical AM/AM')
    plt.xlabel('Normalized input')
    plt.ylabel('Normalized output')
    plt.subplot(212)
    plt.plot(np.abs(tx), np.angle(np.conj(tx) * ty), '.')
    plt.title('Dynamical AM/PM')
    plt.xlabel('Normalized input')
    plt.ylabel('Dephase (deg)')

###################################################


def identif_dpd_GMP(x, y, dpdModel):

    dpdModel = estimGMPopt(y, x, dpdModel)
    dpdModel['nmseTimedB'] = nmse(x, dpdModel['outEst'])

    return dpdModel

#################################################


def estimGMPopt(x, y, gmpModel):
    """
    To estimate GMP coefficients 
    x    : reference signal
    y    : measured signal
    gmpModel  : object containing model parameters

    returns
    gmpModel : gmpModel with the estimated coefficients
    """

    Ka = gmpModel['parameters']['Ka']
    La = gmpModel['parameters']['La']
    Kb = gmpModel['parameters']['Kb']
    Lb = gmpModel['parameters']['Lb']
    Mb = gmpModel['parameters']['Mb']
    Kc = gmpModel['parameters']['Kc']
    Lc = gmpModel['parameters']['Lc']
    Mc = gmpModel['parameters']['Mc']
    gmpModel['nbCoef'] = Ka * La + Kb * Lb * Mb + Kc * Lc * Mc
    nSample = len(x)
    PPHI = np.zeros((nSample, gmpModel['nbCoef']), dtype=np.complex128)

    for l in range(La):
        it = np.concatenate((np.zeros(l), x[:nSample-l]))
        for k in range(Ka):
            PPHI[:, (La * k) + l] = (np.abs(it))**k * it

    for l in range(Lb):
        it2 = np.concatenate((np.zeros(l), x[:nSample-l]))
        for m in range(1, Mb + 1):
            it1 = np.concatenate((np.zeros(l + m), x[:nSample-l-m]))
            for k in range(1, Kb + 1):
                PPHI[:, La * Ka + (Lb * Mb * (k - 1)) +
                     (Mb * l) + (m-1)] = (np.abs(it1))**k * it2

    for l in range(Lc):
        it2 = np.concatenate((np.zeros(l), x[:nSample-l]))
        for m in range(1, Mc + 1):
            if l >= m:
                it1 = np.concatenate((np.zeros(l - m), x[:nSample-l+m]))
            else:
                it1 = np.concatenate((np.zeros(m - l), x[m-l:nSample])) * 0
            for k in range(1, Kc + 1):
                PPHI[:, La * Ka + Lb * Kb * Mb +
                     (Lc * Mc * (k - 1)) + (Mc * l) + (m-1)] = (np.abs(it1))**k * it2

    start = 0
    gmpModel['coefGMP'], gmpModel['H'] = algoLMS(PPHI, y, start, nSample)
    gmpModel['outEst'] = PPHI @ gmpModel['coefGMP']

    return gmpModel

######################################################


def simGMPopt(x, gmpModel):
    """
    Simulates GMP model 
    x : input signal (stimulus)
    gmpModel  : object containing model parameters

    returns
    outEst : simulated output    
    """

    Ka = gmpModel['parameters']['Ka']
    La = gmpModel['parameters']['La']
    Kb = gmpModel['parameters']['Kb']
    Lb = gmpModel['parameters']['Lb']
    Mb = gmpModel['parameters']['Mb']
    Kc = gmpModel['parameters']['Kc']
    Lc = gmpModel['parameters']['Lc']
    Mc = gmpModel['parameters']['Mc']

    nSample = len(x)
    PPHI = np.zeros((nSample, gmpModel['nbCoef']), dtype=np.complex128)

    for l in range(La):
        it = np.concatenate((np.zeros(l), x[:nSample-l]))
        for k in range(Ka):
            PPHI[:, (La * k) + l] = (np.abs(it))**k * it

    for l in range(Lb):
        it2 = np.concatenate((np.zeros(l), x[:nSample-l]))
        for m in range(1, Mb + 1):
            it1 = np.concatenate((np.zeros(l + m), x[:nSample-l-m]))
            for k in range(1, Kb + 1):
                PPHI[:, La * Ka + (Lb * Mb * (k - 1)) +
                     (Mb * l) + (m-1)] = (np.abs(it1))**k * it2

    for l in range(Lc):
        it2 = np.concatenate((np.zeros(l), x[:nSample-l]))
        for m in range(1, Mc + 1):
            if l >= m:
                it1 = np.concatenate((np.zeros(l - m), x[:nSample-l+m]))
            else:
                it1 = np.concatenate((np.zeros(m - l), x[m-l:nSample])) * 0
            for k in range(1, Kc + 1):
                PPHI[:, La * Ka + Lb * Kb * Mb +
                     (Lc * Mc * (k - 1)) + (Mc * l) + (m-1)] = (np.abs(it1))**k * it2

    outEst = np.zeros(nSample, dtype=np.complex128)
    outEst = PPHI @ gmpModel['coefGMP']

    return outEst

##################################################
