#!/usr/bin/env python
# coding: utf-8

##############################################################################
#
# 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.
#
##############################################################################

# # pyPlutoDPD
# Pre-Distortion for Linearization of RF power amplifier in python using ADALM PLUTO.

# ## Overview
# The pyPlutoDPD is a python implementation of Digital Pre-Distortion, a most used technique to linearize the distortions introiduced by the RF Power Amplifiers.
#
# This has been developed under python by:
# - Smail Bachir, University of Poitiers, XLIM-Laboratory UMR-CNRS 7252
#
# in collaboration with:
# - Claude Duvanaud, University of Poitiers, XLIM-Laboratory UMR-CNRS 7252
# - Jean-Marc Ouvrard, Rapid Space
# - Thomas Gambier, Rapid Space
#
# LMS Algorithm running with a polular ADALM-PlutoSDR has been implemneted.
#
# <b> Pre-requises </b> : pyPlutoDPD requires only free dependances :
# - Python with packages : numpy, adi, matplotlib, scipy
# - Install PlutoSDR in Python :  see https://pysdr.org/content/pluto.html
# - Cheap board : Adalm PlutoSDR
#
# <b> Setup </b> : pyPlutoDPD can be used with two configuations
# 1. Minimale configuration with only ADALM-Pluto SDR connected to himself (SMA wire bettween TX and RX). This allows to test the program and see the steps.
# 2. Configuration with DUT (Device Under Test) : Here, insert your RF-PA bettween TX and RX. In our case, we linearize a RF PA of .....

# # Main program

# ### Required packages

# In[47]:


from commonFunctions import *  # XLIM used functions
import numpy as np
import adi  # API for Adalm-Pluto
import matplotlib.pyplot as plt
from scipy.io import loadmat, savemat
from scipy.signal import get_window
from scipy.fft import fft, fftshift
from scipy import signal
import time

start = time.time()    # start time to measure the execution time


# Load the used functions located in the ./commonFunctions.py

# In[48]:


# ## Transmission without DPD

# ### Transmit original samples
# Load transmitted samples en .bin format
# The preambule samples are used for synchronization
# You can use "lte-5G/BW-10MHz with 61.44MHz of sampling rate"  or  "WLAN/BW-10MHz with 60MHz of sampling sate"

# In[49]:


# Read the frame
path = ""
osf = 3
filePreambule = "preambule_BW20osf" + str(osf) + "_single.bin"
fileRx = "lteBW20osf" + str(osf) + "_single.bin"
sample_rate = 61.44e6  # 61.44e6 46.08e6 30.72e6 23.04e6 11.52e6
# fileRx = "wlanBW20osf" + str(osf) + "_single.bin"
# sample_rate = 60e6

fp = open(path + filePreambule, 'rb')
preambule = np.fromfile(fp, dtype=np.complex64)
fp = open(path + fileRx, 'rb')
wave = np.fromfile(fp, dtype=np.complex64)
samples = np.ndarray(len(preambule) + len(wave), dtype=np.complex64)
samples[:len(preambule)] = preambule
samples[len(preambule):] = wave
del preambule, wave

txPower = np.round(10*np.log10(np.mean(np.abs(samples)**2)), 3)
print("TX avg-power = ", txPower, " dBm")  # Print the txPower

num_samps = len(samples)
# The PlutoSDR expects samples between -2^14 and +2^14
coeff = 2**14
samples *= coeff


# Connect the sdr pluto

# In[50]:


sdr = adi.Pluto("ip:192.168.2.1")
sdr.sample_rate = sample_rate


# Config Tx and start transmitting

# In[51]:


freq = int(3650)
win_fft = 'flattop'  # The used window for FFT ans spectrum plot

# Here, the initial gain of TX is set to the minimal value for safety
initialGain = int(0)
sdr.tx_rf_bandwidth = int(sample_rate)  # filter cutoff
sdr.tx_lo = freq*10**6
sdr.tx_hardwaregain_chan0 = initialGain
sdr.tx_destroy_buffer()  # Stop transmitting for safe
sdr.tx_cyclic_buffer = True  # Enable cyclic buffers

safetyGain = float(0)  # set your tx gain limitation according to your DUT
sdr.tx(samples)  # start transmitting


# Loop for tx_gain monitoring in real time

# In[52]:


valid = False
while (valid):
    valid = input("1 to set txGain  ||  0 to compute DPD : ")
    valid = bool(int(valid))
    if valid == True:
        value = input(
            f"enter new txGain (Actual txGain = {sdr.tx_hardwaregain_chan0}) : ")
        gain = float(value)
        if gain < safetyGain:  # Gain limitation for safety
            sdr.tx_hardwaregain_chan0 = gain
        else:
            sdr.tx_hardwaregain_chan0 = safetyGain


# ### Receive the original samples

# In[53]:


sdr.rx_lo = sdr.tx_lo
sdr.rx_rf_bandwidth = sdr.tx_rf_bandwidth
sdr.rx_buffer_size = 3*num_samps  # 3 times to capture one sequence
sdr.gain_control_mode_chan0 = 'manual'
# set the receive gain, but be careful not to saturate the ADC
sdr.rx_hardwaregain_chan0 = 8


# Clear buffer just to be safe

# In[54]:


for i in range(0, 10):
    raw_data = sdr.rx()


# Receive samples

# In[55]:


rx_samples = sdr.rx()
avg_pwr = np.mean(np.abs(rx_samples/coeff)**2)
avg_pwr_dB = round(10*np.log10(avg_pwr), 2)
print("check_rx = ", avg_pwr_dB,
      ": Best in -30 +/- 5. Else, adjust the rx_hardwaregain")


# Synchronization and coarse phase correction using the preambule (zeroing AM/PM)

# In[56]:


# Synchronization using cross-corelation
[paIn, paOut] = corelation(samples, rx_samples, osf)


# In[57]:


# Coarse phase correction using the preambule (zeroing AM/PM)
paOut = coarsePhaseCorrection(paIn, paOut, osf)
paIn /= np.max(np.abs(paIn))  # scaling
paOut /= np.max(np.abs(paOut))  # scaling


# Plot freq domain

# In[58]:


[pxxWithout, f] = toPlotSpectrum(paOut, sample_rate, win_fft, 20, "all")
plt.figure(10)
plt.plot((f + freq)/1e6, 10*np.log10(pxxWithout))
plt.grid()
plt.xlabel('Frequency (MHz)')
plt.ylabel('Spectrum')
plt.legend(['Original signal'])
plt.show()


# Plot AM/AM and AM/PM

# In[59]:


amam_ampm_plot(paIn, paOut)
plt.show()


# ## DPD Computation

# $$
#  {y_{dpd}}(n) = \displaystyle \sum_{k=0}^{K_a-1} \sum_{l=0}^{L_a-1} a_{k l} \, x(n-l) \cdot |x(n-l)|^{k} + \sum_{k=1}^{K_b} \sum_{l=0}^{L_b-1} \sum_{m=1}^{M_b} b_{k l m} \, x(n-l) \cdot |x(n-l-m)|^{k} + \sum_{k=1}^{K_c} \sum_{l=0}^{L_c-1} \sum_{m=1}^{M_c} c_{k l m} \, x(n-l) \cdot |x(n-l+m)|^{k}
# $$
#
# $x(n)$ : dpd input (complex envelope) <br />
# $y_{dpd}(n)$ : dpd output (complex envelope)

# In[37]:


dpdModel = {'fs': sample_rate}
# set the dpd orders and memory depths
[Ka, La, Kb, Lb, Mb, Kc, Lc, Mc] = [7, 1, 0, 1, 1, 0, 1, 1]
dpdModel['parameters'] = {'Ka': Ka, 'La': La, 'Kb': Kb,
                          'Lb': Lb, 'Mb': Mb, 'Kc': Kc, 'Lc': Lc, 'Mc': Mc}

gainATT = 1

dpdModel = identif_dpd_GMP(paIn, paOut/gainATT, dpdModel)
print("nmse of identification :", np.round(dpdModel['nmseTimedB'], 2), " dB")

dpdOutput = simGMPopt(paIn, dpdModel)  # generate DPD signal
dpdOutput = dpdOutput/np.max(np.abs(dpdOutput))
txPower = np.round(10*np.log10(np.mean(np.abs(dpdOutput)**2)), 3)
print("dpd tx-power = ", txPower, " dBm")

dpdOutput *= coeff  # scaling for ADALM-Pluto


# ## Transmission with DPD

# ### Transmit DPD signal

# In[38]:


sdr.tx_destroy_buffer()  # Stop transmitting for safe
sdr.tx_cyclic_buffer = True  # Enable cyclic buffers

sdr.tx(dpdOutput)  # start transmitting with DPD

valid = False
while (valid):
    valid = input("1 to set txGain  ||  0 to compute DPD : ")
    valid = bool(int(valid))
    if valid == True:
        value = input(
            f"enter new txGain (Actual txGain = {sdr.tx_hardwaregain_chan0}) : ")
        gain = float(value)
        if gain < safetyGain:  # Gain limitation for safety
            sdr.tx_hardwaregain_chan0 = gain
        else:
            sdr.tx_hardwaregain_chan0 = safetyGain

# savemat('36dBm.mat', {'paInput':paIn, 'paOutput':paOut, 'dpdOutput':dpdOutput, 'dpdModel':dpdModel})
# np.savetxt('outfile.txt', rx_samples.view(complex))


# ### Receive with DPD signal

# In[39]:


rx_samples = sdr.rx()
avg_pwr = np.mean(np.abs(rx_samples/coeff)**2)
avg_pwr_dB = round(10*np.log10(avg_pwr), 2)
print("check_rx = ", avg_pwr_dB,
      ": Best in -30 +/- 5. Else, adjust the rx_hardwaregain")


# Synchronization of input and output

# In[40]:


[paIn, paLin] = corelation(samples, rx_samples, osf)
paLin = coarsePhaseCorrection(paIn, paLin, osf)

paIn /= np.max(np.abs(paIn))  # scaling
paLin /= np.max(np.abs(paLin))  # scaling


# Plot Linearized AM/AM and AM/PM

# In[41]:


amam_ampm_plot(paIn, paLin)
plt.show()
print("nmse after linearization : ", np.round(nmse(paIn, paLin), 2), " dB")


# In[42]:


plt.figure()
plt.subplot(221)
plt.plot(np.real(paIn))
plt.plot(np.real(paLin))
plt.title("I/O in phase")
plt.xlabel("Samples")
plt.subplot(222)
plt.plot(np.imag(paIn))
plt.plot(np.imag(paLin))
plt.title("I/O in quadrature")
plt.xlabel("Samples")
plt.subplot(223)
plt.plot(np.real(paIn)-np.real(paLin))
plt.title("Error in phase (I)")
plt.xlabel("Samples")
plt.subplot(224)
plt.plot(np.imag(paIn)-np.imag(paLin))
plt.title("Error in qudadrature (Q)")
plt.xlabel("Samples")
plt.show()


# Compare spectrum with and without DPD

# In[43]:


[pxxWith, f] = toPlotSpectrum(paLin, sample_rate, win_fft, 20, "all")
plt.figure(10)
plt.plot((f + freq)/1e6, 10*np.log10(pxxWith))
plt.grid()
plt.xlabel('Frequency (MHz)')
plt.ylabel('Spectrum')
plt.legend(['With DPD'])
plt.show()


# Stop transmitting

# In[44]:


sdr.tx_destroy_buffer()


# Print the DPD coefficients

# In[45]:


end = time.time()
elapsed = end - start
print(f'Elapsed time : {elapsed:.2}ms')


# In[46]:


for i in dpdModel['coefGMP']:
    print(np.round(i, 3))
