import os
import pandas as pd
import torch
import torchvision
import torch.utils.data
from torchvision import transforms
from collections import Counter
from torch.nn import DataParallel
from datetime import datetime
from PIL import Image
from torch.optim.lr_scheduler import MultiStepLR
from nts_net.config import BATCH_SIZE, PROPOSAL_NUM, SAVE_FREQ, LR, WD, resume, save_dir
from nts_net import model
from nts_net.utils import init_log, progress_bar

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
start_epoch = 1

test_path = f'data/TestFolder'
testset = torchvision.datasets.ImageFolder(root=test_path, transform=transforms.ToTensor())
label_dist = testset.class_to_idx

# read dataset
transform_test = transforms.Compose([
    # transforms.Resize((600, 600), Image.BILINEAR),
    # transforms.CenterCrop((448, 448)),
    transforms.RandomCrop((224, 224), Image.BILINEAR),
    transforms.ToTensor(),
    # transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])

origin_path = f'data/rock'
rocks = os.listdir(origin_path)
rocks = list(filter(lambda r : r[-5] == '1', rocks))

rock_label = pd.read_csv('data/rock_label_1.csv')
rock_label = rock_label.iloc[:,:].values

data_dist = {}
for rl in rock_label:
    data_dist[str(rl[0])] = rl[1]

# define model
net = model.attention_net(topN=PROPOSAL_NUM)
if resume:
    ckpt = torch.load(resume)
    net.load_state_dict(ckpt['net_state_dict'])
    start_epoch = ckpt['epoch'] + 1

# define optimizers
net = net.cuda()
net = DataParallel(net)

net.eval()
# evaluate on test set
test_correct = 0
total = 0
for i, rock in enumerate(rocks):
    class_name = data_dist[int(rock[:-6])]
    label = label_dist[class_name]

    img = Image.open(os.path.join(origin_path, rock))
    data = []
    for j in range(BATCH_SIZE):
        data.append(transform_test(img))

    data = torch.stack(data, dim=0)

    with torch.no_grad():
        data = data.cuda()
        batch_size = data.size(0)
        _, concat_logits, _, _, _ = net(data)
        # calculate accuracy
        _, concat_predict = torch.max(concat_logits, 1)

        predict = concat_predict.reshape(-1)
        predict_label = Counter(predict.numpy()).most_common(1)[0][0]
        if predict_label == label: test_correct += 1
        total += 1
        progress_bar(i, len(rocks), 'eval test set')

test_acc = float(test_correct) / total
print(
    'test acc: {:.3f} total sample: {}'.format(
        test_acc,
        total))
print('finishing testing')