import multiprocessing as mp
import pickle
import shelve
from collections import defaultdict

from pathlib import Path
from typing import Union

from tqdm import trange

from lib.objects.map import Map
from lib.utils.parallel import run_n_in_parallel

class Generator:

    possible_modes = ['one_patching']

    def __init__(self, data_root, map_obj, binary=True):
        self.binary: bool = binary
        self.map: Map = map_obj

        self.data_root = Path(data_root)

    def generate_n_trajectories_m_alternatives(self, n, m, dataset_name='', **kwargs):
        trajectories_with_alternatives = list()
        for _ in trange(n, desc='Processing Trajectories'):
            trajectory = self.map.get_random_trajectory()
            alternatives, labels = self.generate_n_alternatives(trajectory, m, dataset_name=dataset_name, **kwargs)
            if not alternatives or labels:
                continue
            else:
                trajectories_with_alternatives.append(
                    dict(trajectory=trajectory, alternatives=alternatives, labels=labels)
                )
        return trajectories_with_alternatives

    def generate_alternatives(self, trajectory, output: Union[mp.
                              Queue, None] = None, mode='one_patching'):
        start, dest = trajectory.endpoints
        if mode == 'one_patching':
            patch = self.map.get_valid_position()
            alternative = self.map.get_trajectory_from_vertices(start, patch, dest)
        else:
            raise RuntimeError(f'mode checking went wrong...')

        if output:
            output.put(alternative)
        return alternative

    def generate_n_alternatives(self, trajectory, n, dataset_name: Union[str, Path] = '',
                                mode='one_patching', equal_samples=True, binary_check=True):
        assert mode in self.possible_modes, f'Parameter "mode" must be either {self.possible_modes}, but was {mode}.'
        # Define an output queue
        #output = mp.Queue()

        results = run_n_in_parallel(self.generate_alternatives, n, trajectory=trajectory, mode=mode)  # , output=output)

        # Get process results from the output queue
        #results = [output.get() for _ in range(n)]

        # label per homotopic class
        homotopy_classes = defaultdict(list)
        homotopy_classes[0].append(trajectory)
        for i in range(len(results)):
            alternative = results[i]
            class_not_found = True
            # check for homotopy class
            for label in homotopy_classes.keys():
                if self.map.are_homotopic(homotopy_classes[label][0], alternative):
                    homotopy_classes[label].append(alternative)
                    class_not_found = False
                    break
            if class_not_found:
                label = 1 if binary_check else len(homotopy_classes)
                homotopy_classes[label].append(alternative)

        # There should be as much homotopic samples as non-homotopic samples
        if equal_samples:
            homotopy_classes = self._remove_unequal(homotopy_classes)
            if not homotopy_classes:
                return None, None

        # Compose lists of alternatives with labels
        alternatives, labels = list(), list()
        for key in homotopy_classes.keys():
            alternatives.extend(homotopy_classes[key])
            labels.extend([key] * len(homotopy_classes[key]))

        # Write to disk
        if dataset_name:
            self.write_to_disk(dataset_name, trajectory, alternatives, labels)

        # Return
        return alternatives, labels

    def write_to_disk(self, filepath, trajectory, alternatives, labels):
        dataset_name = filepath if filepath.endswith('.pik') else f'{filepath}.pik'
        self.data_root.mkdir(exist_ok=True, parents=True)
        with shelve.open(str(self.data_root / dataset_name), protocol=pickle.HIGHEST_PROTOCOL) as f:
            new_key = len(f)
            f[f'trajectory_{new_key}'] = dict(alternatives=alternatives,
                                              trajectory=trajectory,
                                              labels=labels)
            if 'map' not in f:
                f['map'] = dict(map=self.map, name=f'map_{self.map.name}')

    @staticmethod
    def _remove_unequal(hom_dict):
        # We argue, that there will always be more non-homotopic routes than homotopic alternatives.
        # TODO: Otherwise introduce a second condition / loop
        hom_dict = hom_dict.copy()
        if len(hom_dict[0]) <= 1:
            return None
        counter = len(hom_dict)
        while sum([len(hom_dict[class_id]) for class_id in range(1, len(hom_dict))]) > len(hom_dict[0]):
            if counter == 0:
                counter = len(hom_dict)
            if counter in hom_dict:
                if len(hom_dict[counter]) == 0:
                    del hom_dict[counter]
                else:
                    del hom_dict[counter][-1]
            counter -= 1
        return hom_dict