from __future__ import annotations

import pathlib

# from functools import partial
from itertools import compress
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union

import cytoolz as tz
import numpy as np
import tqdm

from robustnessgym.core.constants import (
    ATTACK,
    CURATION,
    GENERIC,
    SLICEBUILDERS,
    SUBPOPULATION,
    TRANSFORMATION,
)
from robustnessgym.core.dataset import Batch, BatchOrDataset, Dataset
from robustnessgym.core.identifier import Identifier
from robustnessgym.core.slice import Slice
from robustnessgym.core.storage import StorageMixin
from robustnessgym.core.tools import recmerge, strings_as_json

# from multiprocess.pool import Pool


class SliceBuilder(StorageMixin):
    """Base class for builders that output slices."""

    # Path to a log directory
    logdir: pathlib.Path = pathlib.Path.home() / "robustnessgym" / SLICEBUILDERS

    # Create the log directory
    logdir.mkdir(parents=True, exist_ok=True)

    CATEGORIES = [
        GENERIC,
        SUBPOPULATION,
        ATTACK,
        TRANSFORMATION,
        CURATION,
    ]

    def __init__(
        self,
        category: str,
        identifiers: List[Identifier],
        apply_fn: Callable = None,
        *args,
        **kwargs,
    ):

        super(SliceBuilder, self).__init__(*args, **kwargs)

        # The SliceMaker belongs to a category
        assert (
            category in self.CATEGORIES
        ), f"argument category must be one of {self.CATEGORIES}"
        self.category = category

        # Each identifier corresponds to a single output Slice generated by this
        # SliceBuilder
        self.identifiers = identifiers

        # Keep track of the CachedOperation dependencies
        self.prerequisites = (
            set() if "prerequisites" not in kwargs else kwargs["prerequisites"]
        )
        # TODO(karan): remove comment
        # for base in self.__class__.__bases__:
        #     for cls in base.__mro__:
        #         if str(CachedOperation.__name__) in str(cls):
        #             self.prerequisites.add(base)

        if apply_fn:
            # Assign to the method
            self.apply = apply_fn

    def __repr__(self):
        return (
            f"{self.category}[{self.__class__.__name__}(num_slices={self.num_slices})]"
        )

    @property
    def num_slices(self):
        return len(self.identifiers)

    def __getitem__(self, item: int):
        return self.identifiers[item]

    def __iter__(self):
        yield from self.identifiers

    def __call__(
        self,
        batch_or_dataset: BatchOrDataset,
        columns: List[str],
        mask: List[int] = None,
        num_proc: int = None,
        batch_size: int = 100,
        *args,
        **kwargs,
    ):

        # Check that prerequisites are satisfied
        self.prerequisites_handler(batch_or_dataset)

        if isinstance(batch_or_dataset, Dataset):
            # Prepare the dataset
            self.prepare_dataset(
                dataset=batch_or_dataset,
                columns=columns,
                batch_size=batch_size,
                *args,
                **kwargs,
            )

            # Slice a dataset
            slices, slice_membership = self.process_dataset(
                dataset=batch_or_dataset,
                columns=columns,
                # # Automatically infer the mask from the Dataset if it's not specified
                # mask=batch_or_dataset.check_tape(
                #     path=[SLICEBUILDERS, self.category],
                #     identifiers=self.identifiers,
                #     columns=columns,
                # )
                # if not mask
                # else mask,
                num_proc=num_proc,
                *args,
                **kwargs,
            )

            # Update the Dataset's history
            # TODO(karan): use mask to figure out what is actually applied
            batch_or_dataset.update_tape(
                path=[SLICEBUILDERS, self.category],
                identifiers=self.identifiers,
                columns=columns,
            )

            return slices, slice_membership

        elif isinstance(batch_or_dataset, Dict):
            # Slice a batch
            return self.process_batch(
                batch=batch_or_dataset,
                columns=columns,
                *args,
                **kwargs,
            )
        else:
            raise NotImplementedError

    def prerequisites_handler(
        self,
        batch_or_dataset: BatchOrDataset,
    ):
        if isinstance(batch_or_dataset, Dataset):
            batch = batch_or_dataset[:2]
        else:
            batch = batch_or_dataset

        # Check if prerequisites are satisfied
        # TODO(karan): move to a method
        pending = {
            prerequisite
            for prerequisite in self.prerequisites
            if not prerequisite.available(batch)
        }

        # TODO(karan): Automatically run the pending pre-requisites
        if pending:
            raise RuntimeError(
                f"Cannot run SliceBuilder, prerequisites {pending} not satisfied."
            )

    @staticmethod
    def store(batch: Dict[str, List], updates: List[Dict]) -> Dict[str, List]:
        """Update a batch of examples with slice information."""
        if "slices" not in batch:
            batch["slices"] = [{} for _ in range(len(batch["index"]))]

        # For each example, recursively merge the example's original cache dictionary
        # with the update dictionary
        batch["slices"] = [
            recmerge(example_dict, update_dict, merge_sequences=True)
            for example_dict, update_dict in zip(batch["slices"], updates)
        ]

        return batch

    def prepare_batch(
        self,
        batch: Batch,
        columns: List[str],
        *args,
        **kwargs,
    ) -> None:
        """Apply a preparation function to a batch. Use this to update
        attributes of `self`.

        Args:
            batch: batch of data
            columns: list of columns
            *args: optional additional arguments
            **kwargs: optional additional keyword arguments
        """
        raise NotImplementedError("Implemented `prepare_batch` should return a batch.")

    def _filter_prerequisite_columns(
        self,
        columns: List[str],
        all_columns: List[str],
    ) -> List[str]:
        # Simple filtering that doesn't use columns
        # TODO(karan): improve this by using `columns` to filter further
        return [
            col
            for col in all_columns
            if any(
                [
                    col.startswith(
                        prereq.__name__ if isinstance(prereq, type) else str(prereq)
                    )
                    for prereq in self.prerequisites
                ]
            )
        ]

    def prepare_dataset(
        self,
        dataset: Dataset,
        columns: List[str],
        batch_size: int = 32,
        *args,
        **kwargs,
    ) -> None:
        """Apply a preparation function to the dataset. Use this to update
        attributes of `self`.

        Args:
            dataset: dataset
            columns: list of columns
            batch_size: batch size for preparation
            *args: optional additional arguments
            **kwargs: optional additional keyword arguments
        """
        # Set the data format
        with dataset.format(
            columns + self._filter_prerequisite_columns(columns, dataset.column_names)
        ):
            # Batch the dataset, and prepare each batch
            for batch in dataset.batch(batch_size):
                try:
                    # Check if the `prepare_batch` function has been implemented
                    self.prepare_batch(
                        batch=batch,
                        columns=columns,
                        *args,
                        **kwargs,
                    )
                except NotImplementedError:
                    break

    def process_batch(
        self,
        batch: Dict[str, List],
        columns: List[str],
        *args,
        **kwargs,
    ) -> Tuple[List[Dict[str, List]], Optional[np.ndarray]]:
        """Apply a SliceBuilder to a batch of data.

        Args:
            batch: a batch of data
            columns: list of columns
            *args: optional additional arguments
            **kwargs: optional additional keyword arguments

        Returns: tuple of (list of slices (as batches), matrix of (example,
        slice) membership))
        """
        return [batch], None

    def process_dataset(
        self,
        dataset: Dataset,
        columns: List[str],
        batch_size: int = 32,
        # mask: List[int] = None,
        num_proc: int = None,
        *args,
        **kwargs,
    ) -> Tuple[List[Slice], np.ndarray]:
        """Apply a SliceBuilder to a dataset.

        Args:
            dataset: Dataset
            columns: list of columns
            batch_size: integer batch size
            # mask: boolean or integer mask array, mask[i] = True means that the ith
            # slice will be masked out
            num_proc: num processes for multiprocessing
            *args: optional additional arguments
            **kwargs: optional additional keyword arguments

        Returns: tuple of (Dataset, list of Slices, matrix of (example,
        slice) membership)
        """

        # # Compute a hash
        # val = persistent_hash(str(dataset.identifier)) ^ dataset.hash_interactions()
        # for i, identifier in enumerate(self.identifiers):
        #     if not mask[i]:
        #         val ^= persistent_hash(str(identifier)
        #         + str(strings_as_json(columns)))

        # try:
        #     # Map the SliceBuilder over the dataset
        #     all_sliced_batches = []
        #     all_slice_memberships = []
        #
        #     def _map_fn(batch):
        #         """Map function for processing batches.
        #
        #         Note that using this map_fn in a stateful way is
        #         dangerous, since every invocation of this function
        #         appends to the all_slice_batches list. The .map()
        #         function will invoke this once for testing before
        #         performing the map, so we discard the first entry
        #         inserted into all_sliced_batches.
        #         """
        #         batch, sliced_batches, slice_membership = self.process_batch(
        #             batch=batch,
        #             columns=columns,
        #             mask=mask,
        #             store_compressed=store_compressed,
        #             store=store,
        #             *args,
        #             **kwargs,
        #         )
        #         all_sliced_batches.append(sliced_batches)
        #         all_slice_memberships.append(slice_membership)
        #         return batch
        #
        #     dataset = dataset.map(
        #         _map_fn,
        #         batched=True,
        #         batch_size=batch_size,
        #         # FIXME(karan): enable this by adding logic for generating
        #         #  all_sliced_batches and all_slice_memberships
        #         #  when loading from cache file
        #         load_from_cache_file=False,
        #         # The cache file name is a XOR of the interaction history and the
        #         # current operation
        #         cache_file_name=str(
        #             dataset.logdir / ("cache-" + str(abs(val)) + ".arrow")
        #         ),
        #     )
        #
        #     # Remove the first entry (see _map_fn)
        #     all_sliced_batches = all_sliced_batches[1:]
        #     all_slice_memberships = all_slice_memberships[1:]
        #
        # except:  # noqa
        # all_batches, all_sliced_batches, all_slice_memberships = zip(
        #     *[
        #         self.process_batch(
        #             batch=batch,
        #             columns=columns,
        #             mask=mask,
        #             store_compressed=store_compressed,
        #             store=store,
        #             *args,
        #             **kwargs,
        #         )
        #         for batch in dataset.batch(batch_size)
        #     ]
        # )
        # # Update the dataset efficiently by reusing all_batches
        # dataset = dataset.map(
        #     lambda examples, indices: all_batches[indices[0] // batch_size],
        #     batched=True,
        #     batch_size=batch_size,
        #     with_indices=True,
        #     load_from_cache_file=False,
        #     # The cache file name is a XOR of the interaction history and the
        #     # current operation
        #     cache_file_name=str(
        #         dataset.logdir / ("cache-" + str(abs(val)) + ".arrow")
        #     ),
        # )

        # Create slices
        slices = [Slice() for _ in range(len(self.identifiers))]
        all_slice_memberships = []
        # Batch the dataset, and process each batch
        for batch in dataset.batch(batch_size):
            # Process the batch
            sliced_batches, slice_memberships = self.process_batch(
                batch=batch,
                columns=columns,
                *args,
                **kwargs,
            )

            # Incrementally build the slices
            for sl, sl_batch in zip(slices, sliced_batches):
                sl._dataset.append(sl_batch)

            # Keep track of the slice memberships
            all_slice_memberships.append(slice_memberships)

        # Create a single slice label matrix
        slice_membership = np.concatenate(all_slice_memberships, axis=0)

        # slice_cache_hashes = []
        # for identifier in self.identifiers:
        #     slice_cache_hashes.append(val ^ persistent_hash(str(identifier)))

        # if not num_proc or num_proc == 1:
        #     # Construct slices
        #     slices = []
        #     for i, slice_batches in enumerate(zip(*all_sliced_batches)):
        #         slices.append(
        #             create_slice(
        #                 (
        #                     dataset,
        #                     slice_membership,
        #                     slice_batches,
        #                     i,
        #                     batch_size,
        #                     slice_cache_hashes[i],
        #                 )
        #             )
        #         )
        # else:
        #     # Parallelized slice construction
        #     with Pool(num_proc) as pool:
        #         slices = pool.map(
        #             create_slice,
        #             [
        #                 (
        #                     dataset,
        #                     slice_membership,
        #                     slice_batches,
        #                     i,
        #                     batch_size,
        #                     slice_cache_hashes[i],
        #                 )
        #                 for i, slice_batches in enumerate(zip(*all_sliced_batches))
        #             ],
        #         )

        for i, sl in enumerate(slices):
            # Set the Slice category using the SliceBuilder's category
            sl.category = self.category

            # Append the the lineage
            sl.add_to_lineage(
                category=str(self.category.capitalize()),
                identifier=self.identifiers[i],
                columns=strings_as_json(columns),
            )

            # # Create the lineage
            # sl.lineage = [
            #     (str(Dataset.__name__), dataset.identifier),
            #     (
            #         str(self.category.capitalize()),
            #         self.identifiers[i],
            #         strings_as_json(columns),
            #     ),
            # ]
            # if isinstance(dataset, Slice):
            #     # Prepend the Slice's lineage instead, if the dataset was a slice
            #     sl.lineage = dataset.lineage + [
            #         (
            #             str(self.category.capitalize()),
            #             self.identifiers[i],
            #             strings_as_json(columns),
            #         )
            #     ]

        return slices, slice_membership

    def apply(self, *args, **kwargs):
        raise NotImplementedError("Must implement apply.")

    @classmethod
    def join(cls, *slicemakers: SliceBuilder) -> Sequence[SliceBuilder]:
        """Join many slicemakers.

        By default, just returns the slicemakers.
        """
        return slicemakers

    def masked(self, mask: List[int]):
        pass

    def unmasked(self):
        pass

    @staticmethod
    def filter_batch_by_slice_membership(
        batch: Dict[str, List], slice_membership: np.ndarray
    ) -> List[Dict[str, List]]:
        """Use a matrix of slice membership labels to select the subset of
        examples in each slice.

        Returns a list. Each element in the list corresponds to a single
        slice, and contains the subset of examples in 'batch' that lies
        in that slice.
        """
        return [
            tz.valmap(lambda v: list(compress(v, s)), batch) for s in slice_membership.T
        ]

    @classmethod
    def retrieve(
        cls,
        batch: Batch,
        columns: Union[List[str], List[List[str]]],
        proc_fns: Union[str, Callable, List[Union[str, Callable]]] = None,
        identifier: Union[str, Identifier] = None,
        reapply: bool = False,
        **kwargs,
    ) -> Optional[Union[Batch, List[Batch]]]:
        if not reapply:
            if "slices" not in batch:
                return None

            # Infer the most relevant key to retrieve if an identifier is not specified
            if not identifier:
                for ident_key in batch["slices"][0].keys():
                    # Pick the first key that matches the cls name
                    if ident_key.startswith(cls.__name__):
                        identifier = ident_key
                        break

            try:
                if isinstance(columns[0], str):
                    retrieval = {
                        strings_as_json(columns): [
                            cls.decode(cache[str(identifier)][strings_as_json(columns)])
                            for cache in batch["cache"]
                        ]
                    }
                else:
                    retrieval = {
                        strings_as_json(cols_): [
                            cls.decode(cache[str(identifier)][strings_as_json(cols_)])
                            for cache in batch["cache"]
                        ]
                        for cols_ in columns
                    }
            except KeyError:
                raise ValueError("Could not retrieve information for all keys.")

            # Check if the retrieved information needs to be processed
            if not proc_fns:
                return retrieval
            pass
        else:
            pass


class SliceBuilderCollection(SliceBuilder):
    """Collection of Slice Builders."""

    def __init__(self, slicebuilders: List[SliceBuilder], *args, **kwargs):
        super(SliceBuilderCollection, self).__init__(
            category=GENERIC,
            identifiers=list(
                tz.concat([slicebuilder.identifiers for slicebuilder in slicebuilders])
            ),
            *args,
            **kwargs,
        )

        # TODO(karan): some slicebuilders aren't compatible with each other (e.g.
        #  single column vs. multi column):
        #  add some smarter logic here to handle this

        # Store the subpopulations
        self.slicebuilders = slicebuilders

    def __repr__(self):
        # TODO(karan): format this nicely
        return (
            f"{self.__class__.__name__}("
            f"{[str(slicebuilder) for slicebuilder in self.slicebuilders]})]"
        )

    def __call__(
        self,
        batch_or_dataset: BatchOrDataset,
        columns: List[str],
        mask: List[int] = None,
        store_compressed: bool = None,
        store: bool = None,
        *args,
        **kwargs,
    ):

        if mask:
            raise NotImplementedError(
                "Mask not supported for SliceBuilderCollection yet."
            )

        slices = []
        slice_membership = []
        # Apply each slicebuilder in sequence
        for i, slicebuilder in tqdm.tqdm(enumerate(self.slicebuilders)):
            # Apply the slicebuilder
            batch_or_dataset, slices_i, slice_membership_i = slicebuilder(
                batch_or_dataset=batch_or_dataset,
                columns=columns,
                mask=mask,
                store_compressed=store_compressed,
                store=store,
                *args,
                **kwargs,
            )

            # Add in the slices and slice membership
            slices.extend(slices_i)
            slice_membership.append(slice_membership_i)

        slice_membership = np.concatenate(slice_membership, axis=1)

        return batch_or_dataset, slices, slice_membership


def create_slice(args):
    # Unpack args
    dataset, slice_membership, slice_batches, i, batch_size, slice_cache_hash = args

    # Create a new empty slice
    sl = Slice.from_dict({})

    # Create a Slice "copy" of the Dataset
    sl.__dict__.update(dataset.__dict__)
    sl._identifier = None

    # Filter
    sl = sl.filter(
        lambda example, idx: bool(slice_membership[idx, i]),
        with_indices=True,
        input_columns=["index"],
        batch_size=batch_size,
        cache_file_name=str(
            dataset.logdir / ("cache-" + str(abs(slice_cache_hash)) + "-filter.arrow")
        ),
    )

    slice_batch = tz.merge_with(tz.compose(list, tz.concat), slice_batches)

    # FIXME(karan): interaction tape history is wrong here, esp with augmenation/attacks

    # Map
    if len(sl):
        sl = sl.map(
            lambda batch, indices: tz.valmap(
                lambda v: v[indices[0] : indices[0] + batch_size], slice_batch
            ),
            batched=True,
            batch_size=batch_size,
            with_indices=True,
            remove_columns=sl.column_names,
            cache_file_name=str(
                dataset.logdir / ("cache-" + str(abs(slice_cache_hash)) + ".arrow")
            ),
        )

    return sl
