###################################################################################################
# Measure the detection performance: reference code is https://github.com/ShiyuLiang/odin-pytorch #
###################################################################################################

import argparse
from pathlib import Path

import numpy as np
import tensorflow.keras.backend as K

from sdenet.data import data_loader
from sdenet.eval import calculate_log
from sdenet.models import sdenet, resnet
from sdenet.models.helpers import load_weights


def main():
    parser = argparse.ArgumentParser(description='Test code - measure the detection peformance')
    parser.add_argument('--dataset', required=True, help='in domain dataset')
    parser.add_argument('--eva_iter', default=10, type=int, help='number of passes when evaluation')
    parser.add_argument('--network', type=str, choices=['resnet', 'sdenet', 'mc_dropout'], default='resnet')
    parser.add_argument('--batch-size', type=int, default=256, help='batch size')
    parser.add_argument('--imageSize', type=int, default=28, help='the height / width of the input image to network')
    parser.add_argument('--out_dataset', required=True, help='out-of-dist dataset: cifar10 | svhn | imagenet | lsun')
    parser.add_argument('--num_classes', type=int, default=10, help='number of classes (default: 10)')
    parser.add_argument('--pre_trained_net', required=True, help="path to pre trained_net h5")
    parser.add_argument('--test_batch_size', type=int, default=1000)
    parser.add_argument('--net_sigma', default=20, type=int, help='sigma coefficient for the diffusion')
    args = parser.parse_args()
    print(args)

    outf = Path('test') / args.network
    outf.mkdir(parents=True, exist_ok=True)

    print('Load model')
    if args.network == 'resnet':
        model = resnet.ResidualNet()
        args.eva_iter = 1
    elif args.network == 'sdenet':
        model = sdenet.SDENet(layer_depth=6, num_classes=10, dim=64, task=args.dataset)
        model.set_sigma(args.net_sigma)
    else:
        raise Exception('Model not found.')
    # elif args.network == 'mc_dropout':
    #     model = models.Resnet_dropout()

    if Path(args.pre_trained_net).exists():
        load_weights(model, args.pre_trained_net)
    else:
        raise Exception('Could not find the weights file.')

    apply_grayscale = args.dataset == 'mnist'
    print('load target data: ', args.dataset)
    _, test_loader = data_loader.getDataSet(args.dataset, args.batch_size, args.test_batch_size, args.imageSize,
                                            apply_grayscale=apply_grayscale)

    print('load non target data: ', args.out_dataset)
    nt_train_loader, nt_test_loader = data_loader.getDataSet(args.out_dataset, args.batch_size, args.test_batch_size,
                                                             args.imageSize, apply_grayscale=apply_grayscale)

    def generate_target():
        correct = 0
        total = 0
        f1 = open('%s/confidence_Base_In.txt' % outf, 'w')
        f3 = open('%s/confidence_Base_Succ.txt' % outf, 'w')
        f4 = open('%s/confidence_Base_Err.txt' % outf, 'w')

        for data, targets in test_loader:
            data, targets = np.array(data), np.array(targets)
            total += len(data)
            data = np.transpose(data, (0, 2, 3, 1))
            batch_output = 0
            for j in range(args.eva_iter):
                batch_output = batch_output + K.softmax(model(data), axis=-1)
            batch_output = batch_output / args.eva_iter
            # compute the accuracy
            predicted = K.argmax(batch_output, axis=1).numpy()
            correct += (predicted == targets).sum()
            correct_index = (predicted == targets)
            for i in range(len(data)):
                # confidence score: max_y p(y|x)
                # output = batch_output[i].view(1, -1)
                output = K.expand_dims(batch_output[i], axis=0)
                soft_out = K.max(output).numpy()
                f1.write("{}\n".format(soft_out))
                if correct_index[i] == 1:
                    f3.write("{}\n".format(soft_out))
                elif correct_index[i] == 0:
                    f4.write("{}\n".format(soft_out))
        f1.close()
        f3.close()
        f4.close()
        print('\n Final Accuracy: {}/{} ({:.2f}%)\n '.format(correct, total, 100. * correct / total))

    def generate_non_target():
        total = 0
        f2 = open('%s/confidence_Base_Out.txt' % outf, 'w')
        for data, targets in nt_test_loader:
            data, targets = np.array(data), np.array(targets)
            data = np.transpose(data, (0, 2, 3, 1))
            total += len(data)
            batch_output = 0
            for j in range(args.eva_iter):
                batch_output = batch_output + K.softmax(model(data), axis=-1)
            batch_output = batch_output / args.eva_iter
            for i in range(len(data)):
                # confidence score: max_y p(y|x)
                output = K.expand_dims(batch_output[i], axis=0)
                soft_out = K.max(output).numpy()
                f2.write("{}\n".format(soft_out))
        f2.close()

    print('generate log from in-distribution data')
    generate_target()
    print('generate log  from out-of-distribution data')
    generate_non_target()
    print('calculate metrics for OOD')
    calculate_log.metric(outf, 'OOD')
    print('calculate metrics for mis')
    calculate_log.metric(outf, 'mis')


if __name__ == '__main__':
    main()
