import argparse
import logging
from pathlib import Path

import torch
from omegaconf import OmegaConf

from ..datasets import get_dataset
from ..models import get_model
from ..settings import DATA_PATH
from ..utils.export_predictions import export_predictions

resize = 1600

sp_keys = ["keypoints", "descriptors", "keypoint_scores"]

# SuperPoint
n_kpts = 2048
configs = {
    "sp": {
        "name": f"r{resize}_SP-k{n_kpts}-nms3",
        "keys": ["keypoints", "descriptors", "keypoint_scores"],
        "gray": True,
        "conf": {
            "name": "gluefactory_nonfree.superpoint",
            "nms_radius": 3,
            "max_num_keypoints": n_kpts,
            "detection_threshold": 0.000,
        },
    },
    "sift": {
        "name": f"r{resize}_SIFT-k{n_kpts}",
        "keys": ["keypoints", "descriptors", "keypoint_scores", "oris", "scales"],
        "gray": True,
        "conf": {
            "name": "sift",
            "max_num_keypoints": n_kpts,
            "options": {
                "peak_threshold": 0.001,
            },
            "peak_threshold": 0.001,
            "device": "cpu",
        },
    },
    "disk": {
        "name": f"r{resize}_DISK-k{n_kpts}-nms6",
        "keys": ["keypoints", "descriptors", "keypoint_scores"],
        "gray": False,
        "conf": {
            "name": "disk",
            "max_num_keypoints": n_kpts,
        },
    },
}


def run_export(feature_file, images, args):
    conf = {
        "data": {
            "name": "image_folder",
            "grayscale": configs[args.method]["gray"],
            "preprocessing": {
                "resize": resize,
            },
            "images": str(images),
            "batch_size": 1,
            "num_workers": args.num_workers,
        },
        "split": "train",
        "model": configs[args.method]["conf"],
    }

    conf = OmegaConf.create(conf)

    keys = configs[args.method]["keys"]
    dataset = get_dataset(conf.data.name)(conf.data)
    loader = dataset.get_data_loader(conf.split or "test")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = get_model(conf.model.name)(conf.model).eval().to(device)

    export_predictions(loader, model, feature_file, as_half=True, keys=keys)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("dataset", type=str)
    parser.add_argument("--export_prefix", type=str, default="")
    parser.add_argument("--method", type=str, default="sp")
    parser.add_argument("--scenes", type=str, default=None)
    parser.add_argument("--num_workers", type=int, default=0)
    args = parser.parse_args()

    export_name = configs[args.method]["name"]

    if args.dataset == "megadepth":
        data_root = Path(DATA_PATH, "megadepth/Undistorted_SfM")
        export_root = Path(DATA_PATH, "exports", "megadepth-undist-" + export_name)
        export_root.mkdir(parents=True, exist_ok=True)

        if args.scenes is None:
            scenes = [p.name for p in data_root.iterdir() if p.is_dir()]
        else:
            with open(DATA_PATH / "megadepth" / args.scenes, "r") as f:
                scenes = f.read().split()
        for i, scene in enumerate(scenes):
            # print(f'{i} / {len(scenes)}', scene)
            print(scene)
            feature_file = export_root / (scene + ".h5")
            if feature_file.exists():
                continue
            if not (data_root / scene / "images").exists():
                logging.info("Skip " + scene)
                continue
            logging.info(f"Export local features for scene {scene}")
            run_export(feature_file, data_root / scene / "images", args)
    else:
        data_root = Path(DATA_PATH, args.dataset)
        feature_file = Path(
            DATA_PATH, "exports", args.export_prefix + export_name + ".h5"
        )
        feature_file.parent.mkdir(exist_ok=True, parents=True)
        logging.info(
            f"Export local features for dataset {args.dataset} "
            f"to file {feature_file}"
        )
        run_export(feature_file, data_root)