# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import json
import logging
from typing import List, Optional
import torch
from torch import nn

from detectron2.utils.file_io import PathManager

from densepose.structures.mesh import create_mesh


class MeshAlignmentEvaluator:
    """
    Class for evaluation of 3D mesh alignment based on the learned vertex embeddings
    """

    def __init__(self, embedder: nn.Module, mesh_names: Optional[List[str]]):
        self.embedder = embedder
        # use the provided mesh names if not None and not an empty list
        self.mesh_names = mesh_names if mesh_names else embedder.mesh_names
        self.logger = logging.getLogger(__name__)
        with PathManager.open(
            "https://dl.fbaipublicfiles.com/densepose/data/cse/mesh_keyvertices_v0.json", "r"
        ) as f:
            self.mesh_keyvertices = json.load(f)

    def evaluate(self):
        ge_per_mesh = {}
        gps_per_mesh = {}
        for mesh_name_1 in self.mesh_names:
            avg_errors = []
            avg_gps = []
            embeddings_1 = self.embedder(mesh_name_1)
            keyvertices_1 = self.mesh_keyvertices[mesh_name_1]
            keyvertex_names_1 = list(keyvertices_1.keys())
            keyvertex_indices_1 = [keyvertices_1[name] for name in keyvertex_names_1]
            for mesh_name_2 in self.mesh_names:
                if mesh_name_1 == mesh_name_2:
                    continue
                embeddings_2 = self.embedder(mesh_name_2)
                keyvertices_2 = self.mesh_keyvertices[mesh_name_2]
                sim_matrix_12 = embeddings_1[keyvertex_indices_1].mm(embeddings_2.T)
                vertices_2_matching_keyvertices_1 = sim_matrix_12.argmax(axis=1)
                mesh_2 = create_mesh(mesh_name_2, embeddings_2.device)
                geodists = mesh_2.geodists[
                    vertices_2_matching_keyvertices_1,
                    [keyvertices_2[name] for name in keyvertex_names_1],
                ]
                Current_Mean_Distances = 0.255
                gps = (-(geodists**2) / (2 * (Current_Mean_Distances**2))).exp()
                avg_errors.append(geodists.mean().item())
                avg_gps.append(gps.mean().item())

            ge_mean = torch.as_tensor(avg_errors).mean().item()
            gps_mean = torch.as_tensor(avg_gps).mean().item()
            ge_per_mesh[mesh_name_1] = ge_mean
            gps_per_mesh[mesh_name_1] = gps_mean
        ge_mean_global = torch.as_tensor(list(ge_per_mesh.values())).mean().item()
        gps_mean_global = torch.as_tensor(list(gps_per_mesh.values())).mean().item()
        per_mesh_metrics = {
            "GE": ge_per_mesh,
            "GPS": gps_per_mesh,
        }
        return ge_mean_global, gps_mean_global, per_mesh_metrics