

from typing import Iterable, Any, Callable, List, Dict

import numpy as np
import torchdata.datapipes as dp
from itertools import repeat, chain

from .base import PostProcessor
from ..fields import Field


__all__ = [
    "LeftPruningRow", "RightPruningRow",
    "AddingRow",
    "LeftPaddingRow", "RightPaddingRow",
]


#==================================Filter==================================
class RowFilter(PostProcessor):
    r"""
    Apply a function to specific indices of each row in an IterableDataset.

    Parameters:
    -----------
    source_dp : tud.IterableDataset
        The source IterableDataset.
    fn : Callable
        The function to apply to the specified indices of each row.
    checked_fields : Iterable[Field]
        The fields to be checked.

    Raises:
    -------
    TypeError
        If `indices` is not an iterable.
    """

    def __init__(
        self, source: dp.iter.IterableWrapper,
        fn: Callable, checked_fields: Iterable[Field]
    ):
        super().__init__(source)
        self.fn = fn
        self.checked_fields = set(self.sure_input_fields()) & set(checked_fields)

    def _check(self, row: Dict[Field, Any]) -> bool:
        return all(self.fn(field, row[field]) for field in self.checked_fields)

    def __iter__(self):
        for row in self.source:
            if self._check(row):
                yield row


#==================================Mapper==================================
class RowMapper(PostProcessor):
    r"""
    Apply a function to specific indices of each row in an IterableDataset.

    Parameters:
    -----------
    source_dp : tud.IterableDataset
        The source IterableDataset.
    fn : Callable
        The function to apply to the specified indices of each row.
    modified_fields : Iterable[Field]
        The fields to be modified.

    Raises:
    -------
    TypeError
        If `indices` is not an iterable.
    """

    def __init__(
        self, source: dp.iter.IterableWrapper,
        fn: Callable, modified_fields: Iterable[Field]
    ):
        super().__init__(source)

        self.fn = fn
        self.modified_fields = set(self.sure_input_fields()) & set(modified_fields)

    def __iter__(self):
        for row in self.source:
            for field in self.modified_fields:
                row[field] = self.fn(field, row[field])
            yield row


@dp.functional_datapipe("lprune_")
class LeftPruningRow(RowMapper):
    r"""
    A functional datapipe that prunes the left side of a given datapipe to a specified maximum length.

    Parameters:
    -----------
    maxlen: int 
        The maximum length to prune the input data to.
    modifields_fields: Iterable[Field]
        The fields to be modified.

    Flows:
    ------
    [1, 2, 3, 4] --(maxlen=3)--> [2, 3, 4]
    [3, 4] --(maxlen=3)--> [3, 4]

    Examples:
    ---------
    >>> dataset: RecDataSet
    >>> ISeq = dataset[ITEM, ID].fork(SEQUENCE)
    >>> datapipe = dataset.valid().ordered_user_ids_source(
    ).valid_sampling_().lprune_(
        3, modified_fields=(ISeq,)
    )
    >>> next(iter(datapipe))
    {Field(USER:ID,USER): 0,
    Field(ITEM:ID,ITEM,SEQUENCE): (9839, 10076, 11155),
    Field(ITEM:ID,ITEM,UNSEEN): (11752,),
    Field(ITEM:ID,ITEM,SEEN): (9449, 9839, 10076, 11155)}
    """

    def __init__(
        self, source: dp.iter.IterDataPipe, maxlen: int, modified_fields: Iterable[Field]
    ) -> None:

        self.maxlen = maxlen

        super().__init__(
            source, self._prune,
            modified_fields
        )

    def _prune(self, field: Field, x: Iterable) -> Iterable:
        return x[-self.maxlen:]


@dp.functional_datapipe("rprune_")
class RightPruningRow(LeftPruningRow):
    r"""
    A functional datapipe that prunes the right side of a given datapipe to a specified maximum length.

    Parameters:
    -----------
    maxlen: int 
        The maximum length to prune the input data to.
    modifields_fields: Iterable[Field]
        The fields to be modified.

    Flows:
    ------
    [1, 2, 3, 4] --(maxlen=3)--> [1, 2, 3]
    [3, 4] --(maxlen=3)--> [3, 4]

    Examples:
    ---------
    >>> dataset: RecDataSet
    >>> ISeq = dataset[ITEM, ID].fork(SEQUENCE)
    >>> datapipe = dataset.valid().ordered_user_ids_source(
    ).valid_sampling_().rprune_(
        3, modified_fields=(ISeq,)
    )
    >>> next(iter(datapipe))
    {Field(USER:ID,USER): 0,
    Field(ITEM:ID,ITEM,SEQUENCE): (9449, 9839, 10076),
    Field(ITEM:ID,ITEM,UNSEEN): (11752,),
    Field(ITEM:ID,ITEM,SEEN): (9449, 9839, 10076, 11155)}
    """

    def _prune(self, field: Field, x: Iterable) -> Iterable:
        return x[:self.maxlen]


@dp.functional_datapipe("add_")
class AddingRow(RowMapper):
    r"""
    Mapper that adds the input data by a specified offset.

    Parameters:
    -----------
    offset: int
        Amount to add the input data by.   
    modifields_fields: Iterable[Field]
        The fields to be modified.

    Flows:
    ------
    [1, 2, 3, 4] --(offset=1)--> [2, 3, 4, 5]

    Examples:
    ---------
    >>> dataset: RecDataSet
    >>> ISeq = dataset[ITEM, ID].fork(SEQUENCE)
    >>> datapipe = dataset.valid().ordered_user_ids_source(
    ).valid_sampling_().rprune_(
        3, modified_fields=(ISeq,)
    ).add_(
        1, modified_fields=(ISeq,)
    )
    """

    def __init__(
        self, source: dp.iter.IterDataPipe, offset: int, modified_fields: Iterable[Field]
    ) -> None:

        self.offset = offset

        super().__init__(
            source, self._add,
            modified_fields
        )

    def _add(self, field: Field, x: Iterable) -> List:
        r"""
        Examples:
        ---------
        >>> x = [1, 2, 3] 
        >>> _add(x) # self.offset = 1
        [2, 3, 4]
        >>> _add(x) # self.offset = -1
        [0, 1, 2]
        """
        return np.add(x, self.offset).tolist()


@dp.functional_datapipe("lpad_")
class LeftPaddingRow(RowMapper):
    r"""
    A functional data pipeline component that left pads sequences to a maximum length.

    Parameters:
    -----------
    maxlen : int
        The maximum length to pad the sequences to.
    modifields_fields: Iterable[Field]
        The fields to be modified.
    padding_value : int, optional (default=0)
        The value to use for padding.

    Flows:
    ------
    [1, 2, 3, 4] --(maxlen=7, padding_value=0)--> [0, 0, 0, 1, 2, 3, 4]
    [1, 2, 3, 4] --(maxlen=4, padding_value=0)--> [1, 2, 3, 4]

    Examples:
    ---------
    >>> dataset: RecDataSet
    >>> ISeq = dataset[ITEM, ID].fork(SEQUENCE)
    >>> datapipe = dataset.valid().ordered_user_ids_source(
    ).valid_sampling_().lprune_(
        3, modified_fields=(ISeq,)
    ).add_(
        1, modified_fields=(ISeq,)
    ).lpad_(
        5, modified_fields=(ISeq,)
    )
    >>> next(iter(datapipe))
    {Field(USER:ID,USER): 0,
    Field(ITEM:ID,ITEM,SEQUENCE): [0, 0, 9840, 10077, 11156],
    Field(ITEM:ID,ITEM,UNSEEN): (11752,),
    Field(ITEM:ID,ITEM,SEEN): (9449, 9839, 10076, 11155)}
    """

    def __init__(
        self, source: dp.iter.IterDataPipe, maxlen: int, modified_fields: Iterable[Field], padding_value: int = 0
    ) -> None:

        self.maxlen = maxlen
        self.padding_value = int(padding_value)

        super().__init__(
            source, self._pad,
            modified_fields
        )

        self.sure_zero_elems()

    def sure_zero_elems(self):
        def guess_zero(field: Field, value: Iterable):
            if isinstance(value, Iterable):
                if len(value) == 0:
                    return self.padding_value
                elif isinstance(value[0], Iterable):
                    return (self.padding_value,) * len(value[0])
                else:
                    return self.padding_value
            else:
                raise ValueError(f"{value} for {field} is non-iterable ...")

        row = next(iter(self.source))
        self.zeros = dict()
        for field in self.modified_fields:
            self.zeros[field] = guess_zero(field, row[field])

    def _pad(self, field: Field, x: Iterable) -> List:
        return list(chain(repeat(self.zeros[field], self.maxlen - len(x)), x))


@dp.functional_datapipe("rpad_")
class RightPaddingRow(LeftPaddingRow):
    r"""
    A functional data pipeline component that right pads sequences to a maximum length.

    Parameters:
    -----------
    maxlen : int
        The maximum length to pad the sequences to.
    modifields_fields: Iterable[Field]
        The fields to be modified.
    padding_value : int, optional (default=0)
        The value to use for padding.

    Flows:
    ------
    [1, 2, 3, 4] --(maxlen=7, padding_value=0)--> [1, 2, 3, 4, 0, 0, 0]
    [1, 2, 3, 4] --(maxlen=4, padding_value=0)--> [1, 2, 3, 4]

    Examples:
    ---------
    >>> dataset: RecDataSet
    >>> ISeq = dataset[ITEM, ID].fork(SEQUENCE)
    >>> datapipe = dataset.valid().ordered_user_ids_source(
    ).valid_sampling_().rprune_(
        3, modified_fields=(ISeq,)
    ).add_(
        1, modified_fields=(ISeq,)
    ).rpad_(
        5, modified_fields=(ISeq,)
    )
    >>> next(iter(datapipe))
    {Field(USER:ID,USER): 0,
    Field(ITEM:ID,ITEM,SEQUENCE): [9450, 9840, 10077, 0, 0],
    Field(ITEM:ID,ITEM,UNSEEN): (11752,),
    Field(ITEM:ID,ITEM,SEEN): (9449, 9839, 10076, 11155)}
    """

    def _pad(self, field: Field, x: Iterable) -> List:
        return list(chain(x, repeat(self.zeros[field], self.maxlen - len(x))))