%matplotlib inline

import math

import numpy as np
from matplotlib import pyplot as plt

from pyphysim.modulators.fundamental import BPSK, QAM, Modulator
from pyphysim.simulations import Result, SimulationResults, SimulationRunner
from pyphysim.util.conversion import dB2Linear, linear2dB
from pyphysim.util.misc import pretty_time, randn_c, count_bit_errors

np.set_printoptions(precision=2, linewidth=120)

"""
Explanation: Simulating transmission with Rayleigh and AWGN channels and different constellation sizes
This notebook illustrates how to perform a simulation of a transmission through Rayleigh channel using different modulations. 
Furthermore, it describe how smulations can be run in parallel.
This is a modification of the Transmission_with_AWGN_channel notebook.
End of explanation
"""


bpsk = BPSK()
qam4 = QAM(4)
qam16 = QAM(16)
qam64 = QAM(64)
qam256 = QAM(256)

fig_bpsk, ax_bpsk = plt.subplots(figsize=(4, 4))
ax_bpsk.set_title("BPSK")
ax_bpsk.plot(bpsk.symbols.real, bpsk.symbols.imag, "*r", label="4-QAM")
ax_bpsk.axis("equal")

fig, [[ax11, ax12], [ax21, ax22]] = plt.subplots(figsize=(8, 8),
                                                 nrows=2,
                                                 ncols=2)

ax11.set_title("4-QAM")
ax11.plot(qam4.symbols.real, qam4.symbols.imag, "*r", label="4-QAM")
ax11.axis("equal")

ax12.set_title("16-QAM")
ax12.plot(qam16.symbols.real, qam16.symbols.imag, "*r", label="16-QPSK")
ax12.axis("equal")

ax21.set_title("64-QAM")
ax21.plot(qam64.symbols.real, qam64.symbols.imag, "*r", label="64-QAM")
ax21.axis("equal")

ax22.set_title("256-QAM")
ax22.plot(qam256.symbols.real, qam256.symbols.imag, "*r", label="256-QAM")
ax22.axis("equal");

"""
Explanation: Visualize the constellation
First let's visualize the different constellations that will be employed.
For $M=2$, BPSK modulation will be used, with symbols being $-1$ and $1$. For other sizes QAM modulation will be used.
End of explanation
"""


%%writefile rayleigh_qam_simulator.py

import math

import numpy as np

from pyphysim.modulators.fundamental import BPSK, QAM, Modulator
from pyphysim.simulations import Result, SimulationResults, SimulationRunner
from pyphysim.util.conversion import dB2Linear
from pyphysim.util.misc import randn_c, count_bit_errors


class RayleighOrAwgnSimulator(SimulationRunner):

    def __init__(self, SINR_dB_values, simulate_with_rayleigh=False):
        # If simulate_with_rayleigh is false only AWGN is used
        super().__init__()

        self._simulate_with_rayleigh = simulate_with_rayleigh

        # Add the simulation parameters to the `params` attribute.
        self.params.add('EbN0_db', SINR_dB_values)
        self.params.set_unpack_parameter('EbN0_db')

        # Note that when M=2 BPSK modulation will be used, while other values will use QAM
        self.params.add("M", [2, 4, 16, 64, 256])
        self.params.set_unpack_parameter('M')

        # Number of times the `_run_simulation` method will run when `simulate` method is called.
        # We are using a value 100 times larger than before, but the simulation will not take
        # 100 times the previous elapsed time to finish thanks to the implementation of the
        # `_keep_going` method that will allow us to skip many of these iterations for low EbN0 values
        self.rep_max = 50000

        # Number of symbols generated for this realization
        self.num_symbols = 1000

        # Used in the implementation of `_keep_going` method. This is the maximum numbers of symbol
        # errors we allow before `_run_simulation` is stoped for a given configuration
        self.max_symbol_errors = 1. / 300. * self.num_symbols * self.rep_max

        # Set a nice message for the progressbar
        self.progressbar_message = f"Simulating for {self.params.get_num_unpacked_variations()} configurations"

        self.update_progress_function_style = "text1"  # "ipython"

    def _keep_going(self, current_params, current_sim_results, current_rep):
        # Note that we have added a "symbol_errors" result in `_run_simulation` to use here

        # Get the last value in the "symbol_errors" results list, which corresponds to the current configuration
        cumulated_symbol_errors \
            = current_sim_results['symbol_errors'][-1].get_result()
        return cumulated_symbol_errors < self.max_symbol_errors

    def _run_simulation(self, current_parameters):
        # Since EbN0_db is an "unpacked parameter" a single value is passed to `_run_simulation`.
        # We can get the current value as below
        sinr_dB = current_parameters['EbN0_db']
        M = current_parameters['M']

        modulator = BPSK() if M == 2 else QAM(M)

        # Find the noise power from the EbN0 value (in dB)
        EbN0_linear = dB2Linear(sinr_dB)
        snr_linear = EbN0_linear * math.log2(M)
        noise_power = 1 / snr_linear

        # Generate random transmit data and modulate it
        data = np.random.randint(0, modulator.M, size=self.num_symbols)
        modulated_data = modulator.modulate(data)

        # Noise vector
        n = math.sqrt(noise_power) * randn_c(self.num_symbols)

        if self._simulate_with_rayleigh:
            # Rayleigh channel
            h = randn_c(modulated_data.size)

            # Receive the corrupted data
            received_data = h * modulated_data + n

            # Equalization
            received_data /= h

        else:
            # Receive the corrupted data
            received_data = modulated_data + n

        # Demodulate the received data and compute the number of symbol errors
        demodulated_data = modulator.demodulate(received_data)
        symbol_errors = sum(demodulated_data != data)

        num_bit_errors = count_bit_errors(data, demodulated_data)

        # Create a SimulationResults object and save the symbol error rate.
        # Note that the symbol error rate is given by the number of symbol errors divided by the number of
        # transmited symbols. We want to combine the symbol error rate for the many calls of `_run_simulation`.
        # Thus, we choose `Result.RATIOTYPE` as the "update_type". See the documentation of the `Result` class
        # for more about it.
        simResults = SimulationResults()
        simResults.add_new_result("symbol_error_rate",
                                  Result.RATIOTYPE,
                                  value=symbol_errors,
                                  total=self.num_symbols)
        simResults.add_new_result("symbol_errors",
                                  Result.SUMTYPE,
                                  value=symbol_errors)
        simResults.add_new_result("num_bit_errors",
                                  Result.SUMTYPE,
                                  value=num_bit_errors)
        simResults.add_new_result("bit_error_rate",
                                  Result.RATIOTYPE,
                                  value=num_bit_errors,
                                  total=int(np.log2(modulator.M)) * data.size)
        return simResults

"""
Explanation: Implement a simulator
In this notebook we will use the run_in_parallel method to run the simulation, which requires our simulator to be importable (the reason is that it needs to work with pickle, which is not possible if the class is define interactively in the notebook). For that, we add the code in the cell below and used the %%writefile cell magic to save the content of the cell to a file. The cell content is NOT executed. After that, we can import from the file.
End of explanation
"""


from rayleigh_qam_simulator import RayleighOrAwgnSimulator

"""
Explanation: Running the simulation
Now we import the simulator with the line below
End of explanation
"""


from ipyparallel import Client
import sys
import os

c = Client()
# Get a "direct view" -> we will use it to run some code in all engines before we start the simulation
# Partucularly, we want to clear the engine's workspace and to append the current working directory to each engine's path
dview = c.direct_view()  # Get a view of all engines in the cluster
# This will execute the command in ALL engines
dview.execute('%reset')
dview.execute('import sys')
dview.execute('sys.path.append("{0}")'.format(os.getcwd()), block=True)

# Get a load balanced view, which is the one we will pass to the `simulate_in_parallel` method.
lview = c.load_balanced_view()

"""
Explanation: The parallel computation os powered by the ipyparallel library. Before you can run the simulation in parallel you must start the ipython cluster and engines. This can be done by running the ipcluster start command in a separated terminal.
After that, in our python code we can import from the ipyparallel library and create a Client. From the client we get a "view" of the engines that can be used to run code in parallel. This view is passed to the simulate_in_parallel method of our simulator.
The code below illustrates how to create a client and how to get a view. Note that each engine corresponds to a different python process (possible running in a different machine) and the working directory in the engines is where the cluster was started and not necessarelly our working directory. Assuming you started the engines in a shell with the virtual environment with pyphysim installed everything in pyphysim should be importable in the engines. However, code in the same folder of this notebook might not be importable in the engines. To solve that we add the current directory of this notebook in each engine's path.
End of explanation
"""


EbN0_db_rayleigh = np.linspace(-5, 30, 8)

runner_rayleigh = RayleighOrAwgnSimulator(EbN0_db_rayleigh,
                                          simulate_with_rayleigh=True)
runner_rayleigh.set_results_filename("qam_rayleigh")
print("Simulation parameters:", runner_rayleigh.params)

runner_rayleigh.simulate_in_parallel(lview)
# runner_rayleigh.simulate()
print(runner_rayleigh.elapsed_time)

"""
Explanation: Now we can run the simulation. Let's create an RayleighSimulator and call the simulate_in_parallel method.
End of explanation
"""


EbN0_db_awgn = np.linspace(-5, 30, 8)

runner_awgn = RayleighOrAwgnSimulator(EbN0_db_awgn,
                                      simulate_with_rayleigh=False)
runner_awgn.set_results_filename("qam_awgn")
print("Simulation parameters:", runner_awgn.params)

runner_awgn.simulate_in_parallel()
# runner_awgn.simulate()
print(runner_awgn.elapsed_time)

"""
Explanation: Now let's simulate transmission in an AWGN channel. However, this time we will not pass an ipyparallel view. If a view is not passed, then one will be created inside the simulate_in_parallel method for us like the one we have previously defined. Usually that is enough, but it is useful to know how to pass a view manually as before if, for instance, you want to use a different profile (we have not specified one when we created the client and the default profile was used) or when you want a view that does not include all started engines.
End of explanation
"""


symbol_error_rate_bpsk_rayleigh = runner_rayleigh.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 2})
symbol_error_rate_qam4_rayleigh = runner_rayleigh.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 4})
symbol_error_rate_qam16_rayleigh = runner_rayleigh.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 16})
symbol_error_rate_qam64_rayleigh = runner_rayleigh.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 64})
symbol_error_rate_qam256_rayleigh = runner_rayleigh.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 256})

print("Symbol Errors Rayleigh Channel (BPSK):\n",
      symbol_error_rate_bpsk_rayleigh)
print("Symbol Errors Rayleigh Channel (4-QAM):\n",
      symbol_error_rate_qam4_rayleigh)
print("Symbol Errors Rayleigh Channel (16-QAM):\n",
      symbol_error_rate_qam16_rayleigh)
print("Symbol Errors Rayleigh Channel (64-QAM):\n",
      symbol_error_rate_qam64_rayleigh)
print("Symbol Errors Rayleigh Channel (256-QAM):\n",
      symbol_error_rate_qam256_rayleigh)

symbol_error_rate_bpsk_awgn = runner_awgn.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 2})
symbol_error_rate_qam4_awgn = runner_awgn.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 4})
symbol_error_rate_qam16_awgn = runner_awgn.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 16})
symbol_error_rate_qam64_awgn = runner_awgn.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 64})
symbol_error_rate_qam256_awgn = runner_awgn.results.get_result_values_list(
    "symbol_error_rate", fixed_params={"M": 256})

print("Symbol Errors AWGN Channel (BPSK):\n", symbol_error_rate_qam4_awgn)
print("Symbol Errors AWGN Channel (4-QAM):\n", symbol_error_rate_qam4_awgn)
print("Symbol Errors AWGN Channel (16-QAM):\n", symbol_error_rate_qam16_awgn)
print("Symbol Errors AWGN Channel (64-QAM):\n", symbol_error_rate_qam64_awgn)
print("Symbol Errors AWGN Channel (256-QAM):\n", symbol_error_rate_qam256_awgn)

"""
Explanation: If you run this notebook you might notice that the simulate_in_parallel method of the runner_rayleigh runner blocks and only when it finishes is that the simulate_in_parallel method of the runner_awgn runner in the subsequent cell can start. That means that at some point when only a few (or only one) variations of the parameters for runner_rayleigh runner are missing not all processors in your computer will be used.
When you want to simulate in parallel using multiple runners as in this notebook, you might want to pass the optional wait argument as False. In that case, you need to manually call the wait_parallel_simulation method of both simulators at some point.
Visualizing the results
To get the obtained results we can use the runner.results.get_result_values_list as before. However, since we have two parameters that were unpacked we will get the results of all possible combinations. We likely want to specify which subset we want. This can be done by passing the fixed_params argument to the get_result_values_list with a dictionary indicating the fixed values as below.
End of explanation
"""


# Now let's plot the results
fig, [ax1, ax2] = plt.subplots(figsize=(15, 6), ncols=2)

ax1.semilogy(EbN0_db_rayleigh,
             symbol_error_rate_bpsk_rayleigh,
             "-",
             color="red",
             label="BPSK Simulation")
ax1.semilogy(EbN0_db_rayleigh,
             symbol_error_rate_qam4_rayleigh,
             "-",
             color="orange",
             label="4-QAM Simulation")
ax1.semilogy(EbN0_db_rayleigh,
             symbol_error_rate_qam16_rayleigh,
             "-",
             color="blue",
             label="16-QAM Simulation")
ax1.semilogy(EbN0_db_rayleigh,
             symbol_error_rate_qam64_rayleigh,
             "-",
             color="lightgreen",
             label="64-QAM Simulation")
ax1.semilogy(EbN0_db_rayleigh,
             symbol_error_rate_qam256_rayleigh,
             "-",
             color="magenta",
             label="256-QAM Simulation")
# ax1.semilogy(EbN0_db, QAM(256).calcTheoreticalSER(EbN0_db), color="magenta", label="M=256 Theory")
ax1.set_title("Symbol Error Rate (Rayleigh channel)")
ax1.set_ylabel("Symbol Error Rate")
ax1.set_xlabel("EbN0 (dB)")
ax1.grid(True, which="both")
ax1.set_ylim((1e-4, 1e0))
ax1.legend()

ax2.semilogy(EbN0_db_awgn,
             symbol_error_rate_bpsk_awgn,
             "o",
             color="red",
             label="BPSK Simulation")
ax2.semilogy(EbN0_db_awgn,
             BPSK().calcTheoreticalSER(EbN0_db_awgn),
             color="red",
             label="BPSK Theory")
ax2.semilogy(EbN0_db_awgn,
             symbol_error_rate_qam4_awgn,
             "o",
             color="orange",
             label="M=4 Simulation")
ax2.semilogy(EbN0_db_awgn,
             QAM(4).calcTheoreticalSER(EbN0_db_awgn + linear2dB(2)),
             color="orange",
             label="4-QAM Theory")
ax2.semilogy(EbN0_db_awgn,
             symbol_error_rate_qam16_awgn,
             "o",
             color="blue",
             label="16-QAM Simulation")
ax2.semilogy(EbN0_db_awgn,
             QAM(16).calcTheoreticalSER(EbN0_db_awgn + linear2dB(4)),
             color="blue",
             label="16-QAM Theory")
ax2.semilogy(EbN0_db_awgn,
             symbol_error_rate_qam64_awgn,
             "o",
             color="lightgreen",
             label="64-QAM Simulation")
ax2.semilogy(EbN0_db_awgn,
             QAM(64).calcTheoreticalSER(EbN0_db_awgn + linear2dB(6)),
             color="lightgreen",
             label="64-QAM Theory")
ax2.semilogy(EbN0_db_awgn,
             symbol_error_rate_qam256_awgn,
             "o",
             color="magenta",
             label="256-QAM Simulation")
ax2.semilogy(EbN0_db_awgn,
             QAM(256).calcTheoreticalSER(EbN0_db_awgn + linear2dB(8)),
             color="magenta",
             label="256-QAM Theory")
ax2.set_title("Symbol Error Rate (AWGN channel)")
ax2.set_ylabel("Symbol Error Rate")
ax2.set_xlabel("EbN0 (dB)")
ax2.grid(True, which="both")
ax2.set_ylim((1e-5, 1e0))
ax2.legend();

"""
Explanation: Now we can finally see the plot and compare the simulated results with theoretical values.
End of explanation
"""


