# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

from __future__ import annotations

import dataclasses
import itertools
from operator import itemgetter
from typing import Optional, List

import numpy as np
from toolz import first, keymap, valmap, dissoc, merge

from gluonts_customized import maybe
from gluonts_customized.itertools import pluck_attr, rows_to_columns

from ._base import Pad
from ._period import Periods, period
from ._repr import html_table
from ._util import _replace


@dataclasses.dataclass
class SplitFrame:
    _past: dict
    _future: dict
    index: Optional[Periods]
    static: dict
    past_length: int
    future_length: int
    tdims: dict
    metadata: Optional[dict] = None
    default_tdim: int = -1
    _pad: Pad = Pad()

    def __post_init__(self):
        for column in itertools.chain(self._past, self._future):
            self.tdims.setdefault(column, self.default_tdim)

        # this triggers checks for past_length and future_length
        _, _ = self.past, self.future

    @property
    def past(self):
        return TimeFrame(
            self._past,
            index=maybe.map(
                self.index, lambda index: index[: self.past_length]
            ),
            static=self.static,
            length=self.past_length,
            tdims=self.tdims,
            metadata=self.metadata,
            _pad=Pad(self._pad.left, 0),
        )

    @property
    def future(self):
        return TimeFrame(
            self._future,
            index=maybe.map(
                self.index, lambda index: index[self.past_length :]
            ),
            static=self.static,
            length=self.future_length,
            tdims=self.tdims,
            metadata=self.metadata,
            _pad=Pad(0, self._pad.right),
        )

    def __getitem__(self, name):
        if name in self._past:
            if name in self._future:
                return np.concatenate(
                    self._past[name], self._future[name], axis=self.tdims[name]
                )
            return self._past[name]
        return self._future[name]

    def __len__(self):
        return self.past_length + self.future_length

    def set(self, name, value, tdim=None):
        tdim = maybe.unwrap_or(tdim, self.default_tdim)
        assert value.shape[tdim] == len(self)

        past, future = np.split(
            value,
            [self.past_length],
            axis=tdim,
        )

        return _replace(
            past=merge(self.past, {name: past}),
            future=merge(self.future, {name: future}),
            tdims=merge(self.tdims, {name: tdim}),
        )

    def set_like(self, ref: str, column, value, tdim=None):
        is_past = ref in self._past
        is_future = ref in self._future

        if is_past:
            if is_future:
                return self.set(column, value, tdim)
            else:
                return self.set_past(column, value, tdim)
        elif is_future:
            return self.set_future(column, value, tdim)

        raise KeyError(f"Ref {ref} is neither past nor future")

    def set_past(self, name, value, tdim=None):
        tdim = maybe.unwrap_or(tdim, self.default_tdim)
        assert value.shape[tdim] == self.past_length
        assert self.tdims.get(name, tdim) == tdim

        return _replace(
            self,
            _past=merge(self._past, {name: value}),
            tdims=merge(self.tdims, {name: tdim}),
        )

    def set_future(self, name, value, tdim=None):
        tdim = maybe.unwrap_or(tdim, self.default_tdim)
        assert value.shape[tdim] == self.future_length
        assert self.tdims.get(name, tdim) == tdim

        return _replace(
            self,
            _future=merge(self.future, {name: value}),
            tdims=merge(self.tdims, {name: tdim}),
        )

    def remove(self, column):
        return _replace(
            self,
            columns=dissoc(self.columns, column),
            tdims=dissoc(self.tdims, column),
        )

    def _repr_html_(self):
        past = self.past._table_columns()
        future = self.future._table_columns()

        length = max(
            len(first(past.values())) if past else 0,
            len(first(future.values())) if future else 0,
        )

        def pad(col):
            to_pad = length - len(col)

            return list(col) + [""] * to_pad

        past = valmap(pad, past)
        future = valmap(pad, future)

        past = keymap(lambda key: f"past_{key}" if key else "past", past)
        future = keymap(
            lambda key: f"future_{key}" if key else "future", future
        )

        return html_table({**past, "|": ["|"] * length, **future})

    def as_dict(self):
        past = keymap(lambda key: f"past_{key}", self._past)
        future = keymap(lambda key: f"future_{key}", self._future)

        return {**past, **future, **self.static}

    def resize(
        self,
        past_length: Optional[int] = None,
        future_length: Optional[int] = None,
        pad_value=0.0,
    ) -> SplitFrame:
        index = self.index

        past_length = maybe.unwrap_or(past_length, self.past_length)
        future_length = maybe.unwrap_or(future_length, self.future_length)

        if index is not None:
            # Calculate new start. If current past_length is larger than the
            # the new one, we shift it to the right, if it's smaller, we need
            # to go further into the past (shift to the left)
            start = index[0] + (self.past_length - past_length)
            index = start.periods(past_length + future_length)

        return _replace(
            self,
            _past=self.past.resize(
                past_length, pad_value, pad="l", skip="l"
            ).columns,
            past_length=maybe.unwrap_or(past_length, self.past_length),
            _future=self.future.resize(
                future_length, pad_value, pad="r", skip="r"
            ).columns,
            future_length=maybe.unwrap_or(future_length, self.future_length),
            index=index,
        )

    def with_index(self, index):
        return _replace(self, index=index)

    @staticmethod
    def _batch(split_frames: List[SplitFrame]) -> BatchSplitFrame:
        ref = split_frames[0]
        pluck = pluck_attr(split_frames)

        return BatchSplitFrame(
            _past=rows_to_columns(pluck("_past"), np.stack),  # type: ignore
            _future=rows_to_columns(pluck("_future"), np.stack),  # type: ignore
            index=pluck("index"),
            static=rows_to_columns(pluck("static"), np.stack),  # type: ignore
            past_length=ref.past_length,
            future_length=ref.future_length,
            tdims=ref.tdims,
            metadata=pluck("metadata"),
            _pad=pluck("_pad"),
        )


@dataclasses.dataclass
class BatchSplitFrame:
    _past: dict
    _future: dict
    index: List[Optional[Periods]]
    static: dict
    past_length: int
    future_length: int
    tdims: dict
    metadata: List[Optional[dict]]
    _pad: List[Pad]

    @property
    def batch_size(self):
        return len(self.index)

    def __len__(self):
        return self.past_length + self.future_length

    @property
    def past(self):
        return BatchTimeFrame(
            columns=self._past,
            static=self.static,
            length=self.past_length,
            index=[
                maybe.map(index, itemgetter(slice(None, self.past_length)))
                for index in self.index
            ],
            tdims=self.tdims,
            metadata=self.metadata,
            _pad=self._pad,
        )

    @property
    def future(self):
        return BatchTimeFrame(
            columns=self._future,
            static=self.static,
            length=self.future_length,
            index=[
                maybe.map(index, itemgetter(slice(self.past_length, None)))
                for index in self.index
            ],
            tdims=self.tdims,
            metadata=self.metadata,
            _pad=self._pad,
        )

    def items(self):
        return BatchSplitFrameItems(self)

    def as_dict(self):
        return SplitFrame.as_dict(self)


@dataclasses.dataclass(repr=False)
class BatchSplitFrameItems:
    data: BatchSplitFrame

    def __len__(self):
        return self.data.batch_size

    def __getitem__(self, idx):
        tdims = self.data.tdims

        if isinstance(idx, int):
            cls = SplitFrame
            tdims = valmap(lambda tdim: tdim - 1 if tdim >= 0 else tdim, tdims)
        else:
            cls = BatchSplitFrame

        return cls(
            _past=valmap(itemgetter(idx), self.data._past),
            _future=valmap(itemgetter(idx), self.data._future),
            index=self.data.index[idx],
            static=valmap(itemgetter(idx), self.data.static),
            tdims=tdims,
            metadata=self.data.metadata[idx],
            _pad=self.data._pad[idx],
            past_length=self.data.past_length,
            future_length=self.data.future_length,
        )


def split_frame(
    full=None,
    *,
    past=None,
    future=None,
    past_length=None,
    future_length=None,
    static=None,
    index=None,
    start=None,
    freq=None,
    tdims=None,
    metadata=None,
    default_tdim=-1,
):
    """
    Create a ``zebras.SplitFrame`` where columns can either be `past`, `future`
    or `full`, which spans both past and future.

    ``past_length`` and ``future_length`` is derived from the input data if
    possible or default to ``0`` in case no respective data is available. It is
    possible to set these values explicitly for enforcing consistency or to
    provide a length even though no time series spans that range.

    Parameters
    ----------
    full, optional
        Time series columns that span past and future.
    past, optional
        Time series columns that are past only.
    future, optional
        Time series columns that are future only.
    past_length, optional
        Set length of the past section, derived from data if not provided.
    future_length, optional
        Set length of the future section, derived from data if not provided.
    static, optional
        Values that are independent of time.
    index, optional
        A ``zebras.Periods`` object representing timestamps.
        Must have the same length as full range.
    start, optional
        The start time represented by a string (e.g., "2023-01-01"),
        or a ``zebras.Period`` object. An index will be constructed using
        this start time and the specified frequency
    freq, optional
        The frequency to use for constructing the index.
    tdims, optional
        A dictionary specifying the time dimension for each column, this
        applies to past, future and full.
    default_tdim, optional
        The default time dimension, by default -1
    metadata, optional
        A dictionary of metadata associated with the TimeFrame, by default None
    Returns
    -------
        A ``zebras.SplitFrame`` object.
    """
    full = valmap(np.array, maybe.unwrap_or_else(full, dict))
    past = valmap(np.array, maybe.unwrap_or_else(past, dict))
    future = valmap(np.array, maybe.unwrap_or_else(future, dict))
    static = valmap(np.array, maybe.unwrap_or_else(static, dict))
    tdims = maybe.unwrap_or_else(tdims, dict)

    # Resolve `past_length` and `future_length` if not directly set from
    # provided data. If no data is passed for either field, the value is still
    # `None` after this.
    if past_length is None and past:
        column = first(past)
        past_length = past[column].shape[tdims.get(column, default_tdim)]

    if future_length is None and future:
        column = first(future)
        future_length = future[column].shape[tdims.get(column, default_tdim)]

    if full:
        column = first(full)
        full_length = full[column].shape[tdims.get(column, default_tdim)]

        if maybe.or_(past_length, future_length) is None:
            raise ValueError(
                "Cannot determine past and future length if only "
                "`full` is provided."
            )
    # No data and no lengths are passed
    elif maybe.or_(past_length, future_length) is None:
        past_length = 0
        future_length = 0
        full_length = 0
    else:
        # If past_length and/or future_length is provided, but no `full` data
        # is given, then we first resolve past and future length and then just
        # calculate full_length later.
        full_length = None

    if past_length is None:
        past_length = maybe.map_or(
            full_length,
            lambda fl: fl - future_length,
            0,
        )
    elif future_length is None:
        future_length = maybe.map_or(
            full_length,
            lambda fl: fl - past_length,
            0,
        )
    full_length = past_length + future_length

    # create copies to not mutate user provided dicts
    past = dict(past)
    future = dict(future)

    for name, data in full.items():
        tdim = tdims.get(name, default_tdim)
        assert data.shape[tdim] == full_length

        past_data, future_data = np.split(data, [past_length], axis=tdim)
        past[name] = past_data
        future[name] = future_data

    sf = SplitFrame(
        _past=past,
        _future=future,
        index=index,
        static=static,
        tdims=tdims,
        past_length=past_length,
        future_length=future_length,
        default_tdim=default_tdim,
        metadata=metadata,
    )

    if sf.index is None and start is not None:
        if freq is not None:
            start = period(start, freq)

        return sf.with_index(start.periods(len(sf)))

    return sf


# We defer these imports to avoid circular imports.
from ._time_series import TimeSeries  # noqa
from ._time_frame import BatchTimeFrame, TimeFrame  # noqa
