# backend\ml\train_scene_model.py

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from torchvision import models, transforms


def load_encoder():
    weights = models.ResNet18_Weights.DEFAULT
    model = models.resnet18(weights=weights)
    model.fc = torch.nn.Identity()
    model.eval()

    preprocess = weights.transforms()
    return model, preprocess


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 {
        "peopleCount": int(pred.get("peopleCount") or 0),
        "maleCount": int(pred.get("maleCount") or 0),
        "femaleCount": int(pred.get("femaleCount") or 0),
        "unknownCount": int(pred.get("unknownCount") or 0),
        "sexPosition": str(pred.get("sexPosition") or "unknown"),
        "bodyPartsPresent": [x.get("label") for x in pred.get("bodyPartsPresent") or [] if x.get("label")],
        "objectsPresent": [x.get("label") for x in pred.get("objectsPresent") or [] if x.get("label")],
        "clothingPresent": [x.get("label") for x in pred.get("clothingPresent") or [] if x.get("label")],
    }


def correction_target(annotation):
    corr = annotation.get("correction") or {}

    return {
        "peopleCount": int(corr.get("peopleCount") or 0),
        "maleCount": int(corr.get("maleCount") or 0),
        "femaleCount": int(corr.get("femaleCount") or 0),
        "unknownCount": int(corr.get("unknownCount") or 0),
        "sexPosition": str(corr.get("sexPosition") or "unknown"),
        "bodyPartsPresent": list(corr.get("bodyPartsPresent") or []),
        "objectsPresent": list(corr.get("objectsPresent") or []),
        "clothingPresent": list(corr.get("clothingPresent") or []),
    }


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

    return correction_target(annotation)


def embed_image(model, preprocess, image_path: Path):
    img = Image.open(image_path).convert("RGB")
    x = preprocess(img).unsqueeze(0)

    with torch.no_grad():
        emb = model(x).cpu().numpy()[0].astype("float32")

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

    return emb


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)

    model, preprocess = load_encoder()

    embeddings = []
    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

        try:
            emb = embed_image(model, preprocess, image_path)
            target = target_from_annotation(row)

            embeddings.append(emb)
            targets.append(target)
            used += 1
        except Exception:
            skipped += 1

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

    embeddings_np = np.stack(embeddings).astype("float32")

    np.savez_compressed(
        model_dir / "scene_knn.npz",
        embeddings=embeddings_np,
    )

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

    with (model_dir / "status.json").open("w", encoding="utf-8") as f:
        json.dump(
            {
                "ok": True,
                "usedSamples": used,
                "skippedSamples": skipped,
                "model": "resnet18_knn",
            },
            f,
            ensure_ascii=False,
            indent=2,
        )

    print(json.dumps({
        "ok": True,
        "usedSamples": used,
        "skippedSamples": skipped,
        "modelPath": str(model_dir / "scene_knn.npz"),
    }))


if __name__ == "__main__":
    main()