#%%
import gc
from pathlib import Path
from argparse import ArgumentParser

import numpy as np
import torch
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed

from tqdm import tqdm

import sys
sys.path.append('.')
from experiments.imagenet_discrete import test_loaders, load_model, parse_args, show
from experiments.imagenet_ddu import get_train_loader


if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")


def embeddings(loader, model):
    # switch to evaluate mode
    model.eval()
    logits = []
    targets = []

    with torch.no_grad():
        for i, (images, target) in tqdm(enumerate(loader)):
            images = images.to(device)

            # compute output
            output = model(images)
            logits.append(output.cpu().detach().numpy())
            targets.extend(target.cpu().detach().numpy())

            # if i == 1:
            #     break
    return np.concatenate(logits), np.array(targets)


def main():
    args = parse_args()
    val_loader, ood_loader = test_loaders(args.data_folder, args.ood_folder, args.b)
    train_loader = get_train_loader(batch_size=args.b)

    args.dir_to_save = Path('checkpoint')

    model = load_model(args.net)

    def save(arr, name):
        with open(args.dir_to_save / f"{name}_{args.ood_name}.npy", 'wb') as f:
            np.save(f, arr)

    train_embeddings, train_targets = embeddings(train_loader, model)
    save(train_embeddings, 'train_embeddings')
    save(train_targets, 'train_targets')
    del train_embeddings
    gc.collect()

    ood_embeddings, _ = embeddings(ood_loader, model)
    save(ood_embeddings, 'ood_embeddings')
    del ood_embeddings
    gc.collect()

    val_embeddings, val_targets = embeddings(val_loader, model)
    save(val_embeddings, 'val_embeddings')
    save(val_targets, 'val_targets')


if __name__ == '__main__':
    main()
