"""
dino_similarity.py — objective, numeric image-similarity checks using
DINOv2 embeddings. Built specifically to catch the failure mode confirmed
on real hardware (cannon1): a retry loop can run all MAX_RETRIES attempts
with the judge's TEXT critique never being specific enough for Qwen to
actually produce a meaningfully different result — the loop burns through
every attempt with no real change, and nothing catches this because the
only feedback signal was itself another LLM's subjective read.

DINOv2 gives an OBJECTIVE alternative for one specific, narrow question:
did anything actually change, in pixel/feature space — independent of
whether any model's text says so clearly enough to act on.

Two distinct uses, kept separate since they answer different questions:
  - stagnation_score(): same keyframe index, ACROSS two attempts. High
    similarity = the retry did nothing, regardless of what the judge said.
  - temporal_consistency(): CONSECUTIVE keyframes, WITHIN one attempt.
    Reported as diagnostic numbers only for now — not used to gate
    behavior, since there's no real data yet to calibrate what threshold
    actually separates "smooth motion" from "erratic" or "static" in this
    embedding space. Log first, calibrate once real runs exist to look at.
"""

import numpy as np


def load_dino_model(model_name="facebook/dinov2-base"):
    from transformers import AutoImageProcessor, AutoModel
    print(f"loading DINOv2 from {model_name} ...")
    processor = AutoImageProcessor.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)
    model.eval()
    print("DINOv2 loaded.")
    return model, processor


def embed_image(model, processor, image):
    """image: a PIL Image. Returns a 1D numpy embedding vector (CLS token)."""
    import torch
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    # CLS token embedding — standard choice for a single whole-image descriptor
    cls_embedding = outputs.last_hidden_state[:, 0, :].squeeze(0)
    return cls_embedding.cpu().numpy()


def cosine_similarity(a, b):
    a, b = np.asarray(a), np.asarray(b)
    denom = (np.linalg.norm(a) * np.linalg.norm(b))
    if denom == 0:
        return 0.0
    return float(np.dot(a, b) / denom)


def stagnation_score(model, processor, previous_images, current_images):
    """
    previous_images, current_images: lists of PIL Images, same length,
    SAME keyframe index expected at each position (e.g. both length 5,
    index i = kf{i} from each attempt).

    Returns: dict with per-keyframe cosine similarity and an overall
    mean — high similarity (close to 1.0) means the two attempts produced
    near-identical images at that keyframe, i.e. the retry did not
    meaningfully change anything there.
    """
    if len(previous_images) != len(current_images):
        raise ValueError(f"length mismatch: {len(previous_images)} previous vs {len(current_images)} current")

    per_keyframe = []
    for i, (prev_img, curr_img) in enumerate(zip(previous_images, current_images)):
        prev_emb = embed_image(model, processor, prev_img)
        curr_emb = embed_image(model, processor, curr_img)
        sim = cosine_similarity(prev_emb, curr_emb)
        per_keyframe.append(sim)
        print(f"  kf{i}: attempt-to-attempt similarity = {sim:.4f}")

    mean_sim = float(np.mean(per_keyframe))
    return {"per_keyframe_similarity": per_keyframe, "mean_similarity": mean_sim}


def is_stagnant(stagnation_result, threshold=0.98):
    """
    threshold=0.98 is a STARTING GUESS, not a calibrated value — DINOv2
    cosine similarities for genuinely different images are typically well
    below this, but the exact right cutoff for "meaningfully changed vs.
    not" depends on real data this pipeline hasn't collected yet. Treat
    this as adjustable once you have real stagnant vs. non-stagnant runs
    to compare against, not as a validated constant.
    """
    return stagnation_result["mean_similarity"] >= threshold


def temporal_consistency(model, processor, images):
    """
    images: list of PIL Images for ONE attempt, in keyframe order.
    Returns: list of consecutive-pair cosine similarities (kf0-kf1,
    kf1-kf2, kf2-kf3, kf3-kf4) — DIAGNOSTIC ONLY, not currently used to
    gate any retry decision. Print/log these; don't threshold on them yet
    without real calibration data.
    """
    similarities = []
    embeddings = [embed_image(model, processor, img) for img in images]
    for i in range(len(embeddings) - 1):
        sim = cosine_similarity(embeddings[i], embeddings[i + 1])
        similarities.append(sim)
        print(f"  kf{i}->kf{i+1}: temporal similarity = {sim:.4f}")
    return similarities