# backend\ml\predict_scene_model.py

import argparse
import json
from pathlib import Path

import joblib
import numpy as np
import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor


CLIP_MODEL_NAME = "openai/clip-vit-base-patch32"


def empty_prediction(source="no_model"):
    return {
        "modelAvailable": False,
        "source": source,
        "peopleCount": 0,
        "maleCount": 0,
        "femaleCount": 0,
        "unknownCount": 0,
        "sexPosition": "unknown",
        "sexPositionScore": 0,
        "bodyPartsPresent": [],
        "objectsPresent": [],
        "clothingPresent": [],
        "boxes": [],
    }


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

    processor = CLIPProcessor.from_pretrained(CLIP_MODEL_NAME)
    model = CLIPModel.from_pretrained(CLIP_MODEL_NAME)
    model.eval()
    model.to(device)

    return model, processor, device


def image_features_to_tensor(model, out):
    if torch.is_tensor(out):
        return out

    if hasattr(out, "image_embeds") and out.image_embeds is not None:
        return out.image_embeds

    if hasattr(out, "pooler_output") and out.pooler_output is not None:
        emb = out.pooler_output

        # Nur projizieren, wenn pooler_output noch die erwartete Eingangsgröße hat.
        # Bei manchen Transformers-Versionen ist pooler_output bereits 512-dimensional.
        projection = getattr(model, "visual_projection", None)
        if projection is not None and hasattr(projection, "in_features"):
            if emb.shape[-1] == projection.in_features:
                emb = projection(emb)

        return emb

    if isinstance(out, (tuple, list)) and len(out) > 0:
        first = out[0]
        if torch.is_tensor(first):
            return first

    raise TypeError(f"Unsupported CLIP image feature output: {type(out)!r}")


def embed_image(model, processor, device, image_path: Path):
    img = Image.open(image_path).convert("RGB")

    inputs = processor(images=img, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        try:
            out = model.get_image_features(**inputs)
        except Exception:
            out = model.vision_model(pixel_values=inputs["pixel_values"])

        emb = image_features_to_tensor(model, out)

    emb = emb.detach().cpu().numpy()[0].astype("float32")

    norm = np.linalg.norm(emb)
    if norm > 0:
        emb = emb / norm

    return emb.reshape(1, -1)


def predict_with_model(model, emb):
    label = str(model.predict(emb)[0])

    score = 0.0
    if hasattr(model, "predict_proba"):
        probs = model.predict_proba(emb)[0]
        classes = [str(x) for x in model.classes_]

        if label in classes:
            score = float(probs[classes.index(label)])
        elif len(probs) > 0:
            score = float(np.max(probs))

    return label, score


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--image", required=True)
    args = parser.parse_args()

    root = Path(args.root)
    image_path = Path(args.image)

    model_dir = root / "model"
    lr_path = model_dir / "scene_clip_lr.joblib"
    knn_path = model_dir / "scene_clip_knn.joblib"

    if not lr_path.exists() and not knn_path.exists():
        print(json.dumps(empty_prediction("scene_clip_missing"), ensure_ascii=False))
        return

    try:
        clip_model, processor, device = load_clip()
        emb = embed_image(clip_model, processor, device, image_path)
    except Exception as e:
        out = empty_prediction("scene_clip_embed_failed")
        out["error"] = repr(e)
        print(json.dumps(out, ensure_ascii=False))
        return

    # Bevorzugt Logistic Regression, weil sie stabilere Wahrscheinlichkeiten liefert.
    # KNN bleibt Fallback, wenn nur eine Klasse oder sehr wenig Daten vorhanden sind.
    source = "scene_position_clip_lr"

    try:
        if lr_path.exists():
            model = joblib.load(lr_path)
            sex_position, score = predict_with_model(model, emb)
        else:
            source = "scene_position_clip_knn"
            model = joblib.load(knn_path)
            sex_position, score = predict_with_model(model, emb)
    except Exception as e:
        out = empty_prediction("scene_clip_predict_failed")
        out["error"] = repr(e)
        print(json.dumps(out, ensure_ascii=False))
        return

    if not sex_position:
        sex_position = "unknown"

    pred = {
        "modelAvailable": True,
        "source": source,
        "peopleCount": 0,
        "maleCount": 0,
        "femaleCount": 0,
        "unknownCount": 0,
        "sexPosition": sex_position,
        "sexPositionScore": float(score),
        "bodyPartsPresent": [],
        "objectsPresent": [],
        "clothingPresent": [],
        "boxes": [],
    }

    print(json.dumps(pred, ensure_ascii=False))


if __name__ == "__main__":
    main()