# qosst-sim - Simulation module of the Quantum Open Software for Secure Transmissions.
# Copyright (C) 2021-2024 Mayeul Chavanne
# Copyright (C) 2021-2024 Yoann Piétri

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""
Example script for qosst-sim.
"""

import time
from typing import List

import matplotlib.pyplot as plt

from qosst_sim.modulation.gaussian_qam import GaussianQAM
from qosst_sim.detector import NoisyHeterodyneDetector
from qosst_sim.channel import GaussianChannel
from qosst_sim.simulator.finite_size_simulator import FiniteSizeSimulator
from qosst_sim.simulator.gaussian_channel_asymptotic_calculator import (
    GaussianChannelAsymptoticCalculator,
)


def linear_range(start: float, stop: float, num_points: int) -> List[float]:
    """
    Return a linear range from start to stop with num_points.

    Args:
        start (float): start point of the linear range.
        stop (float): end point of the linear range.
        num_points (int): number of points in the range

    Returns:
        List[float]: range.
    """
    step = (stop - start) / num_points
    return [start + x * step for x in range(num_points + 1)]


def transmission(distance: float) -> float:
    """
    Return the transmittance in a fiber with an attenuation coefficient of 0.2dB/km.

    Args:
        distance (float): distance in km.

    Returns:
        float: transmittance in a fiber at 0.2dB/km.
    """
    return 10 ** (-0.02 * distance)


# pylint: disable=too-many-locals
def main():
    """
    Main function for the example.
    """
    varying_parameter = "Distance (km)"
    varying_range = linear_range(0, 40, 4 * 40)
    beta = 0.95
    dim = 105
    modulation_size = 8
    variance = 5
    nu = 0.0746269
    num_symbols = 500000
    xi_bob = 0.02
    electronic_noise = 0.01
    eta = 0.65

    simulated_skr = []
    asymptotic_skr = []

    print("Start")
    time1 = time.time()

    modulation = GaussianQAM(dim, modulation_size, variance, nu)
    label = (
        " PCS "
        + str(modulation_size**2)
        + "-QAM, dim ="
        + str(dim)
        + ", nu = "
        + str(nu)
    )
    detector = NoisyHeterodyneDetector(eta, electronic_noise)
    label += (
        " (noisy detector eta =" + str(eta) + " Vel = " + str(electronic_noise) + ")"
    )

    for distance in varying_range:
        # initialize the channel of the desired type
        transmittance = transmission(distance)
        channel = GaussianChannel(
            transmittance, xi_bob / (transmittance * detector.eta)
        )

        # initialize the simulator of the desired type
        simulator = FiniteSizeSimulator(
            modulation, channel, detector, num_symbols, beta
        )
        calculator = GaussianChannelAsymptoticCalculator(
            modulation, channel, detector, beta
        )

        print(varying_parameter, distance)
        print(
            "Finite size simulated nB =",
            simulator.n_B,
            "VS asymptotic nB =",
            calculator.n_B,
        )
        print(
            "Finite size simulated c1 =",
            simulator.c1,
            "VS asymptotic c1 =",
            calculator.c1,
        )
        print(
            "Finite size simulated c2 =",
            simulator.c2,
            "VS asymptotic c2 =",
            calculator.c2,
        )
        current_simulated_skr = simulator.skr()
        current_asymptotic_skr = calculator.skr()
        print(
            "Finite size simulated SKR =",
            current_simulated_skr,
            "VS asymptotic SKR",
            current_asymptotic_skr,
        )
        simulated_skr.append(current_simulated_skr)
        asymptotic_skr.append(current_asymptotic_skr)

    time2 = time.time()
    print("Stop")
    print("computation time =", time2 - time1)

    # plotting tools
    _, axes = plt.subplots(figsize=(12, 6))  # plt.subplots(figsize=(12, 6))

    for side in axes.spines.keys():  # 'top', 'bottom', 'left', 'right'
        axes.spines[side].set_linewidth(1)

    # plotting of the SKR
    axes.plot(
        varying_range,
        simulated_skr,
        color="r",
        linestyle="-",
        linewidth=1,
        label=str(num_symbols) + "-symbols simulation," + label,
    )
    axes.plot(
        varying_range,
        asymptotic_skr,
        color="b",
        linestyle="--",
        linewidth=1,
        label="Asymptotic calculation" + label,
    )

    title = "SKR (VA = " + str(variance) + ", xi_bob = " + str(xi_bob) + ")"

    plt.xlabel(varying_parameter, fontsize=6, fontweight="bold")
    plt.ylabel("SKR", fontsize=6, fontweight="bold")
    plt.xticks(fontsize=6, fontweight="bold")
    plt.yticks(fontsize=6, fontweight="bold")
    plt.grid(color="0.8", linestyle="--", linewidth=0.5)

    plt.legend(loc="upper right", fontsize=8)

    plt.title(title, fontweight="bold", fontsize=10)
    plt.show()


if __name__ == "__main__":
    main()
