import random
from copy import deepcopy
import numpy as np
import torch
from torch.optim import Adam
from numpy import linalg as LA
import gym
import d4rl
import argparse
import json
from utils import redirect_stdout
import torch.nn as nn
import itertools
import torch.nn.functional as F
from torch.distributions.normal import Normal

def load_dataset(env, traj_type, sample_type, pref_type):
    npzfile = np.load('../dataset/%s_%s_%s_%s.npz' %(env, traj_type, sample_type, pref_type))
    dataset = [npzfile['traj_obs'], npzfile['traj_act'], npzfile['traj_rew'], npzfile['traj_idx_1'], npzfile['traj_idx_2'], npzfile['pref'], traj_type]
    return dataset

def mlp(sizes, activation, output_activation=nn.Identity):
    layers = []
    for j in range(len(sizes)-1):
        act = activation if j < len(sizes)-2 else output_activation
        layers += [nn.Linear(sizes[j], sizes[j+1]), act()]
    return nn.Sequential(*layers).to(torch.device('cuda'))
