__author__ = "Tomasz Rybotycki"

from typing import List

from numpy import arange, ndarray
from numpy.random import choice
from scipy import special

from ..boson_sampling_simulator import BosonSamplingSimulator
from .fixed_loss_simulation_strategy import FixedLossSimulationStrategy
from .simulation_strategy_interface import SimulationStrategyInterface


class UniformLossSimulationStrategy(SimulationStrategyInterface):
    def __init__(self, interferometer_matrix: ndarray,
                 number_of_modes: int, transmissivity: float) \
            -> None:
        self.interferometer_matrix = interferometer_matrix
        self.number_of_modes = number_of_modes
        self.transmissivity = transmissivity

    def simulate(self, input_state: ndarray, samples_number: int = 1) -> List[ndarray]:
        initial_number_of_particles = int(sum(input_state))

        # Using n, eta, l notation from the paper.
        n = initial_number_of_particles
        eta = self.transmissivity

        separable_states_weights = [pow(eta, l) * special.binom(n, l) * pow(1.0 - eta, n - l) for l in range(n + 1)]

        samples = []
        while len(samples) < samples_number:
            number_of_particles_left_in_selected_separable_state = choice(arange(0, n + 1), p=separable_states_weights)

            strategy = FixedLossSimulationStrategy(self.interferometer_matrix,
                                                   number_of_particles_left_in_selected_separable_state,
                                                   self.number_of_modes)

            simulator = BosonSamplingSimulator(strategy)

            samples.append(simulator.get_classical_simulation_results(input_state)[0])

        return samples
