"""comfyui_score_aesthetic — puntuacion estetica LAION-V2 (0-10) de una imagen.

Juez objetivo de calidad del panel `comfyui-judge`. Usa el predictor estetico
LAION "improved aesthetic" (head MLP lineal sobre embeddings CLIP ViT-L/14
normalizados): el .pth `sac+logos+ava1-l14-linearMSE.pth` entrenado sobre AVA +
SAC + LOGOS, que devuelve un score continuo ~1-10 donde mayor = mas atractiva.

Impura: el scoring necesita torch + open_clip + CUDA, que NO estan en el venv del
registry (`python/.venv`). Por eso esta funcion delega la inferencia a un SUBPROCESO
contra el python del venv de ComfyUI (`~/ComfyUI/.venv/bin/python3`), que si los
tiene. El subproceso reconstruye el MLP head, carga CLIP ViT-L-14-quickgelu
(pretrained openai, ya cacheado), embebe la imagen, la normaliza L2 y aplica el head.
La funcion captura un JSON marcado en stdout y nunca lanza excepcion.
"""
import json
import os
import subprocess
import sys

DEFAULT_MODEL_PATH = "/mnt/2tb/comfyui_models/aesthetic/sac+logos+ava1-l14-linearMSE.pth"
DEFAULT_VENV_PYTHON = "~/ComfyUI/.venv/bin/python3"

# Script autonomo ejecutado por el python del venv de ComfyUI (tiene torch + open_clip).
# Reconstruye el MLP head LAION-V2 y puntua la imagen. Emite el resultado en una
# linea marcada con _MARKER para aislarlo de warnings / barras de descarga.
_MARKER = "__AESTHETIC_RESULT__"
_INFER_SCRIPT = r'''
import json, sys
import numpy as np
import torch
import torch.nn as nn
import open_clip
from PIL import Image

MARKER = "__AESTHETIC_RESULT__"


class AestheticMLP(nn.Module):
    def __init__(self, input_size=768):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_size, 1024),
            nn.Dropout(0.2),
            nn.Linear(1024, 128),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.Dropout(0.1),
            nn.Linear(64, 16),
            nn.Linear(16, 1),
        )

    def forward(self, x):
        return self.layers(x)


def _l2norm(a, axis=-1):
    l2 = np.atleast_1d(np.linalg.norm(a, 2, axis))
    l2[l2 == 0] = 1
    return a / np.expand_dims(l2, axis)


def run(image_path, model_path, device):
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    model, _, preprocess = open_clip.create_model_and_transforms(
        "ViT-L-14-quickgelu", pretrained="openai")
    model = model.to(device).eval()
    head = AestheticMLP(768)
    head.load_state_dict(torch.load(model_path, map_location="cpu"))
    head = head.to(device).eval()

    img = preprocess(Image.open(image_path).convert("RGB")).unsqueeze(0).to(device)
    with torch.no_grad():
        feats = model.encode_image(img).cpu().numpy()
    norm = _l2norm(feats)
    with torch.no_grad():
        pred = head(torch.from_numpy(norm).to(device).float())
    return float(pred.cpu().numpy()[0][0])


if __name__ == "__main__":
    image_path, model_path, device = sys.argv[1], sys.argv[2], sys.argv[3]
    try:
        score = run(image_path, model_path, device)
        print(MARKER + json.dumps({"ok": True, "score_0_10": score, "error": ""}))
    except Exception as exc:
        print(MARKER + json.dumps({"ok": False, "score_0_10": 0.0, "error": repr(exc)}))
'''


def comfyui_score_aesthetic(
    image_path: str,
    *,
    model_path: str = DEFAULT_MODEL_PATH,
    venv_python: str = DEFAULT_VENV_PYTHON,
    device: str = "auto",
    timeout: float = 180.0,
) -> dict:
    """Puntua la calidad estetica (0-10) de una imagen con el predictor LAION-V2.

    Args:
        image_path: ruta a la imagen (PNG/JPG/WebP) en disco local.
        model_path: ruta al .pth del head estetico LAION (sac+logos+ava1-l14-linearMSE).
            keyword-only.
        venv_python: python del venv de ComfyUI con torch + open_clip; se expande `~`.
            La inferencia corre ahi por subproceso porque el venv del registry no tiene
            torch. keyword-only.
        device: "auto" (cuda si hay, si no cpu), "cuda" o "cpu". keyword-only.
        timeout: timeout del subproceso en segundos (la primera vez puede descargar CLIP).
            keyword-only.

    Returns:
        dict ``{ok, score_0_10, error}``. En exito ``ok=True`` y ``score_0_10`` es el
        score continuo (~1-10, mayor = mejor). En error ``ok=False``, ``score_0_10=0.0``
        y ``error`` describe la causa. Nunca lanza excepcion.
    """
    img = os.path.expanduser(image_path)
    if not os.path.isfile(img):
        return {"ok": False, "score_0_10": 0.0, "error": f"imagen no encontrada: {img}"}

    py = os.path.expanduser(venv_python)
    if not os.path.isfile(py):
        return {"ok": False, "score_0_10": 0.0,
                "error": f"python del venv ComfyUI no encontrado: {py}"}

    mp = os.path.expanduser(model_path)
    if not os.path.isfile(mp):
        return {"ok": False, "score_0_10": 0.0,
                "error": f"modelo estetico no encontrado: {mp}"}

    try:
        proc = subprocess.run(
            [py, "-c", _INFER_SCRIPT, img, mp, device],
            capture_output=True, text=True, timeout=timeout,
        )
    except subprocess.TimeoutExpired:
        return {"ok": False, "score_0_10": 0.0,
                "error": f"subproceso de inferencia excedio {timeout}s"}

    for line in proc.stdout.splitlines():
        if line.startswith(_MARKER):
            try:
                return json.loads(line[len(_MARKER):])
            except json.JSONDecodeError as exc:
                return {"ok": False, "score_0_10": 0.0,
                        "error": f"JSON invalido del subproceso: {exc}"}

    tail = (proc.stderr or proc.stdout or "")[-400:]
    return {"ok": False, "score_0_10": 0.0,
            "error": f"sin resultado del subproceso (rc={proc.returncode}): {tail}"}


def _main(argv):
    if not argv:
        sys.stderr.write("uso: comfyui_score_aesthetic IMAGEN [--device cpu|cuda]\n")
        return 2
    device = "auto"
    image_path = ""
    i = 0
    while i < len(argv):
        if argv[i] == "--device" and i + 1 < len(argv):
            device = argv[i + 1]
            i += 2
        else:
            image_path = argv[i]
            i += 1
    res = comfyui_score_aesthetic(image_path, device=device)
    print(json.dumps(res, indent=2, ensure_ascii=False))
    return 0 if res["ok"] else 1


if __name__ == "__main__":
    sys.exit(_main(sys.argv[1:]))
