"""
CLIP + face embeddings. Same models/config as your Embedding_expansion
notebook: CLIP ViT-H-14 (laion2b_s32b_b79k, 1024-d, L2-normalized) and
InsightFace buffalo_l (512-d). Lazily loaded once on first use instead
of at notebook-cell/import time, since a Gradio app shouldn't pay this
cost unless Part 2 is actually run.
"""
import numpy as np
import torch

CLIP_MODEL_NAME = "ViT-H-14"
CLIP_PRETRAINED = "laion2b_s32b_b79k"
CLIP_DIM = 1024
FACE_DIM = 512

_clip_model = None
_clip_preprocess = None
_device = None
_face_app = None


def _get_device():
    global _device
    if _device is None:
        _device = "cuda" if torch.cuda.is_available() else "cpu"
    return _device


def _get_clip():
    global _clip_model, _clip_preprocess
    if _clip_model is None:
        import open_clip
        device = _get_device()
        _clip_model, _, _clip_preprocess = open_clip.create_model_and_transforms(
            CLIP_MODEL_NAME, pretrained=CLIP_PRETRAINED
        )
        _clip_model = _clip_model.to(device).eval()
    return _clip_model, _clip_preprocess


def _get_face_app():
    global _face_app
    if _face_app is None:
        from insightface.app import FaceAnalysis
        _face_app = FaceAnalysis(
            name="buffalo_l",
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
        )
        _face_app.prepare(ctx_id=0, det_size=(640, 640))
    return _face_app


@torch.no_grad()
def embed_clip(pil_image):
    """Returns a (1024,) float32 L2-normalized CLIP embedding."""
    model, preprocess = _get_clip()
    x = preprocess(pil_image).unsqueeze(0).to(_get_device())
    f = model.encode_image(x)
    f = f / f.norm(dim=-1, keepdim=True)
    return f.squeeze(0).cpu().numpy().astype(np.float32)


def embed_face(pil_image):
    """
    Returns (emb, det_score, bbox_xywh, n_faces, status).
    status: "ok" | "ok_multi" (largest face used) | "no_face"
    """
    face_app = _get_face_app()
    rgb = np.array(pil_image.convert("RGB"))
    bgr = rgb[:, :, ::-1].copy()
    faces = face_app.get(bgr)
    if not faces:
        return None, None, None, 0, "no_face"

    def area(f):
        x1, y1, x2, y2 = f.bbox
        return (x2 - x1) * (y2 - y1)

    face = max(faces, key=area) if len(faces) > 1 else faces[0]
    emb = face.normed_embedding.astype(np.float32)
    x1, y1, x2, y2 = [float(v) for v in face.bbox]
    bbox = [x1, y1, x2 - x1, y2 - y1]
    status = "ok" if len(faces) == 1 else "ok_multi"
    return emb, float(face.det_score), bbox, len(faces), status