# Copyright (c) Megvii, Inc. and its affiliates.

"""do the evaluation work with single gpu
"""
import argparse
import os

import megengine as mge
import megengine.data as data
import megengine.data.transform as T
import megengine.functional as F
import numpy as np
from tqdm.auto import tqdm

from recognition.datasets import get_eval_dataset
from recognition.models import FaceRecognitionModel
from recognition.tools.utils import load_config_from_path

logger = mge.get_logger(__name__)


def get_inference_func(configs):
    """load checkpoint and construct inference function

    Args:
        configs (dict): configuration, required fields include:
            base_dir: base directory of experiment outputs
            evaluate_epoch: model of evaluate_epoch to evaluate

    Raises:
        FileNotFoundError: model of given epoch is not found

    Returns:
        inference_func (function): inference function mapping image to embedding
    """
    model = FaceRecognitionModel(configs)
    evaluate_epoch = configs["evaluate_epoch"]
    checkpoint_path = os.path.join(configs["base_dir"], f"epoch-{evaluate_epoch}-checkpoint.pkl")
    if os.path.exists(checkpoint_path):
        checkpoint_data = mge.load(checkpoint_path)
        model.load_state_dict(checkpoint_data["state_dict"], strict=False)
    else:
        raise FileNotFoundError(f"{checkpoint_path} not found!!!")

    def inference_func(images):
        model.eval()
        # classic test-time mirror augment
        embedding_origin = model.forward_embedding_only(images)
        embedding_mirror = model.forward_embedding_only(images[:, :, :, ::-1])
        embedding = embedding_origin + embedding_mirror
        embedding = F.normalize(embedding, axis=1)
        return embedding

    return inference_func


def extract_feature_and_clean_noise(configs, inference_func):
    """extract feature and clean noise. the noise cleaning algorithm is proposed in
    `"ArcFace: Additive Angular Margin Loss for Deep Face Recognition" <https://arxiv.org/pdf/1801.07698.pdf>`_
    please refer to https://github.com/deepinsight/insightface/blob/master/Evaluation/Megaface/remove_noises.py for
    more detail. this implement does basicly the same thing as the above, but with much higher speed

    Args:
        configs (dict): configuration, required fields include:
            batch_size: inference batch size
            feature_dim: model output feature dimension
            base_dir: base directory of experiment outputs
            dataset_dir: directory of dataset root
        inference_func (function): constructed inference function

    Returns:
        facescrub_feature (np.array): noise-cleaned feature of facescrub (shape: n * (feature_dim + 1))
        facescrub_label (np.array): label of facescrub (shape: n)
        megaface_feature (np.array): noise-cleaned feature of megaface (shape: m * (feature_dim + 1))
    """

    def prepare_dataset(name):
        """prepare dataset

        Args:
            name (str): name of the dataset, should be one of {facescrub, megaface}

        Returns:
            dataset (data.Dataset): required dataset
            queue (data.DataLoader): corresponding dataloader
        """
        preprocess = T.Compose([T.Normalize(mean=127.5, std=128), 
T.ToMode("CHW")