__author__ = "Tomasz Rybotycki"

from copy import deepcopy
from typing import List, Iterable

from numpy import ndarray
from scipy import special

from ..boson_sampling_utilities.boson_sampling_utilities import (
    generate_possible_outputs,
)
from ..distribution_calculators.bs_distribution_calculator_interface import (
    BosonSamplingExperimentConfiguration,
)
from ..distribution_calculators.bs_distribution_calculator_with_fixed_losses import (
    BSDistributionCalculatorWithFixedLosses,
    BSPermanentCalculatorInterface,
)
from multiprocessing import cpu_count, Pool


class BSDistributionCalculatorWithUniformLosses(
    BSDistributionCalculatorWithFixedLosses
):
    def __init__(
        self,
        configuration: BosonSamplingExperimentConfiguration,
        permanent_calculator: BSPermanentCalculatorInterface,
    ) -> None:
        super().__init__(configuration, permanent_calculator)
        self.weights = self._initialize_weights()
        self.weightless = False

    def set_weightless(self, weightless: bool) -> None:
        if not weightless:
            self.weights = self._initialize_weights()
        else:
            self.weights = [1 for _ in self.weights]

        self.weightless = weightless

    def _initialize_weights(self) -> List[float]:

        weight = (
            lambda n, l, eta: pow(eta, l) * special.binom(n, l) * pow(1.0 - eta, n - l)
        )
        weights = []

        for number_of_particles_left in range(
            self.configuration.initial_number_of_particles + 1
        ):
            weights.append(
                weight(
                    self.configuration.initial_number_of_particles,
                    number_of_particles_left,
                    self.configuration.uniform_transmissivity,
                )
            )

        return weights

    def calculate_probabilities_of_outcomes(
        self, outcomes: Iterable[Iterable[int]]
    ) -> List[float]:

        with Pool(processes=cpu_count()) as pool:
            outcomes_probabilities = pool.map(
                self._calculate_probability_of_outcome, outcomes
            )

        return outcomes_probabilities

    def _calculate_probability_of_outcome(self, outcome: ndarray) -> float:

        number_of_particles_left = int(sum(outcome))
        l = number_of_particles_left

        if l == 0:
            return self.weights[0]

        n = self.configuration.initial_number_of_particles

        subconfiguration = deepcopy(self.configuration)

        subconfiguration.number_of_particles_left = number_of_particles_left
        subconfiguration.number_of_particles_lost = n - l
        subdistribution_calculator = BSDistributionCalculatorWithFixedLosses(
            subconfiguration, self._permanent_calculator
        )

        probability_of_outcome = subdistribution_calculator.calculate_probabilities_of_outcomes(
            [outcome]
        )[
            0
        ]

        return probability_of_outcome * self.weights[l]

    def get_outcomes_in_proper_order(self) -> List[ndarray]:
        return generate_possible_outputs(
            self.configuration.initial_number_of_particles,
            self.configuration.number_of_modes,
            consider_loses=True,
        )
