#!/usr/bin/env python
# coding: utf-8

import os, glob, cv2
import argparse
from argparse import Namespace
import yaml
from tqdm import tqdm
import torch
from torch.utils.data import Dataset, DataLoader, SequentialSampler

from src.datasets.custom_dataloader import TestDataLoader
from src.utils.dataset import read_img_gray
from configs.data.base import cfg as data_cfg
import viz


def get_model_config(method_name, dataset_name, root_dir="viz"):
    config_file = f"{root_dir}/configs/{method_name}.yml"
    with open(config_file, "r") as f:
        model_conf = yaml.load(f, Loader=yaml.FullLoader)[dataset_name]
    return model_conf


class DemoDataset(Dataset):
    def __init__(self, dataset_dir, img_file=None, resize=0, down_factor=16):
        self.dataset_dir = dataset_dir
        if img_file is None:
            self.list_img_files = glob.glob(os.path.join(dataset_dir, "*.*"))
            self.list_img_files.sort()
        else:
            with open(img_file) as f:
                self.list_img_files = [
                    os.path.join(dataset_dir, img_file.strip())
                    for img_file in f.readlines()
                ]
        self.resize = resize
        self.down_factor = down_factor

    def __len__(self):
        return len(self.list_img_files)

    def __getitem__(self, idx):
        img_path = self.list_img_files[
            idx
        ]  # os.path.join(self.dataset_dir, self.list_img_files[idx])
        img, scale = read_img_gray(
            img_path, resize=self.resize, down_factor=self.down_factor
        )
        return {"img": img, "id": idx, "img_path": img_path}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Visualize matches")
    parser.add_argument("--gpu", "-gpu", type=str, default="0")
    parser.add_argument("--method", type=str, default=None)
    parser.add_argument("--dataset_dir", type=str, default="data/aachen-day-night")
    parser.add_argument("--pair_dir", type=str, default=None)
    parser.add_argument(
        "--dataset_name",
        type=str,
        choices=["megadepth", "scannet", "aachen_v1.1", "inloc"],
        default="megadepth",
    )
    parser.add_argument("--measure_time", action="store_true")
    parser.add_argument("--no_viz", action="store_true")
    parser.add_argument("--compute_eval_metrics", action="store_true")
    parser.add_argument("--run_demo", action="store_true")

    args = parser.parse_args()

    model_cfg = get_model_config(args.method, args.dataset_name)
    class_name = model_cfg["class"]
    model = viz.__dict__[class_name](model_cfg)
    # all_args = Namespace(**vars(args), **model_cfg)
    if not args.run_demo:
        if args.dataset_name == "megadepth":
            from configs.data.megadepth_test_1500 import cfg

            data_cfg.merge_from_other_cfg(cfg)
        elif args.dataset_name == "scannet":
            from configs.data.scannet_test_1500 import cfg

            data_cfg.merge_from_other_cfg(cfg)
        elif args.dataset_name == "aachen_v1.1":
            data_cfg.merge_from_list(
                [
                    "DATASET.TEST_DATA_SOURCE",
                    "aachen_v1.1",
                    "DATASET.TEST_DATA_ROOT",
                    os.path.join(args.dataset_dir, "images/images_upright"),
                    "DATASET.TEST_LIST_PATH",
                    args.pair_dir,
                    "DATASET.TEST_IMGSIZE",
                    model_cfg["imsize"],
                ]
            )
        elif args.dataset_name == "inloc":
            data_cfg.merge_from_list(
                [
                    "DATASET.TEST_DATA_SOURCE",
                    "inloc",
                    "DATASET.TEST_DATA_ROOT",
                    args.dataset_dir,
                    "DATASET.TEST_LIST_PATH",
                    args.pair_dir,
                    "DATASET.TEST_IMGSIZE",
                    model_cfg["imsize"],
                ]
            )

        has_ground_truth = str(data_cfg.DATASET.TEST_DATA_SOURCE).lower() in [
            "megadepth",
            "scannet",
        ]
        dataloader = TestDataLoader(data_cfg)
        with torch.no_grad():
            for data_dict in tqdm(dataloader):
                for k, v in data_dict.items():
                    if isinstance(v, torch.Tensor):
                        data_dict[k] = v.cuda() if torch.cuda.is_available() else v
                img_root_dir = data_cfg.DATASET.TEST_DATA_ROOT
                model.match_and_draw(
                    data_dict,
                    root_dir=img_root_dir,
                    ground_truth=has_ground_truth,
                    measure_time=args.measure_time,
                    viz_matches=(not args.no_viz),
                )

        if args.measure_time:
            print(
                "Running time for each image is {} miliseconds".format(
                    model.measure_time()
                )
            )
        if args.compute_eval_metrics and has_ground_truth:
            model.compute_eval_metrics()
    else:
        demo_dataset = DemoDataset(args.dataset_dir, img_file=args.pair_dir, resize=640)
        sampler = SequentialSampler(demo_dataset)
        dataloader = DataLoader(demo_dataset, batch_size=1, sampler=sampler)

        writer = cv2.VideoWriter(
            "topicfm_demo.mp4",
            cv2.VideoWriter_fourcc(*"mp4v"),
            15,
            (640 * 2 + 5, 480 * 2 + 10),
        )

        model.run_demo(
            iter(dataloader), writer
        )  # , output_dir="demo", no_display=True)