# backend\ml\train_scene_model.py

import argparse
import json
from pathlib import Path

import joblib
import numpy as np
import torch
from PIL import Image
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from transformers import CLIPModel, CLIPProcessor


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


def read_jsonl(path: Path):
    if not path.exists():
        return []

    rows = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            try:
                rows.append(json.loads(line))
            except Exception:
                pass

    return rows


def prediction_target(annotation):
    pred = annotation.get("prediction") or {}
    return str(pred.get("sexPosition") or "unknown").strip() or "unknown"


def correction_target(annotation):
    corr = annotation.get("correction") or {}
    return str(corr.get("sexPosition") or "unknown").strip() or "unknown"


def target_from_annotation(annotation):
    if annotation.get("accepted") is True:
        return prediction_target(annotation)

    return correction_target(annotation)


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


def train_lr_if_possible(x, y):
    classes = sorted(set(y))

    if len(classes) < 2:
        return None

    # Logistic Regression braucht mindestens zwei Klassen.
    # class_weight hilft bei unausgeglichenen Positionen.
    clf = LogisticRegression(
        max_iter=2000,
        class_weight="balanced",
        solver="lbfgs",
    )

    clf.fit(x, y)
    return clf


def train_knn(x, y):
    n_neighbors = min(7, len(y))

    clf = KNeighborsClassifier(
        n_neighbors=n_neighbors,
        metric="cosine",
        weights="distance",
        algorithm="brute",
    )

    clf.fit(x, y)
    return clf


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

    root = Path(args.root)
    feedback_path = root / "feedback.jsonl"
    frames_dir = root / "frames"
    model_dir = root / "model"
    model_dir.mkdir(parents=True, exist_ok=True)

    rows = read_jsonl(feedback_path)

    clip_model, processor, device = load_clip()

    embeddings = []
    labels = []
    targets = []
    used = 0
    skipped = 0

    for row in rows:
        sample_id = str(row.get("sampleId") or "").strip()
        if not sample_id:
            skipped += 1
            continue

        image_path = frames_dir / f"{sample_id}.jpg"
        if not image_path.exists():
            skipped += 1
            continue

        label = target_from_annotation(row)
        if not label:
            label = "unknown"

        try:
            emb = embed_image(clip_model, processor, device, image_path)

            embeddings.append(emb)
            labels.append(label)
            targets.append({
                "sampleId": sample_id,
                "sexPosition": label,
            })
            used += 1
        except Exception as e:
            print(f"skip {sample_id}: {repr(e)}")
            skipped += 1

    if used < 5:
        raise SystemExit(f"too few usable samples: {used}")

    x = np.stack(embeddings).astype("float32")
    y = np.array(labels)

    np.savez_compressed(
        model_dir / "scene_clip_embeddings.npz",
        embeddings=x,
        labels=y,
    )

    with (model_dir / "scene_clip_targets.json").open("w", encoding="utf-8") as f:
        json.dump(targets, f, ensure_ascii=False, indent=2)

    knn = train_knn(x, y)
    joblib.dump(knn, model_dir / "scene_clip_knn.joblib")

    lr_status = "skipped_single_class"
    lr = train_lr_if_possible(x, y)
    if lr is not None:
        joblib.dump(lr, model_dir / "scene_clip_lr.joblib")
        lr_status = "trained"
    else:
        old_lr = model_dir / "scene_clip_lr.joblib"
        if old_lr.exists():
            old_lr.unlink()

    counts = {}
    for label in labels:
        counts[label] = counts.get(label, 0) + 1

    status = {
        "ok": True,
        "usedSamples": used,
        "skippedSamples": skipped,
        "model": "scene_position_clip",
        "clipModel": CLIP_MODEL_NAME,
        "device": device,
        "classes": sorted(counts.keys()),
        "classCounts": counts,
        "logisticRegression": lr_status,
        "knn": "trained",
        "embeddingsPath": str(model_dir / "scene_clip_embeddings.npz"),
        "knnPath": str(model_dir / "scene_clip_knn.joblib"),
        "lrPath": str(model_dir / "scene_clip_lr.joblib"),
    }

    with (model_dir / "status.json").open("w", encoding="utf-8") as f:
        json.dump(status, f, ensure_ascii=False, indent=2)


if __name__ == "__main__":
    main()