import os
import cv2
import argparse
import numpy as np
import torch
import torchvision

from torchvision import datasets, transforms
from torch.autograd import Variable
from network_v0.model import PointModel
from datasets.hp_loader import PatchesDataset
from torch.utils.data import DataLoader
from evaluation.evaluate import evaluate_keypoint_net


def main():
    parser = argparse.ArgumentParser(description="Testing")
    parser.add_argument("--device", default=0, type=int, help="which gpu to run on.")
    parser.add_argument("--test_dir", required=True, type=str, help="Test data path.")
    opt = parser.parse_args()

    torch.manual_seed(0)
    use_gpu = torch.cuda.is_available()
    if use_gpu:
        torch.cuda.set_device(opt.device)

    # Load data in 320x240
    hp_dataset_320x240 = PatchesDataset(
        root_dir=opt.test_dir, use_color=True, output_shape=(320, 240), type="all"
    )
    data_loader_320x240 = DataLoader(
        hp_dataset_320x240,
        batch_size=1,
        pin_memory=False,
        shuffle=False,
        num_workers=4,
        worker_init_fn=None,
        sampler=None,
    )

    # Load data in 640x480
    hp_dataset_640x480 = PatchesDataset(
        root_dir=opt.test_dir, use_color=True, output_shape=(640, 480), type="all"
    )
    data_loader_640x480 = DataLoader(
        hp_dataset_640x480,
        batch_size=1,
        pin_memory=False,
        shuffle=False,
        num_workers=4,
        worker_init_fn=None,
        sampler=None,
    )

    # Load model
    model = PointModel(is_test=True)
    ckpt = torch.load("./checkpoints/PointModel_v0.pth")
    model.load_state_dict(ckpt["model_state"])
    model = model.eval()
    if use_gpu:
        model = model.cuda()

    print("Evaluating in 320x240, 300 points")
    rep, loc, c1, c3, c5, mscore = evaluate_keypoint_net(
        data_loader_320x240, model, output_shape=(320, 240), top_k=300
    )

    print("Repeatability: {0:.3f}".format(rep))
    print("Localization Error: {0:.3f}".format(loc))
    print("H-1 Accuracy: {:.3f}".format(c1))
    print("H-3 Accuracy: {:.3f}".format(c3))
    print("H-5 Accuracy: {:.3f}".format(c5))
    print("Matching Score: {:.3f}".format(mscore))
    print("\n")

    print("Evaluating in 640x480, 1000 points")
    rep, loc, c1, c3, c5, mscore = evaluate_keypoint_net(
        data_loader_640x480, model, output_shape=(640, 480), top_k=1000
    )

    print("Repeatability: {0:.3f}".format(rep))
    print("Localization Error: {0:.3f}".format(loc))
    print("H-1 Accuracy: {:.3f}".format(c1))
    print("H-3 Accuracy: {:.3f}".format(c3))
    print("H-5 Accuracy: {:.3f}".format(c5))
    print("Matching Score: {:.3f}".format(mscore))
    print("\n")


if __name__ == "__main__":
    main()