import numpy as np
import cv2
import torch
from PIL import Image
import os

from .face_recognition import Siamese, FaceClassifier, trnscm, classes

current_path = os.path.dirname(os.path.abspath(__file__))

_SIAMESE = None
_CLASSIFIER = None
_DEVICE = None

def _to_bgr_np(img):
    if img is None:
        return None
    if isinstance(img, str):
        return cv2.imread(img)
    if isinstance(img, Image.Image):
        rgb = np.array(img.convert("RGB"))
        return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    if isinstance(img, np.ndarray):
        arr = img
        if arr.ndim == 2:
            return cv2.cvtColor(arr, cv2.COLOR_GRAY2BGR)
        if arr.ndim == 3 and arr.shape[2] == 4:
            arr = arr[:, :, :3]
        if arr.ndim == 3 and arr.shape[2] == 3:
            return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
        return arr
    return None

def _load_models():
    global _SIAMESE, _CLASSIFIER, _DEVICE
    if _SIAMESE is not None and _CLASSIFIER is not None:
        return _SIAMESE, _CLASSIFIER, _DEVICE

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

    siamese_path = os.path.join(current_path, "siamese_model.t7")
    clf_path = os.path.join(current_path, "face_classifier_model.pth")

    if not os.path.exists(siamese_path):
        raise FileNotFoundError(f"Missing siamese model: {siamese_path}")
    if not os.path.exists(clf_path):
        raise FileNotFoundError(f"Missing classifier model: {clf_path}")

    siamese = Siamese().to(_DEVICE)
    ckpt = torch.load(siamese_path, map_location=_DEVICE)
    if isinstance(ckpt, dict) and "net_dict" in ckpt:
        siamese.load_state_dict(ckpt["net_dict"], strict=True)
    else:
        siamese.load_state_dict(ckpt, strict=True)
    siamese.eval()

    classifier = FaceClassifier(input_dim=5, num_classes=len(classes)).to(_DEVICE)
    classifier.load_state_dict(torch.load(clf_path, map_location=_DEVICE), strict=True)
    classifier.eval()

    _SIAMESE, _CLASSIFIER = siamese, classifier
    return _SIAMESE, _CLASSIFIER, _DEVICE

def detected_face(image_bgr):
    face_haar = os.path.join(current_path, "haarcascade_frontalface_default.xml")
    face_cascade = cv2.CascadeClassifier(face_haar)
    if face_cascade.empty():
        return 0

    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.3, 5)
    if faces is None or len(faces) == 0:
        return 0

    x, y, w, h = max(faces, key=lambda f: f[2]*f[3])
    crop = gray[y:y+h, x:x+w]
    return Image.fromarray(crop)

def get_similarity(img1, img2):
    siamese, classifier, device = _load_models()

    i1 = _to_bgr_np(img1)
    i2 = _to_bgr_np(img2)
    if i1 is None or i2 is None:
        return 0.0

    f1 = detected_face(i1)
    f2 = detected_face(i2)

    if f1 == 0:
        f1 = Image.fromarray(cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY))
    if f2 == 0:
        f2 = Image.fromarray(cv2.cvtColor(i2, cv2.COLOR_BGR2GRAY))

    t1 = trnscm(f1).unsqueeze(0).to(device)
    t2 = trnscm(f2).unsqueeze(0).to(device)

    with torch.no_grad():
        e1 = siamese.forward_once(t1)
        e2 = siamese.forward_once(t2)
        dist = torch.nn.functional.pairwise_distance(e1, e2).item()

    return float(dist)

def get_face_class(img1):
    siamese, classifier, device = _load_models()

    i1 = _to_bgr_np(img1)
    if i1 is None:
        return "UNKNOWN"

    f1 = detected_face(i1)
    if f1 == 0:
        f1 = Image.fromarray(cv2.cvtColor(i1, cv2.COLOR_BGR2GRAY))

    t1 = trnscm(f1).unsqueeze(0).to(device)

    with torch.no_grad():
        emb = siamese.forward_once(t1)
        logits = classifier(emb)
        pred = int(torch.argmax(logits, dim=1).item())

    if pred < 0 or pred >= len(classes):
        return "UNKNOWN"
    return str(classes[pred])