import chainer
from chainer import cuda
import chainer.functions as F
import chainer.links as L
from chainer import optimizers
import numpy as np
import os
import cv2
import logging
import datetime

log_name='predict_log.txt'
logging.basicConfig(filename=log_name,level=logging.INFO)
logging.info('***************')
logging.info(datetime.datetime.today())
logging.info('Predict Log')

# 中間層の数
n_units   = 1000

# 使用するモデル・オプティマイザのepoch数
epoch = 17

x_train = []
x_test = []
y_train = []
y_test = []

# 画像データの読み込み
path = './data/images/test'
IMGSIZE = 28
images = os.listdir(path)
count=0
for image in images:
    src = cv2.imread(path+'/'+image,1)
    src = cv2.resize(src,(IMGSIZE*3,IMGSIZE))
    src = cv2.cvtColor(src,cv2.COLOR_BGR2GRAY)
    src = cv2.bitwise_not(src)
    x_test.append(src)
    file = image.split('.')[0]
    y_test.append(file.split('_')[1])
    count = count + 1

N = len(x_test)
N_test = len(y_test)
print('N is {}'.format(N))

model = chainer.FunctionSet(l1=F.Linear(2352, n_units),
                    l2=F.Linear(n_units, n_units),
                    l3=F.Linear(n_units, 19))
chainer.serializers.load_npz('data/test_100/model_tegaki_{}'.format(epoch),model)
optimizer = chainer.optimizers.Adam()
optimizer.setup(model)

chainer.serializers.load_npz('data/test_100/optimizer_tegaki_{}'.format(epoch),optimizer)

# 認識してみる
print(' \>Predict \>Label \\')
correct = 0
for idx in np.random.permutation(N)[:1000]:
    xxx = x_test[idx].astype(np.float32)
    x = xxx.reshape(1,1,IMGSIZE*3,IMGSIZE)/255

    h1 = F.relu(model.l1(chainer.Variable(x)))
    h2 = F.relu(model.l2(h1))
    y = model.l3(h2)

    if(int(y_test[idx]) == np.argmax(y.data)):
        correct = correct + 1
    else:
        cv2.imshow('img',x_test[idx])
        src = cv2.bitwise_not(x_test[idx])
        cv2.waitKey(0)
        cv2.imwrite('wrong_{}.png'.format(np.argmax(y.data)),src)
        exit()
        print('Missed. \>',np.argmax(y.data),' \>',y_test[idx],'\\')
        logging.debug('missed:{}!={}'.format(np.argmax(y.data),y_test[idx]))

accuracy = correct / 1000
print('accuracy:',accuracy,'({} / {})'.format(correct,1000))
logging.info('accuracy: {}'.format(accuracy))

logging.info('***************')