import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from itertools import chain
from typing import Callable, Union, Type, Optional, Dict, Any

# From https://github.com/DLR-RM/rl-baselines3-zoo/blob/8ea4f4a87afa548832ca17e575b351ec5928c1b0/utils/utils.py
def linear_schedule(initial_value: Union[float, str]) -> Callable[[float], float]:
    """
    Linear learning rate schedule.
    :param initial_value: (float or str)
    :return: (function)
    """
    if isinstance(initial_value, str):
        initial_value = float(initial_value)

    def func(progress_remaining: float) -> float:
        """
        Progress will decrease from 1 (beginning) to 0
        :param progress_remaining: (float)
        :return: (float)
        """
        return progress_remaining * initial_value

    return func

class SADataset(torch.utils.data.Dataset):
    def __init__(self, obs, acts, normalize):
        if normalize:
            obs = np.array(obs)
            self.mean = obs.mean(axis=0)
            self.std = obs.std(axis=0) + 1e-3
            obs = (obs - self.mean) / (self.std)
            self.is_normalized = True
        else:
            self.is_normalized = False
        self.obs = torch.tensor(obs)
        self.acts = torch.tensor(acts)

    def __len__(self):
        return len(self.obs)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        obs = self.obs[idx]
        acts = self.acts[idx]
        sample = {'obs': obs, 'acts': acts}
        return sample

def make_sa_dataloader(envname, max_trajs=None, normalize=False, batch_size=32):
    demos = np.load(
        "../experts/{0}/demos.npz".format(envname), allow_pickle=True)
    num_trajs = demos["num_trajs"]
    if max_trajs is None:
        max_trajs = num_trajs
    obs = []
    acts = []
    for traj in range(min(max_trajs, num_trajs)):
        obs.extend(demos[str(traj)].item()['states'])
        acts.extend(demos[str(traj)].item()['actions'])
    dataset = SADataset(obs, acts, normalize)
    dataloader = DataLoader(dataset, batch_size=batch_size,
                             shuffle=True, num_workers=0)
    return dataloader

class SADSDataset(torch.utils.data.Dataset):
    def __init__(self, obs, acts, next_obs, traj_lens):
        self.obs = torch.tensor(obs)
        self.acts = torch.tensor(acts)
        self.next_obs = torch.tensor(next_obs)
        dones = [[False for _ in range(l - 2)] + [True] for l in traj_lens]
        self.dones = torch.tensor(list(chain.from_iterable(dones)))

    def __len__(self):
        return len(self.obs)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        obs = self.obs[idx]
        acts = self.acts[idx]
        next_obs = self.next_obs[idx]
        dones = self.dones[idx]
        sample = {'obs': obs, 'acts': acts,
                  'next_obs': next_obs, 'dones': dones}
        return sample

def make_sads_dataloader(envname, max_trajs=None):
    demos = np.load(
        "./experts/{0}/demos.npz".format(envname), allow_pickle=True)
    num_trajs = demos["num_trajs"]
    if max_trajs is None:
        max_trajs = num_trajs
    obs = []
    next_obs = []
    acts = []
    lens = []
    for traj in range(min(max_trajs, num_trajs)):
        obs.extend(demos[str(traj)].item()['states'][:-1])
        next_obs.extend(demos[str(traj)].item()['states'][1:])
        acts.extend(demos[str(traj)].item()['actions'][:-1])
        lens.append(len(demos[str(traj)].item()['states']))
    dataset = SADSDataset(obs, acts, next_obs, lens)
    dataloader = DataLoader(dataset, batch_size=32,
                            shuffle=False, num_workers=0, drop_last=True)
    return dataloader

def make_sa_dataset(envname, max_trajs=None):
    demos = np.load("../pillbox/experts/{0}/demos.npz".format(envname), allow_pickle=True)
    num_trajs = demos["num_trajs"]
    if max_trajs is None:
        max_trajs = num_trajs
    expert_states = []
    expert_actions = []
    expert_next_states = []
    expert_dones = []
    for traj in range(min(max_trajs, num_trajs)):
        expert_states.extend(demos[str(traj)].item()['states'][:-1])
        expert_next_states.extend(demos[str(traj)].item()['states'][1:])
        expert_actions.extend(demos[str(traj)].item()['actions'][:-1])
        l = len(demos[str(traj)].item()['states'])
        expert_dones.extend([False for _ in range(l - 2)] + [True])
    expert_data = dict()
    expert_data['obs'] = np.array(expert_states)
    expert_data['acts'] = np.array(expert_actions)
    expert_data['next_obs'] = np.array(expert_next_states)
    expert_data['dones'] = np.array(expert_dones)
    return expert_data