# -*- coding: utf-8 -*-
'''
ReilData class
==============

The main datatype used to communicate `state`s, `action`s, and `reward`s,
between objects in `reil`. `ReilData` is basically a tuple that contains
instances of `BaseData`, `CategoricalData`, and `NumericalData`.
'''
from __future__ import annotations

import dataclasses
import itertools
from dataclasses import field
from typing import (Any, Callable, Dict, Generic, Iterable, Iterator, List,
                    Mapping, Optional, Sequence, Tuple, TypeVar, Union, cast,
                    overload)

from typing_extensions import Literal

T = TypeVar('T')

Categorical = TypeVar('Categorical')
Numerical = TypeVar('Numerical', int, float)

Normal = Union[Literal[0], Literal[1], float]
Normalized = Union[Normal, Sequence[Normal], None]
Normalizer = Callable[..., Normalized]


@dataclasses.dataclass(frozen=True)
class BaseData(Generic[T]):
    '''
    The base class for `ReilData` elements.

    Attributes
    ----------
    name:
        Name of the data.

    value:
        Value of the data. Can be one item or a sequence.

    normalizer:
        A function that gets a `BaseData` object and returns its normalized
        form.

    lazy_evaluation:
        Whether normalization should be done at the point of instantiation
        (False) or at the access time (True).


    :meta private:
    '''
    name: str
    value: Optional[Union[T, Sequence[T]]] = None
    normalizer: Optional[Normalizer] = lambda _: None
    lazy_evaluation: Optional[bool] = field(default=None, compare=False)
    is_numerical: Optional[bool] = field(
        default=None, init=False, repr=False, compare=False)
    _normalized: Normalized = field(
        default=None, init=False, repr=False, compare=False)

    def __post_init__(self):
        # if not isinstance(self.value, Hashable):
        #     warnings.warn('Non-hashable object!')
        self._validate(self)
        if self.normalizer is None:
            object.__setattr__(self, 'normalizer',
                               self._default_normalizer)
        if not self.lazy_evaluation:
            object.__setattr__(self, '_normalized',
                               self.normalizer(self))

    @staticmethod
    def _validate(data: BaseData[T]) -> None:
        '''
        Validate the input data.

        Arguments
        ---------
        data:
            An instance to be validated.

        Raises
        ------
        ValueError:
            if `data.value` does not match the validation criteria.


        :meta public:
        '''
        pass

    @staticmethod
    def _default_normalizer(data: BaseData[T]) -> Normalized:
        '''
        Normalize the data.

        Arguments
        ---------
        data:
            An instance to be normalized.

        Returns
        -------
        :
            The normalized form of `data.value`.
        '''
        return None

    @property
    def normalized(self) -> Normalized:
        if self._normalized is None:
            object.__setattr__(self, '_normalized',
                               self.normalizer(self))

        return self._normalized

    def as_dict(self) -> Dict[str, Any]:
        '''
        Return the data as a dictionary.

        Returns
        -------
        :
            The data as a dictionary.
        '''
        return {'name': self.name, 'value': self.value}


@dataclasses.dataclass(frozen=True)
class CategoricalData(BaseData[Categorical]):
    '''
    A datatype for categorical data.

    Attributes
    ----------
    name:
        Name of the data.

    value:
        Value of the data. Can be one item or a sequence.

    normalizer:
        A function that gets a `BaseData` object and returns its normalized
        form.

    lazy_evaluation:
        Whether normalization should be done at the point of instantiation
        (False) or at the access time (True).

    categories:
        A list of all categories.
    '''
    categories: Optional[Tuple[Categorical, ...]] = None

    def __post_init__(self):
        super().__post_init__()
        object.__setattr__(self, 'is_numerical', False)

    @staticmethod
    def _default_normalizer(data: CategoricalData) -> Normalized:
        value = data.value
        categories = data.categories
        if categories is None or value is None:
            return None

        if isinstance(value, type(categories[0])):
            return tuple(int(x_i == value)
                         for x_i in categories)

        return tuple(int(x_i == v)
                     for v in value
                     for x_i in categories)

    @staticmethod
    def _validate(data: CategoricalData) -> None:
        value = data.value
        if value is not None:
            is_sequence = isinstance(value, (list, tuple))

            categories = data.categories
            if categories is not None:
                if ((not is_sequence and value not in categories) or
                        (is_sequence and any(v not in categories
                                             for v in value))):
                    raise ValueError(
                        f'value={value} is '
                        f'in the categories={categories}.')

    def as_dict(self) -> Dict[str, Any]:
        return {'name': self.name, 'value': self.value,
                'categories': self.categories,
                'is_numerical': self.is_numerical}


@dataclasses.dataclass(frozen=True)
class NumericalData(BaseData[Numerical]):
    '''
    A datatype for numerical data.

    Attributes
    ----------
    name:
        Name of the data.

    value:
        Value of the data. Can be one item or a sequence.

    normalizer:
        A function that gets a `BaseData` object and returns its normalized
        form.

    lazy_evaluation:
        Whether normalization should be done at the point of instantiation
        (False) or at the access time (True).

    lower:
        The lower bound of value.

    upper:
        The upper bound of value.
    '''
    lower: Optional[Numerical] = None
    upper: Optional[Numerical] = None

    def __post_init__(self):
        super().__post_init__()
        object.__setattr__(self, 'is_numerical', True)

    @staticmethod
    def _default_normalizer(data: NumericalData) -> Normalized:
        value = data.value
        lower = data.lower
        upper = data.upper
        if value is None or lower is None or upper is None:
            return None

        denominator = upper - lower

        sequence_check = isinstance(value, (list, tuple))
        try:
            if sequence_check:
                return tuple((v - lower) / denominator
                             for v in value)
            else:
                return (value - lower) / denominator
        except ZeroDivisionError:
            return [1] * len(value) if sequence_check else 1

    @staticmethod
    def _validate(data: NumericalData) -> None:
        value = data.value
        lower = data.lower
        upper = data.upper
        if value is not None:
            is_sequence = isinstance(value, (list, tuple))

            if lower is not None:
                if ((not is_sequence and value < lower) or
                        (is_sequence and any(v < lower
                                             for v in value))):
                    raise ValueError(
                        f'value={value} is less than lower={lower}.')

            if upper is not None:
                if ((not is_sequence and value > upper) or
                        (is_sequence and any(v > upper
                                             for v in value))):
                    raise ValueError(
                        f'value={value} is '
                        f'greater than upper={upper}.')

    def as_dict(self) -> Dict[str, Any]:
        return {'name': self.name, 'value': self.value,
                'lower': self.lower, 'upper': self.upper,
                'is_numerical': self.is_numerical}


ReilDataInput = Union[Mapping[str, Any],
                      BaseData[Any],
                      CategoricalData[Categorical],
                      NumericalData[Numerical]]


class ReilData(Sequence[ReilDataInput], Generic[Categorical, Numerical]):
    '''
    The main datatype used to communicate `state`s, `action`s, and `reward`s,
    between objects in `reil`.
    '''

    def __init__(self,
                 data: Union[ReilDataInput,
                             Sequence[ReilDataInput],
                             Iterator[ReilDataInput]],
                 lazy_evaluation: Optional[bool] = None):
        '''
        Arguments
        ---------
        data:
            One or a sequence of `BaseData`, `CategoricalData`,
            `NumericalData`, or dicts that include 'name'. Other attributes are
            optional. If none of `categories`, `lower` and `upper` are
            provided, the object is assumed `BaseData`.

        lazy_evaluation:
            Whether to compute normalized value at the time of instantiation or
            at the time of the first access. If not provided, `ReilData` looks
            for `lazy_evaluation` in each item. If failed, True is assumed.
        '''
        temp: List[Union[BaseData[Any],
                         CategoricalData[Categorical],
                         NumericalData[Numerical]]] = []
        _data = data if isinstance(data, (Sequence, Iterator)) else [data]
        for d in _data:
            if isinstance(d, BaseData):
                temp.append(d)
            elif isinstance(d, dict):
                name = d['name']
                value = d.get('value')
                normalizer = d.get('normalizer')
                l_e = lazy_evaluation if lazy_evaluation is not None \
                    else d.get('lazy_evaluation', True)

                if 'categories' in d:
                    temp.append(CategoricalData(
                        name=name,
                        value=value,
                        categories=d['categories'],
                        normalizer=normalizer,
                        lazy_evaluation=l_e))
                elif 'lower' in d or 'upper' in d:
                    temp.append(NumericalData(
                        name=name,
                        value=value,
                        lower=d.get('lower'),
                        upper=d.get('upper'),
                        normalizer=normalizer,
                        lazy_evaluation=l_e))
                else:
                    temp.append(BaseData(
                        name=name,
                        value=value,
                        normalizer=normalizer,
                        lazy_evaluation=l_e))
            else:
                raise TypeError(f'Unknow input type {type(d)} for item: {d}')

        self._data = tuple(temp)
        self._clear_temps()

    @classmethod
    def single_base(
            cls,
            name: str,
            value: Optional[Union[T, Sequence[T]]] = None,
            normalizer: Optional[Normalizer] = None,
            lazy_evaluation: Optional[bool] = None) -> ReilData:
        '''
        Create a `ReilData` instance.

        Arguments
        ---------
        name:
            Name of the instance.

        value:
            The value to store.

        normalizer:
            A function that accepts a `BaseData` object and returns the
            normalized value.

        lazy_evaluation:
            Whether normalization should be done at the point of instantiation
            (False) or at the access time (True).
        '''
        instance = cls([])
        instance._data = (
            BaseData(name=name,
                     value=value,
                     normalizer=normalizer,
                     lazy_evaluation=lazy_evaluation),)

        return instance

    @classmethod
    def single_categorical(
            cls,
            name: str,
            value: Optional[Union[Categorical, Sequence[Categorical]]] = None,
            normalizer: Optional[Normalizer] = None,
            lazy_evaluation: Optional[bool] = None,
            categories: Optional[Tuple[Categorical, ...]] = None) -> ReilData:
        '''
        Create a `ReilData` instance.

        Arguments
        ---------
        name:
            Name of the instance.

        value:
            The value to store.

        categories:
            A list of categories if the object is categorical

        normalizer:
            A function that accepts a `BaseData` object and returns the
            normalized value.

        lazy_evaluation:
            Whether normalization should be done at the point of instantiation
            (False) or at the access time (True).
        '''
        instance = cls([])
        instance._data = (
            CategoricalData(
                name=name,
                value=value,
                categories=categories,
                normalizer=normalizer,
                lazy_evaluation=lazy_evaluation),)

        return instance

    @classmethod
    def single_numerical(
            cls,
            name: str,
            value: Optional[Union[Numerical, Sequence[Numerical]]] = None,
            normalizer: Optional[Normalizer] = None,
            lazy_evaluation: Optional[bool] = None,
            lower: Optional[Numerical] = None,
            upper: Optional[Numerical] = None) -> ReilData:
        '''
        Create a `ReilData` instance.

        Arguments
        ---------
        name:
            Name of the instance.

        value:
            The value to store.

        lower:
            The lower bound of value if data is numerical

        upper:
            The upper bound of the value is numerical

        normalizer:
            A function that accepts a `BaseData` object and returns the
            normalized value.

        lazy_evaluation:
            Whether normalization should be done at the point of instantiation
            (False) or at the access time (True).
        '''
        instance = cls([])
        instance._data = (
            NumericalData(
                name=name,
                value=value,
                lower=lower,
                upper=upper,
                normalizer=normalizer,
                lazy_evaluation=lazy_evaluation),)

        return instance

    @classmethod
    def single_item(
            cls,
            name: str,
            value: Optional[Union[
                Numerical, Sequence[Numerical],
                Categorical, Sequence[Categorical]]] = None,
            normalizer: Optional[Normalizer] = None,
            lazy_evaluation: Optional[bool] = None,
            categories: Optional[Tuple[Categorical, ...]] = None,
            lower: Optional[Numerical] = None,
            upper: Optional[Numerical] = None) -> ReilData:
        '''
        Create a `ReilData` instance.

        Arguments
        ---------
        name:
            Name of the instance.

        value:
            The value to store.

        categories:
            A list of categories if the object is categorical

        lower:
            The lower bound of value if data is numerical

        upper:
            The upper bound of the value is numerical

        normalizer:
            A function that accepts a `BaseData` object and returns the
            normalized value.

        lazy_evaluation:
            Whether normalization should be done at the point of instantiation
            (False) or at the access time (True).
        '''
        instance = cls([])
        if lower is not None or upper is not None:
            instance._data = (
                NumericalData(
                    name=name,
                    value=value,  # type: ignore
                    lower=lower,
                    upper=upper,
                    normalizer=normalizer,
                    lazy_evaluation=lazy_evaluation),)
        elif categories is not None:
            instance._data = (
                CategoricalData(
                    name=name,
                    value=value,
                    categories=categories,
                    normalizer=normalizer,
                    lazy_evaluation=lazy_evaluation),)
        else:
            instance._data = (
                CategoricalData(
                    name=name,
                    value=value,
                    categories=categories,
                    normalizer=normalizer,
                    lazy_evaluation=lazy_evaluation),)

        return instance

    def _clear_temps(self):
        self._value: Optional[Dict[str, Any]] = None
        self._lower: Optional[Dict[str, Numerical]] = None
        self._upper: Optional[Dict[str, Numerical]] = None
        self._categories: Optional[Dict[str, Tuple[Categorical, ...]]] = None
        self._is_numerical: Optional[Dict[str, bool]] = None

    def index(self, value: Any,
              start: int = 0, stop: Optional[int] = None) -> int:
        _stop = stop if stop is not None else len(self._data)
        if isinstance(value, BaseData):
            for i in range(start, _stop):
                if self._data[i] == value:
                    return i

            raise ValueError(f'{value} is not on the list.')

        elif isinstance(value, type(self._data[0].name)):
            for i in range(start, _stop):
                if self._data[i].name == value:
                    return i

            raise ValueError(f'{value} is not on the list.')

        else:
            raise ValueError(f'{value} is not on the list.')

    @property
    def value(self) -> Dict[str, Any]:
        '''
        Return a dictionary with elements' names as keys and
        their respective values as values.

        Returns
        -------
        :
            Names of the elements and their values.
        '''
        if self._value is None:
            self._value = dict((v.name, v.value)
                               for v in self._data)

        return self._value

    @property
    def lower(self) -> Dict[str, Numerical]:
        '''
        Return all `lower` attributes.

        Returns
        -------
        :
            `lower` attribute of all `NumericalData` variables with their names
            as keys.
        '''
        if self._lower is None:
            self._lower = dict(
                (v.name, cast(Numerical, v.lower))  # type: ignore
                for v in self._data
                if hasattr(v, 'lower'))

        return self._lower

    @property
    def upper(self) -> Dict[str, Numerical]:
        '''
        Return all `upper` attributes.

        Returns
        -------
        :
            `upper` attribute of all `NumericalData` variables with their names
            as keys.
        '''
        if self._upper is None:
            self._upper = dict(
                (v.name, cast(Numerical, v.upper))  # type: ignore
                for v in self._data
                if hasattr(v, 'upper'))

        return self._upper

    @property
    def categories(self) -> Dict[str, Tuple[Categorical, ...]]:
        '''
        Return all `categories` attributes.

        Returns
        -------
        :
            `categories` attribute of all `CategoricalData` variables with
            their names as keys.
        '''
        if self._categories is None:
            self._categories = dict(
                (v.name, cast(Categorical, v.categories))  # type: ignore
                for v in self._data
                if hasattr(v, 'categories'))

        return self._categories

    @property
    def normalized(self) -> ReilData:
        '''
        Normalize all items in the instance.

        Returns
        -------
        :
            A `ReilData` consist of normalized values of all the items in the
            instance, in the form of `NumericalData` objects.
        '''
        return ReilData(
            NumericalData(name=v.name, value=v.normalized,
                          lower=0, upper=1, lazy_evaluation=True)
            for v in self._data)

    def flatten(self) -> List[BaseData]:
        """Combine values of all items in the instance.

        Returns
        -------
        :
            A list that contains all the values of all the items.
        """
        def make_iterable(x: T) -> Iterable[T]:
            return x if isinstance(x, Iterable) else [x]

        return list(itertools.chain(*[make_iterable(sublist)
                                      for sublist in self.value.values()]))

    def split(self) -> Union[ReilData, List[ReilData]]:
        """Split the `ReilData` into a list of `ReilData`.

        Returns
        -------
        :
            All items in the instance as separate `ReilData` instances.
        """
        if len(self) == 1:
            d = self._data[0]
            if not isinstance(d.value, (list, tuple)):
                splitted_list = ReilData(self._data)
            else:
                temp = d.as_dict()
                cls = type(d)
                value = temp['value']
                del temp['value']
                if 'is_numerical' in temp:
                    del temp['is_numerical']

                splitted_list = [
                    ReilData(cls(
                        value=v,
                        **temp))
                    for v in value]

        else:
            splitted_list = cast(List[ReilData], list(self._data))

        return splitted_list

    @overload
    def __getitem__(self, i: int) -> BaseData:
        ...

    @overload
    def __getitem__(self, s: slice) -> Tuple[BaseData, ...]:
        ...

    def __getitem__(self, x):
        return self._data.__getitem__(x)

    def __len__(self) -> int:
        return self._data.__len__()

    def __hash__(self) -> int:
        return hash(tuple(self._data))

    def __eq__(self, other) -> bool:
        return isinstance(other, type(self)) and (self._data == other._data)

    def __add__(self, other: Union[BaseData[T], ReilData]) -> ReilData:
        if isinstance(other, BaseData):
            if other.name in self._data:
                raise ValueError(
                    'Cannot have items with same names.'
                    ' Use update() if you need to update an item.')
            new_data = [other]
        elif isinstance(other, ReilData):
            for k in other:
                if k in self._data:
                    raise ValueError(
                        'Cannot have items with same names.'
                        ' Use update() if you need to update an item.')
            new_data = other._data
        else:
            raise TypeError(
                'Concatenation of type ReilData'
                f' and {type(other)} not implemented!')

        return ReilData(itertools.chain(self._data, new_data))

    def __neg__(self) -> ReilData:
        temp = [v.as_dict()
                for v in self._data]
        for item in temp:
            item['value'] = -item['value']

        return ReilData(temp)

    def __repr__(self) -> str:
        return f'[{super().__repr__()} -> {self._data}]'

    def __str__(self) -> str:
        return f"[{', '.join((d.__str__() for d in self._data))}]"


if __name__ == "__main__":
    # from timeit import timeit

    def f1():
        t1 = ReilData.single_categorical(name='test A', value=['a', 'b'],
                                         categories=('a', 'b'))
        t1 = ReilData.single_numerical(name='test A', value=[1, 2])
        t1 = ReilData.single_base(name='test A', value=['a', 'b'])
        return t1

    def f2():
        t2 = ReilData(({'name': x, 'value': x,
                        'categories': list('abcdefghijklmnopqrstuvwxyz')}
                       for x in 'abcdefghijklmnopqrstuvwxyz'),
                      lazy_evaluation=True)
        return t2.normalized.flatten()

    def f3():
        t2 = ReilData([{'name': 'A', 'value': 'a', 'categories': ['a', 'b']},
                       {'name': 'B', 'value': [10, 20],
                        'lower': 0, 'upper': 100}], lazy_evaluation=True)
        return t2

    # t1 = f1()
    # t2 = ReilData(t1)
    print(f1() + f3())
    # test = f1().split()
    # print(test)
    # print(timeit(f1, number=1000))
    # print(timeit(f2, number=100))
    # print(timeit(f3, number=1000))
    # print(f2().normalized.as_list())
    # print(a.values)
    # print(a.lower)
    # print(a.categories)
    # a.values = {'test A':'b'}
