import sys
import caffe
import numpy as np
from pylab import *
import matplotlib.pyplot as plt
import tempfile
import os
import lmdb


caffe_root = 'F:/Programs/Caffe/caffe-windows/build/'  # this file should be run from {caffe_root}/examples (otherwise change this line)
sys.path.insert(0, caffe_root + 'python')

caffe.set_device(0)
caffe.set_mode_gpu()

weights = os.path.join(caffe_root, '../models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel')
assert os.path.exists(weights)


def Test_net_on_lmdb(database_name, deploy_net_filename, deploy_weights=weights):
    env = lmdb.open(database_name, readonly=True)
    num_images = env.stat()["entries"]

    testnet_net = caffe.Net(deploy_net_filename, weights=deploy_weights, phase=caffe.TEST)
    numClasses = 4

    # num_batches depends on the number of Images you will test and the batch_size defined on test.prototxt
    batch_size = 256
    num_batches = num_images/ batch_size

    # Initialize the accuracy reported by the net
    acc_top1 = 0
    #acc_top5 = 0

    for t in range(num_batches):
        # Run the net for the current image batch
        testnet_net.forward()

        # Update accuracy with the average accuracy of the current batch
        acc_top1 = acc_top1 + testnet_net.blobs["accuracy"].data
        #acc_top5 = acc_top5 + net.blobs[nameAccuracyTop5].data
    acc_top1 = acc_top1 / num_batches
    #acc_top5 = acc_top5 / num_batches
    return acc_top1


solver_filename = "TestSolver.proto"
test_solver = caffe.get_solver(solver_filename)
test_solver.restore("F://caffe_snapshots/imagenet_128_gen_fntng2_iter_10000.solverstate")
test_solver.net.copy_from("F://caffe_snapshots/imagenet_128_gen_fntng2_iter_10000.caffemodel")
#test_solver.net.copy_from(weights)
#testnet_net = caffe.Net("ImageNetRandom.proto", weights=weights, phase=caffe.TEST)
test_solver.solve()
#acc1 = Test_net_on_lmdb("F://rgbval.lmdb/", "ImageNetRandom.proto")
#print "RandomNet accuracy: " + str(acc1)
#acc2 = Test_net_on_lmdb("F://rgbval.lmdb/", "ImageNetTrain.proto", "F://caffe_snapshots/imagenet4_fntng_iter_70000.caffemodel")
#print "LearnedNet accuracy: " + str(acc2)

#with open("temp_output2", "w") as file:
#    file.write("RandomNet accuracy: " + str(acc1) + "\n")
#    file.write("LearnedNet accuracy: " + str(acc2) + "\n")
