import copy
import logging
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
from torch.utils.data import Dataset

from mohou.encoding_rule import EncodingRule
from mohou.types import EpisodeBundle, TerminateFlag
from mohou.utils import (
    AnyT,
    assert_equal_with_message,
    assert_seq_list_list_compatible,
    flatten_lists,
)

logger = logging.getLogger(__name__)


@dataclass
class SequenceDatasetConfig:
    n_aug: int = 20
    cov_scale: float = 0.1

    def __new__(cls, *args, **kwargs):
        assert len(args) == 0, "please instantiate config only using kwargs"
        return super(SequenceDatasetConfig, cls).__new__(cls)

    def __post_init__(self):
        assert self.n_aug >= 0
        assert self.cov_scale < 1.0
        logger.info("sequence dataset config: {}".format(self))


@dataclass
class SequenceDataAugmentor:
    covmat: np.ndarray
    config: SequenceDatasetConfig

    @classmethod
    def from_seqs(
        cls, seq_list: List[np.ndarray], config: SequenceDatasetConfig, take_diff: bool = True
    ):
        """construct augmentor.
        seq_list: sequence list used to compute covariance matrix
        take_diff: if True, covaraince is computed using state difference (e.g. {s[t] - s[t-1]}_{1:T})
        otherwise covariance is computed using s[t]_{1:T} sequence.
        """
        if take_diff:
            covmat = cls.compute_diff_covariance(seq_list)
        else:
            covmat = cls.compute_covariance(seq_list)
        return cls(covmat, config)

    @staticmethod
    def compute_covariance(state_seq_list: List[np.ndarray]) -> np.ndarray:
        state_diffs = np.vstack(state_seq_list)
        cov_mat = np.cov(state_diffs.T)
        return cov_mat

    @staticmethod
    def compute_diff_covariance(state_seq_list: List[np.ndarray]) -> np.ndarray:
        state_diff_list = []
        for state_seq in state_seq_list:
            diff = state_seq[1:, :] - state_seq[:-1, :]
            state_diff_list.append(diff)
        state_diffs = np.vstack(state_diff_list)
        cov_mat = np.cov(state_diffs.T)
        return cov_mat

    def apply(self, seq: np.ndarray) -> List[np.ndarray]:
        """apply augmentation
        seq: 2dim array (n_seqlen, n_dim)
        """
        assert seq.ndim == 2
        assert seq.shape[1] == self.covmat.shape[0]

        covmat_scaled = self.covmat * self.config.cov_scale**2

        seq_original = copy.deepcopy(seq)
        auged_seq_list: List[np.ndarray] = [seq_original]
        for _ in range(self.config.n_aug):
            n_seqlen, n_dim = seq_original.shape
            mean = np.zeros(n_dim)
            noises = np.random.multivariate_normal(mean, covmat_scaled, n_seqlen)
            noised_seq = seq_original + noises
            auged_seq_list.append(noised_seq)

        return auged_seq_list


@dataclass
class PaddingSequenceAligner:
    n_seqlen_target: int

    @classmethod
    def from_seqs(
        cls, seqs: Union[List[np.ndarray], List[List]], n_after_termination: int
    ) -> "PaddingSequenceAligner":
        n_seqlen_max = max([len(seq) for seq in seqs])
        n_seqlen_target = n_seqlen_max + n_after_termination
        return cls(n_seqlen_target)

    def apply(self, seq: AnyT) -> AnyT:  # TODO: specify type
        """get padded sequence based on seq.
        seq: np.ndarray or List
        if seq is [1, 2, 3, 4] and n_seqlen_target = 10, the output padded sequence looks like
        [1, 2, 3, 4, 4, 4, 4, 4, 4, 4]
        """
        seq = copy.deepcopy(seq)
        n_seqlen = len(seq)
        n_padding = self.n_seqlen_target - n_seqlen
        assert n_padding >= 0

        # TODO(HiroIshida) To remove type-ignores it's better to use singledispatch .
        # however we prefere simpilicy in this case.
        if isinstance(seq, list):
            padded_seq = seq + [copy.deepcopy(seq[-1]) for _ in range(n_padding)]
        elif isinstance(seq, np.ndarray):
            if n_padding == 0:
                padded_seq = seq  # type: ignore
            else:
                elem_last = seq[-1]
                padding_seq_shape = [n_padding] + [1 for _ in range(elem_last.ndim)]
                padding_seq = np.tile(elem_last, padding_seq_shape)
                padded_seq = np.concatenate((seq, padding_seq), axis=0)
        else:
            assert False

        assert len(padded_seq) == self.n_seqlen_target
        return padded_seq  # type: ignore


@dataclass
class AutoRegressiveDatasetConfig(SequenceDatasetConfig):
    n_dummy_after_termination: int = 20


@dataclass
class AutoRegressiveDataset(Dataset):
    state_seq_list: List[np.ndarray]  # with flag info
    episode_index_list: List[int]
    static_context_list: List[np.ndarray]
    encoding_rule: EncodingRule

    def __post_init__(self):
        assert_equal_with_message(
            len(self.static_context_list), len(self.state_seq_list), "length of sequence"
        )

        # state sequence consists of n_seqlen x n_dim
        assert_equal_with_message(self.state_seq_list[0].ndim, 2, "dimension of state sequence")
        # context sequence consists of n_dim (because its static througout the context list)
        assert_equal_with_message(self.static_context_list[0].ndim, 1, "dimension of context")

    def __len__(self) -> int:
        return len(self.state_seq_list)

    def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """return tuple of episode index of the state_seq, state_esq, and static context"""
        state_seq = torch.from_numpy(self.state_seq_list[idx]).float()
        context = torch.from_numpy(self.static_context_list[idx]).float()
        episode_index = torch.tensor(self.episode_index_list[idx], dtype=torch.int32)
        return episode_index, state_seq, context

    @classmethod
    def from_bundle(
        cls,
        bundle: EpisodeBundle,
        encoding_rule: EncodingRule,
        augconfig: Optional[AutoRegressiveDatasetConfig] = None,
        static_context_list: Optional[List[np.ndarray]] = None,
    ) -> "AutoRegressiveDataset":

        last_key_of_rule = list(encoding_rule.keys())[-1]
        assert last_key_of_rule == TerminateFlag
        if augconfig is None:
            augconfig = AutoRegressiveDatasetConfig()

        state_seq_list = encoding_rule.apply_to_episode_bundle(bundle)

        # setting up biases
        if static_context_list is None:  # create sequence of 0-dim vector
            static_context_list = [np.zeros((0)) for _ in range(len(state_seq_list))]
        assert_equal_with_message(
            len(static_context_list), len(state_seq_list), "length of sequence"
        )

        # augmentation
        augmentor = SequenceDataAugmentor.from_seqs(state_seq_list, augconfig)

        episode_indices_list_auged = []
        state_seq_list_auged = []
        for index, seq in enumerate(state_seq_list):
            seq_auged = augmentor.apply(seq)
            episode_indices_list_auged.extend([index] * len(seq_auged))
            state_seq_list_auged.extend(seq_auged)

        static_context_list_auged = flatten_lists(
            [[copy.deepcopy(c) for _ in range(augconfig.n_aug + 1)] for c in static_context_list]
        )  # +1 for original data

        # make all sequence to the same length due to torch batch computation requirement
        aligner = PaddingSequenceAligner.from_seqs(
            state_seq_list_auged, augconfig.n_dummy_after_termination
        )
        state_seq_list_auged_adjusted = [aligner.apply(seq) for seq in state_seq_list_auged]

        return cls(
            state_seq_list_auged_adjusted,
            episode_indices_list_auged,
            static_context_list_auged,
            encoding_rule,
        )


@dataclass
class MarkovControlSystemDataset(Dataset):
    """o_{t+1} = f(o_{t}, u_t{t})"""

    inp_ctrl_seq: np.ndarray
    inp_obs_seq: np.ndarray
    out_obs_seq: np.ndarray

    def __len__(self) -> int:
        return len(self.inp_ctrl_seq)

    def __getitem__(self, idx) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        inp_ctrl = torch.from_numpy(self.inp_ctrl_seq[idx]).float()
        inp_obs = torch.from_numpy(self.inp_obs_seq[idx]).float()
        out_obs = torch.from_numpy(self.out_obs_seq[idx]).float()
        return inp_ctrl, inp_obs, out_obs

    @classmethod
    def from_bundle(
        cls,
        bundle: EpisodeBundle,
        control_encoding_rule: EncodingRule,
        observation_encoding_rule: EncodingRule,
        config: Optional[SequenceDatasetConfig] = None,
        diff_as_control: bool = True,
    ) -> "MarkovControlSystemDataset":
        if config is None:
            config = SequenceDatasetConfig()

        ctrl_seq_list = control_encoding_rule.apply_to_episode_bundle(bundle)
        obs_seq_list = observation_encoding_rule.apply_to_episode_bundle(bundle)
        assert_seq_list_list_compatible([ctrl_seq_list, obs_seq_list])

        ctrl_augmentor = SequenceDataAugmentor.from_seqs(ctrl_seq_list, config, take_diff=False)
        obs_augmentor = SequenceDataAugmentor.from_seqs(obs_seq_list, config, take_diff=True)

        ctrl_seq_list_auged = flatten_lists([ctrl_augmentor.apply(seq) for seq in ctrl_seq_list])
        obs_seq_list_auged = flatten_lists([obs_augmentor.apply(seq) for seq in obs_seq_list])

        inp_ctrl_seq = []
        inp_obs_seq = []
        out_obs_seq = []
        for i in range(len(ctrl_seq_list_auged)):
            ctrl_seq = ctrl_seq_list_auged[i]
            obs_seq = obs_seq_list_auged[i]
            for j in range(len(ctrl_seq) - 1):
                if diff_as_control:
                    inp_ctrl_seq.append(ctrl_seq[j + 1] - ctrl_seq[j])
                else:
                    inp_ctrl_seq.append(ctrl_seq[j])
                inp_obs_seq.append(obs_seq[j])
                out_obs_seq.append(obs_seq[j + 1])

        return cls(np.array(inp_ctrl_seq), np.array(inp_obs_seq), np.array(out_obs_seq))
