import argparse
import torch
import utils
import datetime
import os
import pickle

import numpy as np
import matplotlib.pyplot as plt
import logging
from train_sim import make_pairwise_encoder

from torch.utils import data
import torch.nn.functional as F

import modules


parser = argparse.ArgumentParser()
parser.add_argument('--save-folder', type=str,
                    default='checkpoints',
                    help='Path to checkpoints.')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='Disable CUDA training.')

args_eval = parser.parse_args()


meta_file = os.path.join(args_eval.save_folder, 'metadata.pkl')
model_file = os.path.join(args_eval.save_folder, 'model.pt')
decoder_model_file = os.path.join(args_eval.save_folder, 'decoder.pt')

args = pickle.load(open(meta_file, 'rb'))['args']

args.cuda = not args_eval.no_cuda and torch.cuda.is_available()

np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

device = torch.device('cuda' if args.cuda else 'cpu')

dataset = utils.StateTransitionsDataset(
    hdf5_file=args.dataset)


def worker_init_fc(worker_index):
    # make sure each worker has a different starting random seed
    seed = np.random.get_state()[1][0]
    np.random.seed(seed + worker_index)


train_loader = data.DataLoader(
    dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, worker_init_fn=worker_init_fc
)

# Get data sample
obs = train_loader.__iter__().next()[0]
input_shape = obs[0].size()

# maybe load bisim metric and turn it into torch tensor on the selected device
bisim_metric = None
if args.bisim_metric_path is not None:
    bisim_metric = torch.tensor(np.load(args.bisim_metric_path), dtype=torch.float32, device=device)

# maybe load bisim model
bisim_model = None
if args.bisim_model_path is not None:
    bisim_model = make_pairwise_encoder()
    bisim_model.load_state_dict(torch.load(args.bisim_model_path))

model = modules.ContrastiveSWM(
    embedding_dim=args.embedding_dim,
    hidden_dim=args.hidden_dim,
    action_dim=args.action_dim,
    input_dims=input_shape,
    num_objects=args.num_objects,
    sigma=args.sigma,
    hinge=args.hinge,
    ignore_action=args.ignore_action,
    copy_action=args.copy_action,
    split_mlp=args.split_mlp,
    same_ep_neg=args.same_ep_neg,
    only_same_ep_neg=args.only_same_ep_neg,
    immovable_bit=args.immovable_bit,
    split_gnn=args.split_gnn,
    no_loss_first_two=args.no_loss_first_two,
    gamma=args.gamma,
    bisim_metric=bisim_metric,
    bisim_eps=args.bisim_eps,
    next_state_neg=args.next_state_neg,
    nl_type=args.nl_type,
    encoder=args.encoder,
    use_coord_grid=args.coord_grid).to(device)

model.load_state_dict(torch.load(model_file, map_location=device))

# avoid re-initializing the model and adding to the trainable parameters list
if bisim_model is not None:
    model.bisim_model = bisim_model
    model.bisim_model.to(device)

if args.encoder == 'large':
    decoder = modules.DecoderCNNLarge(
        input_dim=args.embedding_dim,
        num_objects=args.num_objects,
        hidden_dim=args.hidden_dim // 16,
        output_size=input_shape).to(device)
elif args.encoder == 'medium':
    decoder = modules.DecoderCNNMedium(
        input_dim=args.embedding_dim,
        num_objects=args.num_objects,
        hidden_dim=args.hidden_dim // 16,
        output_size=input_shape).to(device)
elif args.encoder == 'small':
    decoder = modules.DecoderCNNSmall(
        input_dim=args.embedding_dim,
        num_objects=args.num_objects,
        hidden_dim=args.hidden_dim // 16,
        output_size=input_shape).to(device)
decoder.apply(utils.weights_init)
optimizer_dec = torch.optim.Adam(
    decoder.parameters(),
    lr=args.learning_rate)


# Train model.
print('Starting model training...')
step = 0
best_loss = 1e9
losses = []

for epoch in range(1, args.epochs + 1):
    model.train()
    train_loss = 0

    # make sure each epoch starts with a different random seed
    # we only need numpy random seed, we use it to make random decision in custom-neg workers
    np.random.seed(args.seed + epoch)

    for batch_idx, data_batch in enumerate(train_loader):
        data_batch = [tensor.to(device) for tensor in data_batch]

        optimizer_dec.zero_grad()
        obs, _, _ = data_batch
        objs = model.obj_extractor(obs)
        state = model.obj_encoder(objs)

        rec = torch.sigmoid(decoder(state))
        loss = F.binary_cross_entropy(
            rec, obs, reduction='sum') / obs.size(0)

        loss.backward()
        train_loss += loss.item()
        optimizer_dec.step()

        losses.append(loss.item())

        if batch_idx % args.log_interval == 0:
            print(
                'Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data_batch[0]),
                    len(train_loader.dataset),
                    100. * batch_idx / len(train_loader),
                    loss.item() / len(data_batch[0])))

        step += 1

    avg_loss = train_loss / len(train_loader.dataset)
    print('====> Epoch: {} Average loss: {:.6f}'.format(
        epoch, avg_loss))

    if avg_loss < best_loss:
        best_loss = avg_loss
        torch.save(decoder.state_dict(), decoder_model_file)

import argparse
import torch
import utils
import datetime
import os
import pickle

import numpy as np
import matplotlib.pyplot as plt
import logging
from train_sim import make_pairwise_encoder

from torch.utils import data
import torch.nn.functional as F

import modules


parser = argparse.ArgumentParser()
parser.add_argument('--save-folder', type=str,
                    default='checkpoints',
                    help='Path to checkpoints.')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='Disable CUDA training.')

args_eval = parser.parse_args()


meta_file = os.path.join(args_eval.save_folder, 'metadata.pkl')
model_file = os.path.join(args_eval.save_folder, 'model.pt')
decoder_model_file = os.path.join(args_eval.save_folder, 'decoder.pt')

args = pickle.load(open(meta_file, 'rb'))['args']

args.cuda = not args_eval.no_cuda and torch.cuda.is_available()

np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

device = torch.device('cuda' if args.cuda else 'cpu')

dataset = utils.StateTransitionsDataset(
    hdf5_file=args.dataset)


def worker_init_fc(worker_index):
    # make sure each worker has a different starting random seed
    seed = np.random.get_state()[1][0]
    np.random.seed(seed + worker_index)


train_loader = data.DataLoader(
    dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, worker_init_fn=worker_init_fc
)

# Get data sample
obs = train_loader.__iter__().next()[0]
input_shape = obs[0].size()

# maybe load bisim metric and turn it into torch tensor on the selected device
bisim_metric = None
if args.bisim_metric_path is not None:
    bisim_metric = torch.tensor(np.load(args.bisim_metric_path), dtype=torch.float32, device=device)

# maybe load bisim model
bisim_model = None
if args.bisim_model_path is not None:
    bisim_model = make_pairwise_encoder()
    bisim_model.load_state_dict(torch.load(args.bisim_model_path))

model = modules.ContrastiveSWM(
    embedding_dim=args.embedding_dim,
    hidden_dim=args.hidden_dim,
    action_dim=args.action_dim,
    input_dims=input_shape,
    num_objects=args.num_objects,
    sigma=args.sigma,
    hinge=args.hinge,
    ignore_action=args.ignore_action,
    copy_action=args.copy_action,
    split_mlp=args.split_mlp,
    same_ep_neg=args.same_ep_neg,
    only_same_ep_neg=args.only_same_ep_neg,
    immovable_bit=args.immovable_bit,
    split_gnn=args.split_gnn,
    no_loss_first_two=args.no_loss_first_two,
    gamma=args.gamma,
    bisim_metric=bisim_metric,
    bisim_eps=args.bisim_eps,
    next_state_neg=args.next_state_neg,
    nl_type=args.nl_type,
    encoder=args.encoder,
    use_coord_grid=args.coord_grid).to(device)

model.load_state_dict(torch.load(model_file, map_location=device))

# avoid re-initializing the model and adding to the trainable parameters list
if bisim_model is not None:
    model.bisim_model = bisim_model
    model.bisim_model.to(device)

if args.encoder == 'large':
    decoder = modules.DecoderCNNLarge(
        input_dim=args.embedding_dim,
        num_objects=args.num_objects,
        hidden_dim=args.hidden_dim // 16,
        output_size=input_shape).to(device)
elif args.encoder == 'medium':
    decoder = modules.DecoderCNNMedium(
        input_dim=args.embedding_dim,
        num_objects=args.num_objects,
        hidden_dim=args.hidden_dim // 16,
        output_size=input_shape).to(device)
elif args.encoder == 'small':
    decoder = modules.DecoderCNNSmall(
        input_dim=args.embedding_dim,
        num_objects=args.num_objects,
        hidden_dim=args.hidden_dim // 16,
        output_size=input_shape).to(device)
decoder.apply(utils.weights_init)
optimizer_dec = torch.optim.Adam(
    decoder.parameters(),
    lr=args.learning_rate)


# Train model.
print('Starting model training...')
step = 0
best_loss = 1e9
losses = []

for epoch in range(1, args.epochs + 1):
    model.train()
    train_loss = 0

    # make sure each epoch starts with a different random seed
    # we only need numpy random seed, we use it to make random decision in custom-neg workers
    np.random.seed(args.seed + epoch)

    for batch_idx, data_batch in enumerate(train_loader):
        data_batch = [tensor.to(device) for tensor in data_batch]

        optimizer_dec.zero_grad()
        obs, _, _ = data_batch
        objs = model.obj_extractor(obs)
        state = model.obj_encoder(objs)

        rec = torch.sigmoid(decoder(state))
        loss = F.binary_cross_entropy(
            rec, obs, reduction='sum') / obs.size(0)

        loss.backward()
        train_loss += loss.item()
        optimizer_dec.step()

        losses.append(loss.item())

        if batch_idx % args.log_interval == 0:
            print(
                'Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    epoch, batch_idx * len(data_batch[0]),
                    len(train_loader.dataset),
                    100. * batch_idx / len(train_loader),
                    loss.item() / len(data_batch[0])))

        step += 1

    avg_loss = train_loss / len(train_loader.dataset)
    print('====> Epoch: {} Average loss: {:.6f}'.format(
        epoch, avg_loss))

    if avg_loss < best_loss:
        best_loss = avg_loss
        torch.save(decoder.state_dict(), decoder_model_file)

# make sure each worker has a different starting random seed # Get data sample # maybe load bisim metric and turn it into torch tensor on the selected device # maybe load bisim model # avoid re-initializing the model and adding to the trainable parameters list # Train model. # make sure each epoch starts with a different random seed # we only need numpy random seed, we use it to make random decision in custom-neg workers