# evaluate a smoothed classifier on a dataset
import argparse
import os
from datasets import get_dataset, DATASETS, get_num_classes
from core import Smooth
from time import time
import torch
import datetime
from architectures import get_architecture
from train_utils import setup_seed
import numpy as np

parser = argparse.ArgumentParser(description='Certify many examples')
parser.add_argument("--dataset", default="AWA",choices=DATASETS, help="which dataset")
parser.add_argument("--base_classifier", type=str, default='main_model_sd025/checkpoint.pth.tar', help="path to saved pytorch model of base classifier")
parser.add_argument('--model', type=str, default = "reasoning_model/", help='folder to save model and training log)')
parser.add_argument("--sigma", type=float,default=0.25, help="noise hyperparameter")
parser.add_argument("--outfile", type=str, default = 'predict_results/pipeline_certify_025_new.txt', help="output file")
parser.add_argument("--batch", type=int, default=1667, help="batch size")
parser.add_argument("--skip", type=int, default=1, help="how many examples to skip")
parser.add_argument("--max", type=int, default= 500, help="stop after this many examples")
parser.add_argument("--split", choices=["train", "test"], default="test", help="train or test set")
parser.add_argument("--N0", type=int, default=100)
parser.add_argument("--N", type=int, default=10000, help="number of samples to use")
parser.add_argument("--alpha", type=float, default=0.001, help="failure probability")
parser.add_argument('--card', action='store_true', default=False)
args = parser.parse_args()

if __name__ == "__main__":
    # load the base classifier

    if args.card:
        print("card mode")
        smoothed_classifier = Smooth(args.dataset, args.base_classifier, get_num_classes(args.dataset), args.sigma, args.card, args.model)
    else:
        checkpoint = torch.load(args.base_classifier)
        base_classifier = get_architecture(checkpoint["arch"], args.dataset)
        base_classifier.load_state_dict(checkpoint['state_dict'])

        # create the smooothed classifier g
        smoothed_classifier = Smooth(base_classifier, get_num_classes(args.dataset), args.sigma, args.card)

    # prepare output file
    f = open(args.outfile, 'w')
    print("idx\tlabel\tpredict\tradius\tcorrect\ttime", file=f, flush=True)

    # iterate through the dataset
    dataset = get_dataset(args.dataset, args.split)
    index = []

    setup_seed(0)
    ## randomly sample 10 images from each leaf node.
    for cls in range( get_num_classes(args.dataset)):
        choice = np.random.choice(torch.nonzero(torch.IntTensor(dataset.targets) == cls, as_tuple=False).squeeze(1).numpy(),10, replace =False).tolist()
        index.extend(choice)
    
    for i in index:

        # only certify every args.skip examples, and stop after args.max examples
        if i % args.skip != 0:
            continue

        (x, label) = dataset[i]

        before_time = time()
        # certify the prediction of g around x
        x = x.cuda()
        prediction, radius = smoothed_classifier.certify(x, args.N0, args.N, args.alpha, args.batch)
        after_time = time()
        correct = int(prediction == label)
        time_elapsed = str(datetime.timedelta(seconds=(after_time - before_time)))
        print("{}\t{}\t{}\t{:.3}\t{}\t{}".format(
            i, label, prediction, radius, correct, time_elapsed), file=f, flush=True)

    f.close()
