import torch
import numpy as np
import cv2
from PIL import Image
from torchvision.transforms import functional as F
from torchvision.models.detection import (
    fasterrcnn_resnet50_fpn_v2,
    keypointrcnn_resnet50_fpn,
)
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor


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


# -----------------------------
# BUILD MODELS
# -----------------------------
def build_face_detector(num_classes=2):
    model = fasterrcnn_resnet50_fpn_v2(weights=None)
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model


def build_landmark_model(num_classes=2, num_keypoints=13):
    model = keypointrcnn_resnet50_fpn(
        weights=None,
        weights_backbone=None,
        num_classes=num_classes,
        num_keypoints=num_keypoints,
    )
    return model


# -----------------------------
# LOAD WEIGHTS
# -----------------------------
def load_models(face_path, landmark_path, num_keypoints):
    face_model = build_face_detector()
    landmark_model = build_landmark_model(num_keypoints=num_keypoints)

    face_model.load_state_dict(torch.load(face_path, map_location=DEVICE))
    landmark_model.load_state_dict(torch.load(landmark_path, map_location=DEVICE))

    face_model.to(DEVICE).eval()
    landmark_model.to(DEVICE).eval()

    return face_model, landmark_model


# -----------------------------
# CASCADE INFERENCE
# -----------------------------
def run_inference(image, face_model, landmark_model,
                  face_score_thr=0.5,
                  kpt_score_thr=0.2):

    image_pil = Image.fromarray(image).convert("RGB")
    img_rgb = np.array(image_pil)
    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)

    with torch.no_grad():
        det_out = face_model([F.to_tensor(image_pil).to(DEVICE)])[0]

    boxes = det_out["boxes"].cpu().numpy()
    scores = det_out["scores"].cpu().numpy()

    keep = np.where(scores >= face_score_thr)[0]

    for i in keep:
        x1, y1, x2, y2 = boxes[i].astype(int)

        cv2.rectangle(img_bgr, (x1, y1), (x2, y2), (255, 180, 0), 2)

        crop = img_rgb[y1:y2, x1:x2]

        with torch.no_grad():
            kp_out = landmark_model([F.to_tensor(crop).to(DEVICE)])[0]

        if len(kp_out["scores"]) == 0:
            continue

        best = torch.argmax(kp_out["scores"]).item()
        if kp_out["scores"][best] < kpt_score_thr:
            continue

        keypoints = kp_out["keypoints"][best].cpu().numpy()

        for kx, ky, kv in keypoints:
            if kv > 0:
                cv2.circle(img_bgr,
                           (int(kx + x1), int(ky + y1)),
                           2,
                           (0, 255, 0),
                           -1)

    return cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)