# ============================================================================
# INFERENCE TEST: Memory-Extended CLIP-L
#
# Tests both local loading (from checkpoint) and AutoModel loading.
# ============================================================================

import torch
import torch.nn.functional as F
import time

REPO_ID = "AbstractPhil/geolip-clip-vit-large-patch14-ctx576"
CHECKPOINT_DIR = "/home/claude/memory_clip_checkpoints"


def test_local():
    """Test with locally trained weights."""
    print("=" * 70)
    print("TEST: Local Checkpoint Loading")
    print("=" * 70)

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

    from memory_model_configuration import MemoryCLIPConfig
    from memory_model_code import MemoryCLIPModel
    from safetensors.torch import load_file

    config = MemoryCLIPConfig()
    model = MemoryCLIPModel(config).to(device).eval()

    # Load trained weights
    ckpt = f"{CHECKPOINT_DIR}/final/memory_system.safetensors"
    try:
        state = load_file(ckpt, device=str(device))
        missing, unexpected = model.load_state_dict(state, strict=False)
        print(f"  Loaded: {ckpt}")
        print(f"  Missing: {len(missing)}, Unexpected: {len(unexpected)}")
    except Exception as e:
        print(f"  No checkpoint found ({e}), running with random weights")

    run_tests(model, device, "local")


def test_automodel():
    """Test via AutoModel from HuggingFace."""
    print("=" * 70)
    print("TEST: AutoModel from HuggingFace")
    print("=" * 70)

    from transformers import AutoModel
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    try:
        model = AutoModel.from_pretrained(REPO_ID, trust_remote_code=True)
        model = model.to(device).eval()
        print(f"  Loaded from: {REPO_ID}")
        run_tests(model, device, "automodel")
    except Exception as e:
        print(f"  AutoModel loading failed: {e}")
        print(f"  (Expected if repo not yet pushed)")


def run_tests(model, device, tag):
    """Shared test suite."""

    # ── Test 1: Single text encoding ──
    print(f"\n  [{tag}] Test 1: Single encoding")
    short = "A photo of a cat sitting on a windowsill"
    long = (
        "A vast sweeping landscape of rolling green hills under dramatic "
        "storm clouds with a lone oak tree in the foreground its branches "
        "bent by wind casting long shadows across a field of wildflowers "
        "in purple yellow and white while in the distance a medieval stone "
        "castle sits atop a cliff overlooking a turbulent sea with waves "
        "crashing against ancient rocks and seabirds wheeling overhead "
        "against a sky painted in shades of grey and gold as the sun "
        "breaks through the clouds illuminating the castle towers"
    )

    with torch.no_grad():
        t0 = time.time()
        emb_short = model.encode(short)
        t_short = time.time() - t0

        t0 = time.time()
        emb_long = model.encode(long)
        t_long = time.time() - t0

    n_short = len(model.clip_tokenizer.encode(short))
    n_long = len(model.clip_tokenizer.encode(long))
    cos = F.cosine_similarity(emb_short.unsqueeze(0), emb_long.unsqueeze(0)).item()

    print(f"    Short ({n_short} tok): {emb_short.shape}, {t_short:.3f}s")
    print(f"    Long  ({n_long} tok): {emb_long.shape}, {t_long:.3f}s")
    print(f"    Cosine(short, long): {cos:.4f}")

    # ── Test 2: Batch encoding ──
    print(f"\n  [{tag}] Test 2: Batch encoding")
    texts = [
        "A fluffy orange cat sleeping on a warm blanket",
        "A black and white dog playing fetch in a park",
        "A tropical beach with palm trees and turquoise water",
        "A snowy mountain peak with climbers ascending the ridge",
        "A bustling city street at night with neon signs",
        "An old library with floor to ceiling bookshelves",
    ]
    with torch.no_grad():
        t0 = time.time()
        embs = model.encode(texts)
        elapsed = time.time() - t0
    print(f"    {len(texts)} texts → {embs.shape} in {elapsed:.3f}s")

    # ── Test 3: Similarity matrix ──
    print(f"\n  [{tag}] Test 3: Similarity matrix")
    normed = F.normalize(embs, dim=-1)
    sim = normed @ normed.T
    labels = ["cat", "dog", "beach", "mountain", "city", "library"]
    print(f"    {'':8s}", end="")
    for l in labels:
        print(f"{l:>9s}", end="")
    print()
    for i, l in enumerate(labels):
        print(f"    {l:8s}", end="")
        for j in range(len(labels)):
            print(f"  {sim[i,j].item():.3f} ", end="")
        print()

    # ── Test 4: Long context retrieval ──
    print(f"\n  [{tag}] Test 4: Long context retrieval")
    query = "medieval castle on a cliff by the sea"
    candidates = [
        "A castle by the ocean",
        "A medieval stone castle perched on a cliff overlooking turbulent waters with seabirds",
        ("A vast sweeping landscape with a medieval stone castle sits atop a cliff "
         "overlooking a turbulent sea with waves crashing against ancient rocks "
         "and seabirds wheeling overhead"),
        "A cat sitting on a windowsill watching birds outside",
    ]
    with torch.no_grad():
        q_emb = model.encode(query)
        c_embs = model.encode(candidates)
        sims = F.cosine_similarity(q_emb.unsqueeze(0), c_embs)

    cand_labels = ["Short match", "Medium match", "Long (>77 tok)", "Distractor"]
    print(f"    Query: '{query}'")
    for label, s in zip(cand_labels, sims):
        print(f"      {s.item():.4f}  {label}")

    # ── Test 5: forward() API ──
    print(f"\n  [{tag}] Test 5: HuggingFace forward() API")
    with torch.no_grad():
        output = model(texts=["A cat on a mat", "A dog in the park"])
    print(f"    output type: {type(output).__name__}")
    print(f"    last_hidden_state: {output.last_hidden_state.shape}")

    # ── Test 6: Memory-CLIP vs standard CLIP ──
    print(f"\n  [{tag}] Test 6: Memory-CLIP vs truncated CLIP")
    long_caption = (
        "A meticulously arranged still life painting in the Dutch Golden Age style "
        "featuring a silver goblet overflowing with deep red wine next to a half "
        "peeled lemon with its rind spiraling downward a cracked walnut revealing "
        "its inner flesh a porcelain plate holding slices of rare roast beef and "
        "a small bouquet of wilting tulips in shades of pink and white all set "
        "against a dark moody background with dramatic chiaroscuro lighting"
    )
    n_tok = len(model.clip_tokenizer.encode(long_caption))
    with torch.no_grad():
        mem_emb = model.encode(long_caption)
        clip_inputs = model.clip_tokenizer(
            long_caption, max_length=77, truncation=True,
            padding="max_length", return_tensors="pt").to(device)
        clip_out = model.clip_text(**clip_inputs)
        clip_emb = clip_out.pooler_output.squeeze(0)
    cos = F.cosine_similarity(mem_emb.unsqueeze(0), clip_emb.unsqueeze(0)).item()
    print(f"    Caption: {n_tok} tokens ({max(0, n_tok-77)} beyond CLIP limit)")
    print(f"    Cosine(memory-CLIP, truncated-CLIP): {cos:.4f}")
    if cos < 0.9:
        print(f"    ★ Memory system adds information beyond truncation")

    print(f"\n  [{tag}] ALL TESTS PASSED")


if __name__ == "__main__":
    #test_local()
    print()
    test_automodel()