#!/usr/bin/env python

from __future__ import annotations

import os
import pathlib
import sys

import cv2
import gradio as gr
import huggingface_hub
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download

torch.jit.script = lambda f: f

sys.path.insert(0, "face_detection")
sys.path.insert(0, "face_parsing")
sys.path.insert(0, "fpage")
sys.path.insert(0, "roi_tanh_warping")

from ibug.age_estimation import AgeEstimator
from ibug.face_detection import RetinaFacePredictor
from ibug.face_parsing.utils import label_colormap

DESCRIPTION = "# [FP-Age](https://github.com/ibug-group/fpage)"


def is_lfs_pointer_file(path: pathlib.Path) -> bool:
    try:
        with open(path, "r") as f:
            # Git LFS pointer files usually start with version line
            version_line = f.readline()
            if version_line.startswith("version https://git-lfs.github.com/spec/"):
                # Check for the presence of oid and size lines
                oid_line = f.readline()
                size_line = f.readline()
                if oid_line.startswith("oid sha256:") and size_line.startswith("size "):
                    return True
    except Exception as e:
        print(f"Error reading file {path}: {e}")
    return False


lfs_model_paths = sorted(pathlib.Path("face_parsing").rglob("*.torch"))
for lfs_model_path in lfs_model_paths:
    if is_lfs_pointer_file(lfs_model_path):
        os.remove(lfs_model_path)
        out_path = hf_hub_download(
            "public-data/ibug-face-parsing",
            filename=lfs_model_path.name,
            repo_type="model",
            subfolder=lfs_model_path.parts[-3],
        )
        os.symlink(out_path, lfs_model_path)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

detector = RetinaFacePredictor(threshold=0.8, device="cpu", model=RetinaFacePredictor.get_model("mobilenet0.25"))
detector.device = device
detector.net.to(device)

model = AgeEstimator(
    device=device,
    ckpt=huggingface_hub.hf_hub_download("hysts/ibug", "fpage/models/fpage-resnet50-fcn-14-97.torch"),
    encoder="resnet50",
    decoder="fcn",
    age_classes=97,
    face_classes=14,
)


@spaces.GPU
def predict(image: np.ndarray, max_num_faces: int) -> np.ndarray:
    colormap = label_colormap(14)

    # RGB -> BGR
    image = image[:, :, ::-1]

    faces = detector(image, rgb=False)
    if len(faces) == 0:
        raise RuntimeError("No face was found.")
    faces = sorted(list(faces), key=lambda x: -x[4])[:max_num_faces][::-1]
    ages, masks = model.predict_img(image, faces, rgb=False)

    mask_image = np.zeros_like(image)
    for mask in masks:
        temp = colormap[mask]
        mask_image[temp > 0] = temp[temp > 0]

    res = image.astype(float) * 0.5 + mask_image[:, :, ::-1] * 0.5
    res = np.clip(np.round(res), 0, 255).astype(np.uint8)

    for face, age in zip(faces, ages):
        bbox = np.round(face[:4]).astype(int)
        cv2.rectangle(
            res,
            (bbox[0], bbox[1]),
            (bbox[2], bbox[3]),
            color=(0, 255, 0),
            thickness=2,
        )

        text_content = f"Age: ({age: .1f})"
        cv2.putText(
            res,
            text_content,
            (bbox[0], bbox[1] - 10),
            cv2.FONT_HERSHEY_DUPLEX,
            1,
            (255, 255, 255),
            lineType=cv2.LINE_AA,
        )

    return res[:, :, ::-1]


with gr.Blocks(css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            image = gr.Image(label="Input")
            max_num_faces = gr.Slider(label="Max Number of Faces", minimum=1, maximum=20, step=1, value=5)
            run_button = gr.Button()
        with gr.Column():
            result = gr.Image(label="Output")
    gr.Examples(
        examples=[[path.as_posix(), 5] for path in sorted(pathlib.Path("images").rglob("*.jpg"))],
        inputs=[image, max_num_faces],
        outputs=result,
        fn=predict,
    )
    run_button.click(
        fn=predict,
        inputs=[image, max_num_faces],
        outputs=result,
        api_name="predict",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch()