# -*- coding: utf-8 -*-
# @Organization  : Alibaba XR-Lab
# @Author        : Lingteng Qiu
# @Email         : 220019047@link.cuhk.edu.cn
# @Time          : 2025-03-03 10:29:00
# @Function      : easy to use FaceSimilarity metric

import os
import pdb
import shutil
import sys

sys.path.append("./")
from collections import defaultdict

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from prettytable import PrettyTable
from torch.utils.data import Dataset
from torchmetrics.image import StructuralSimilarityIndexMeasure
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchvision import transforms
from tqdm import tqdm

from openlrm.models.arcface_utils import ResNetArcFace
from openlrm.utils.face_detector import FaceDetector

device = "cuda"
model_path = "./pretrained_models/gagatracker/vgghead/vgg_heads_l.trcd"
face_detector = FaceDetector(model_path=model_path, device=device)

id_face_net = ResNetArcFace()
id_face_net.cuda()
id_face_net.eval()


def get_image_paths_current_dir(folder_path):
    image_extensions = {
        ".jpg",
        ".jpeg",
        ".png",
        ".gif",
        ".bmp",
        ".tiff",
        ".webp",
        ".jfif",
    }

    return sorted(
        [
            os.path.join(folder_path, f)
            for f in os.listdir(folder_path)
            if os.path.splitext(f)[1].lower() in image_extensions
        ]
    )


def write_json(path, x):
    """write a json file.

    Args:
        path (str): path to write json file.
        x (dict): dict to write.
    """
    import json

    with open(path, "w") as f:
        json.dump(x, f, indent=2)


def crop_face_image(image_path):
    rgb = np.array(Image.open(image_path))
    rgb = torch.from_numpy(rgb).permute(2, 0, 1)
    bbox = face_detector(rgb)
    head_rgb = rgb[:, int(bbox[1]) : int(bbox[3]), int(bbox[0]) : int(bbox[2])]
    head_rgb = head_rgb.permute(1, 2, 0)
    head_rgb = head_rgb.cpu().numpy()
    return head_rgb


def gray_resize_for_identity(out, size=128):
    out_gray = (
        0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] + 0.1140 * out[:, 2, :, :]
    )
    out_gray = out_gray.unsqueeze(1)
    out_gray = F.interpolate(
        out_gray, (size, size), mode="bilinear", align_corners=False
    )
    return out_gray


@torch.no_grad()
def eval(input_folder, target_folder, front_view_idx, device="cuda"):
    src_img = os.path.join(target_folder, f"{front_view_idx:05d}.png")
    if not os.path.exists(src_img):
        return -1

    head_img = crop_face_image(src_img)
    input_imgs = get_image_paths_current_dir(input_folder)
    if "visualization" in input_imgs[-1]:
        input_imgs = input_imgs[:-1]

    to_tensor = transforms.ToTensor()

    head_img = to_tensor(head_img).unsqueeze(0).to(device)
    src_head_tensor = gray_resize_for_identity(head_img)
    src_head_feature = id_face_net(src_head_tensor).detach()

    face_id_loss_list = []

    for input_img in input_imgs:
        try:
            input_img = crop_face_image(input_img)
            input_head_tensor = gray_resize_for_identity(
                to_tensor(input_img).unsqueeze(0).to(device)
            )
            input_head_feature = id_face_net(input_head_tensor).detach()
            face_id_loss = F.l1_loss(input_head_feature, src_head_feature)
            face_id_loss_list.append(face_id_loss.item())

        except:
            continue
    if len(face_id_loss_list) > 0:

        return min(face_id_loss_list)  # return max similarity view.
    else:
        return -1


def get_parse():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-f1", "--folder1", type=str, required=True)
    parser.add_argument("-f2", "--folder2", type=str, required=True)
    parser.add_argument("--pad", action="store_true")
    parser.add_argument("--debug", action="store_true")
    args = parser.parse_args()
    return args


if __name__ == "__main__":

    opt = get_parse()

    input_folder = opt.folder1
    target_folder = opt.folder2

    valid_txt = os.path.join(input_folder, "front_view.txt")

    target_folder = target_folder[:-1] if target_folder[-1] == "/" else target_folder

    target_key = target_folder.split("/")[-2:]

    save_folder = os.path.join("./exps/metrics", "psnr_results", *target_key)
    os.makedirs(save_folder, exist_ok=True)

    with open(valid_txt) as f:
        items = f.read().splitlines()
        items = [x.split(" ") for x in items]

    results_dict = defaultdict(dict)
    face_similarity_list = []

    for item_ in tqdm(items):

        try:
            item, front_view_idx = item_
            front_view_idx = int(front_view_idx)
        except:
            print(item_)

        target_item_folder = os.path.join(input_folder, item)
        input_item_folder = os.path.join(target_folder, item, "rgb")

        if os.path.exists(input_item_folder) and os.path.exists(target_item_folder):

            fs_ = eval(input_item_folder, target_item_folder, front_view_idx)

            if fs_ == -1:
                continue

            face_similarity_list.append(fs_)

            results_dict[item]["face_similarity"] = fs_
            if opt.debug:
                break
            print(results_dict)

    results_dict["all_mean"]["face_similarity"] = np.mean(face_similarity_list)

    write_json(os.path.join(save_folder, "face_similarity.json"), results_dict)