from __future__ import annotations

from abc import ABC
from typing import Optional, List, Union, Any, TypeVar, Type

from pandas import DataFrame
from pymeos_cffi import *

from ..temporal.temporal import Temporal

TBase = TypeVar('TBase')
TG = TypeVar('TG', bound='Temporal[Any]')
TI = TypeVar('TI', bound='TInstant[Any]')
TS = TypeVar('TS', bound='TSequence[Any]')
TSS = TypeVar('TSS', bound='TSequenceSet[Any]')
Self = TypeVar('Self', bound='TSequenceSet[Any]')


class TSequenceSet(Temporal[TBase, TG, TI, TS, TSS], ABC):
    """
    Base class for temporal sequence set types, i.e. temporal values that are
    defined by a set of temporal sequences.
    """

    # ------------------------- Constructors ----------------------------------
    def __init__(self, string: Optional[str] = None, *,
                 sequence_list: Optional[List[Union[str, Any]]] = None,
                 normalize: bool = True, _inner=None):
        assert (_inner is not None) or ((string is not None) != (sequence_list is not None)), \
            "Either string must be not None or sequence_list must be not"
        if _inner is not None:
            self._inner = as_tsequenceset(_inner)
        elif string is not None:
            self._inner = as_tsequenceset(self.__class__._parse_function(string))
        else:
            sequences = [x._inner if isinstance(x, self.ComponentClass) \
                             else self.__class__._parse_function(x) for x in sequence_list]
            count = len(sequences)
            self._inner = tsequenceset_make(sequences, count, normalize)

    @classmethod
    def from_sequences(cls: Type[Self],
                       sequence_list: Optional[List[Union[str, Any]]] = None,
                       normalize: bool = True) -> Self:
        """
        Create a temporal sequence set from a list of sequences.

        Args:
            sequence_list: List of sequences.
            normalize: Whether to normalize the temporal sequence set.

        Returns:
            A temporal sequence set.
        """
        return cls(sequence_list=sequence_list, normalize=normalize)

    # ------------------------- Accessors -------------------------------------
    def num_sequences(self) -> int:
        """
        Returns the number of sequences in ``self``.
        """
        return temporal_num_sequences(self._inner)

    def start_sequence(self) -> TS:
        """
        Returns the first sequence in ``self``.
        """
        return self.ComponentClass(_inner=temporal_start_sequence(self._inner))

    def end_sequence(self) -> TS:
        """
        Returns the last sequence in ``self``.
        """
        return self.ComponentClass(_inner=temporal_end_sequence(self._inner))

    def sequence_n(self, n) -> TS:
        """
        Returns the ``n``-th sequence in ``self``.
        """
        return self.ComponentClass(_inner=temporal_sequence_n(self._inner, n + 1))

    def sequences(self) -> List[TS]:
        """
        Returns the list of sequences in ``self``.
        """
        ss, count = temporal_sequences(self._inner)
        return [self.ComponentClass(_inner=ss[i]) for i in range(count)]

    # ------------------------- Transformations -------------------------------
    def to_dataframe(self) -> DataFrame:
        """
        Returns a pandas DataFrame representation of ``self``.
        """
        sequences = self.sequences()
        data = {
            'sequence': [i for i, seq in enumerate(sequences) \
                         for _ in range(seq.num_instants())],
            'time': [t for seq in sequences for t in seq.timestamps()],
            'value': [v for seq in sequences for v in seq.values()]
        }
        return DataFrame(data).set_index(keys=['sequence', 'time'])

    # ------------------------- Plot Operations -------------------------------
    def plot(self, *args, **kwargs):
        """
        Plot the temporal sequence set.

        See Also:
            :meth:`pymeos.plotters.TemporalSequenceSetPlotter.plot`
        """
        from ..plotters import TemporalSequenceSetPlotter
        return TemporalSequenceSetPlotter.plot(self, *args, **kwargs)
