import argparse

import numpy as np

import imageio

import torch

from tqdm import tqdm

import scipy
import scipy.io
import scipy.misc

from lib.model_test import D2Net
from lib.utils import preprocess_image
from lib.pyramid import process_multiscale

# CUDA
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")

# Argument parsing
parser = argparse.ArgumentParser(description="Feature extraction script")

parser.add_argument(
    "--image_list_file", type=str, required=True, help="path to a file containing a list of images to process"
)

parser.add_argument("--preprocessing", type=str, default="caffe", help="image preprocessing (caffe or torch)")
parser.add_argument("--model_file", type=str, default="models/d2_tf.pth", help="path to the full model")

parser.add_argument("--max_edge", type=int, default=1600, help="maximum image size at network input")
parser.add_argument("--max_sum_edges", type=int, default=2800, help="maximum sum of image sizes at network input")

parser.add_argument("--output_extension", type=str, default=".d2-net", help="extension for the output")
parser.add_argument("--output_type", type=str, default="npz", help="output file type (npz or mat)")

parser.add_argument("--multiscale", dest="multiscale", action="store_true", help="extract multiscale features")
parser.set_defaults(multiscale=False)

parser.add_argument(
    "--no-relu", dest="use_relu", action="store_false", help="remove ReLU after the dense feature extraction module"
)
parser.set_defaults(use_relu=True)

args = parser.parse_args()

print(args)

# Creating CNN model
model = D2Net(model_file=args.model_file, use_relu=args.use_relu, use_cuda=use_cuda)

# Process the file
with open(args.image_list_file, "r") as f:
    lines = f.readlines()
for line in tqdm(lines, total=len(lines)):
    path = line.strip()

    image = imageio.imread(path)
    if len(image.shape) == 2:
        image = image[:, :, np.newaxis]
        image = np.repeat(image, 3, -1)

    # TODO: switch to PIL.Image due to deprecation of scipy.misc.imresize.
    resized_image = image
    if max(resized_image.shape) > args.max_edge:
        resized_image = scipy.misc.imresize(resized_image, args.max_edge / max(resized_image.shape)).astype("float")
    if sum(resized_image.shape[:2]) > args.max_sum_edges:
        resized_image = scipy.misc.imresize(resized_image, args.max_sum_edges / sum(resized_image.shape[:2])).astype(
            "float"
        )

    fact_i = image.shape[0] / resized_image.shape[0]
    fact_j = image.shape[1] / resized_image.shape[1]

    input_image = preprocess_image(resized_image, preprocessing=args.preprocessing)
    with torch.no_grad():
        if args.multiscale:
            keypoints, scores, descriptors = process_multiscale(
                torch.tensor(input_image[np.newaxis, :, :, :].astype(np.float32), device=device), model
            )
        else:
            keypoints, scores, descriptors = process_multiscale(
                torch.tensor(input_image[np.newaxis, :, :, :].astype(np.float32), device=device), model, scales=[1]
            )

    # Input image coordinates
    keypoints[:, 0] *= fact_i
    keypoints[:, 1] *= fact_j
    # i, j -> u, v
    keypoints = keypoints[:, [1, 0, 2]]

    if args.output_type == "npz":
        with open(path + args.output_extension, "wb") as output_file:
            np.savez(output_file, keypoints=keypoints, scores=scores, descriptors=descriptors)
    elif args.output_type == "mat":
        with open(path + args.output_extension, "wb") as output_file:
            scipy.io.savemat(output_file, {"keypoints": keypoints, "scores": scores, "descriptors": descriptors})
    else:
        raise ValueError("Unknown output type.")
