#! /usr/bin/env python3

import torch

import os
import pdb
from util.args_loader import get_args
from util.data_loader import get_loader_in, get_loader_out
from util.model_loader import get_model
import numpy as np
import torch.nn.functional as F

def vit_features(model, x):
    b, c, fh, fw = x.shape
    x = model.patch_embedding(x)  # b,d,gh,gw
    x = x.flatten(2).transpose(1, 2)  # b,gh*gw,d
    if hasattr(model, 'class_token'):
        x = torch.cat((model.class_token.expand(b, -1, -1), x), dim=1)  # b,gh*gw+1,d
    if hasattr(model, 'positional_embedding'): 
        x = model.positional_embedding(x)  # b,gh*gw+1,d 
    x = model.transformer(x)  # b,gh*gw+1,d
    x = model.norm(x)[:, 0]  # b,d
    return x


torch.manual_seed(1)
torch.cuda.manual_seed(1)
np.random.seed(1)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

args = get_args()

os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu

loader_in_dict = get_loader_in(args, config_type="eval", split=('train', 'val'))
trainloaderIn, testloaderIn, num_classes = loader_in_dict.train_loader, loader_in_dict.val_loader, loader_in_dict.num_classes
model = get_model(args, num_classes, load_ckpt=True) # set true to load from ash_ckpt in their repo; essentially the same as torch ckpt
model.to(device)


batch_size = args.batch_size
featdim = {
    'resnet50': 2048,
    'resnet50-supcon': 2048,
    'vit': 768,
}[args.model_arch]

FORCE_RUN = True
ID_RUN = True
OOD_RUN = True

if ID_RUN:
    for split, in_loader in [('val', testloaderIn), ('train', trainloaderIn)]:
        cache_dir = f"cache/{args.in_dataset}_{split}_{args.name}_in"
        if FORCE_RUN or not os.path.exists(cache_dir):
            os.makedirs(cache_dir, exist_ok=True)
            feat_log = np.memmap(f"{cache_dir}/feat.mmap", dtype=float, mode='w+', shape=(len(in_loader.dataset), featdim))
            score_log = np.memmap(f"{cache_dir}/score.mmap", dtype=float, mode='w+', shape=(len(in_loader.dataset), num_classes))
            label_log = np.memmap(f"{cache_dir}/label.mmap", dtype=float, mode='w+', shape=(len(in_loader.dataset),))

            model.to(device)
            model.eval()
            with torch.no_grad():
                for batch_idx, (inputs, targets) in enumerate(in_loader):

                    inputs, targets = inputs.to(device), targets.to(device)
                    start_ind = batch_idx * batch_size
                    end_ind = min((batch_idx + 1) * batch_size, len(in_loader.dataset))

                    if args.model_arch == 'resnet50-supcon':
                        out = model.encoder(inputs)
                    elif args.model_arch == 'resnet50':
                        out = model.features(inputs)
                    elif args.model_arch == 'vit':
                        out = vit_features(model, inputs)
                    if len(out.shape) > 2:
                        out = F.adaptive_avg_pool2d(out, 1)
                        out = out.view(out.size(0), -1)
                    score = model.fc(out)
                    # score = net(inputs)
                    feat_log[start_ind:end_ind, :] = out.data.cpu().numpy()
                    label_log[start_ind:end_ind] = targets.data.cpu().numpy()
                    score_log[start_ind:end_ind] = score.data.cpu().numpy()
                    if batch_idx % 100 == 0:
                        print(f"{batch_idx}/{len(in_loader)}")
        else:
            feat_log = np.memmap(f"{cache_dir}/feat.mmap", dtype=float, mode='r', shape=(len(in_loader.dataset), featdim))
            score_log = np.memmap(f"{cache_dir}/score.mmap", dtype=float, mode='r', shape=(len(in_loader.dataset), num_classes))
            label_log = np.memmap(f"{cache_dir}/label.mmap", dtype=float, mode='r', shape=(len(in_loader.dataset),))

if OOD_RUN:

    for ood_dataset in args.out_datasets:
        loader_test_dict = get_loader_out(args, dataset=(None, ood_dataset), split=('val'))
        out_loader = loader_test_dict.val_ood_loader

        cache_dir = f"cache/{ood_dataset}vs{args.in_dataset}_{args.name}_out"
        if FORCE_RUN or not os.path.exists(cache_dir):
            os.makedirs(cache_dir, exist_ok=True)
            ood_feat_log = np.memmap(f"{cache_dir}/feat.mmap", dtype=float, mode='w+', shape=(len(out_loader.dataset), featdim))
            ood_score_log = np.memmap(f"{cache_dir}/score.mmap", dtype=float, mode='w+', shape=(len(out_loader.dataset), num_classes))
            model.eval()
            with torch.no_grad():
                for batch_idx, (inputs, _) in enumerate(out_loader):
                    inputs = inputs.to(device)
                    start_ind = batch_idx * batch_size
                    end_ind = min((batch_idx + 1) * batch_size, len(out_loader.dataset))

                    if args.model_arch == 'resnet50-supcon':
                        out = model.encoder(inputs)
                    elif args.model_arch == 'resnet50':
                        out = model.features(inputs)
                    elif args.model_arch == 'vit':
                        out = vit_features(model, inputs)
                    if len(out.shape) > 2:
                        out = F.adaptive_avg_pool2d(out, 1)
                        out = out.view(out.size(0), -1)
                    score = model.fc(out)
                    ood_feat_log[start_ind:end_ind, :] = out.data.cpu().numpy()
                    ood_score_log[start_ind:end_ind] = score.data.cpu().numpy()
                    if batch_idx % 100 == 0:
                        print(f"{batch_idx}/{len(out_loader)}")


        else:
            ood_feat_log = np.memmap(f"{cache_dir}/feat.mmap", dtype=float, mode='r', shape=(len(out_loader.dataset), featdim))
            ood_score_log = np.memmap(f"{cache_dir}/score.mmap", dtype=float, mode='r', shape=(len(out_loader.dataset), num_classes))
