from typing import Iterable

import numpy as np
import pandas as pd
import mlflow
from tqdm import tqdm

from poison.experiment import PoisonedDrinksExperiment
from poison.solvers.base import PoisonedDrinksSolver


class PoisonedDrinksLaboratory:
    """ A laboratory to perform experiments of poisoned drinks and record the results. """

    def __init__(self, num_experiments: int, poison_chance_values: Iterable[int],
                 num_glasses: int = 1000, seed: int = 8200):
        """
        Creates a laboratory to be used to run experiments of poisoned drinks.
        :param num_experiments: The number of experiments to create.
        :param num_glasses: The number of glasses to put in each experiment.
        :param poison_chance_values: The list of probabilities to experiment with.
        :param seed: The random seed used as the basis to create experiments of poisoned drinks.
        """
        self.num_experiments = num_experiments
        self.num_glasses = num_glasses
        self.poison_chance_values = poison_chance_values
        self.seed = seed
        # Create the poisoned drinks experiments to evaluate.
        self.experiments = dict([
            (p, [PoisonedDrinksExperiment(num_glasses=self.num_glasses, poison_chance=p, seed=seed + i)
                 for i in range(self.num_experiments)])
            for p in self.poison_chance_values])
        self.solver_results = dict()
        # Create the MLFlow experiment record.
        experiment_name = "Poisoned Drinks"
        if not mlflow.get_experiment_by_name(experiment_name):
            mlflow.create_experiment(name=experiment_name)
        self.mlflow_experiment = mlflow.get_experiment_by_name(name=experiment_name)

    def add_solver(self, solver: PoisonedDrinksSolver) -> None:
        """
        Adds a solver to the laboratory and tests the solver.
        :param solver: The solver to add.
        """
        # Resets the results associated with the solver name.
        self.solver_results[solver.name] = {}
        # Iterating the different poison probabilities we were requested to experiment with.
        for poison_chance in self.experiments.keys():
            status_description_string = "P={poison_percentage}%, experiment".format(
                solver_name=solver.name,
                poison_percentage=str(np.round(poison_chance * 100, 2)).rjust(5))
            # Solve all the experiments associated with the poison probability using the solver.
            self.solver_results[solver.name][poison_chance] = []
            for i in tqdm(range(self.num_experiments), desc=status_description_string):
                solver.load_experiment(experiment=self.experiments[poison_chance][i])
                solver.solve()
                self.solver_results[solver.name][poison_chance].append(solver.num_tests)
        # Logs the results of the solver.
        self.log_results(solver=solver)

    def log_results(self, solver: PoisonedDrinksSolver) -> None:
        with mlflow.start_run(experiment_id=self.mlflow_experiment.experiment_id, run_name=solver.name) as active_run:
            mlflow.log_params(solver.params)
            for p, test_counts in self.solver_results[solver.name].items():
                mlflow.log_metric(key="mean tests count", value=np.mean(test_counts), step=int(p * 1000))
                mlflow.log_metric(key="P_{0:.3f}".format(p), value=np.mean(test_counts))

    def show_results(self, log_axis=False) -> pd.DataFrame:
        """
        Shows the results generated by the tested solvers.
        :param log_axis: Whether to log axis (default = False).
        :return: The average number of tests for each solver and poison probability.
        """
        # Calculate means of each solver and parameters.
        means_list = []
        for solver_name in self.solver_results.keys():
            for p in self.solver_results[solver_name]:
                means_list.append({"Solver": solver_name, "P": p,
                                   "Mean # tests": np.mean(self.solver_results[solver_name][p])})
        means_df = pd.DataFrame(means_list).pivot('P', 'Solver', 'Mean # tests')
        # Plot the results.
        ax = means_df.plot(figsize=(14, 6))
        ax.set_xscale('log' if log_axis else 'linear')
        ax.set_yscale('log' if log_axis else 'linear')
        # Return the results.
        return means_df
