#!/usr/bin/env python3
"""GenomeClip usage examples.

This script demonstrates how to use GenomeClip to encode DNA sequence
embeddings and text embeddings into a shared 512-dim space.

Prerequisites:
    pip install torch transformers

Usage:
    python example.py
"""

import torch
import torch.nn.functional as F
from transformers import AutoModel

# ──────────────────────────────────────────────────────────────────────
# 1. Load GenomeClip
# ──────────────────────────────────────────────────────────────────────
MODEL_PATH = "your-username/GenomeClip-v1"  # or a local path

model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

print(f"Loaded GenomeClip on {device}")
print(f"  seq_input_dim  = {model.config.seq_input_dim}")
print(f"  text_input_dim = {model.config.text_input_dim}")
print(f"  projection_dim = {model.config.projection_dim}")
print(f"  seq_pooling    = {model.config.seq_pooling}")
print()


# ──────────────────────────────────────────────────────────────────────
# 2. Prepare dummy inputs (replace with real embeddings in practice)
# ──────────────────────────────────────────────────────────────────────
# In practice:
#   seq_emb  = AlphaGenome embeddings_128bp for each gene  → (L, 3072)
#   text_emb = OpenAI text-embedding-3-large for gene desc → (3072,)

batch_size = 4
seq_lengths_list = [120, 80, 50, 200]  # variable token counts per gene

# Simulate per-token sequence embeddings (L2-normalized, as used in training)
max_len = max(seq_lengths_list)
seq_emb = torch.randn(batch_size, max_len, 3072, device=device)
seq_emb = F.normalize(seq_emb, dim=-1)  # L2-normalize each token
seq_lengths = torch.tensor(seq_lengths_list, device=device)

# Simulate text embeddings
text_emb = torch.randn(batch_size, 3072, device=device)


# ──────────────────────────────────────────────────────────────────────
# 3. Encode sequence embeddings only
# ──────────────────────────────────────────────────────────────────────
with torch.no_grad():
    seq_repr = model.encode_sequence(seq_emb, seq_lengths)

print("=== Encode sequence only ===")
print(f"  Input:  seq_emb {tuple(seq_emb.shape)}, seq_lengths {tuple(seq_lengths.shape)}")
print(f"  Output: seq_repr {tuple(seq_repr.shape)}")
print(f"  Norms:  {seq_repr.norm(dim=-1).tolist()}  (should be ~1.0)")
print()


# ──────────────────────────────────────────────────────────────────────
# 4. Encode text embeddings only
# ──────────────────────────────────────────────────────────────────────
with torch.no_grad():
    text_repr = model.encode_text(text_emb)

print("=== Encode text only ===")
print(f"  Input:  text_emb {tuple(text_emb.shape)}")
print(f"  Output: text_repr {tuple(text_repr.shape)}")
print(f"  Norms:  {text_repr.norm(dim=-1).tolist()}  (should be ~1.0)")
print()


# ──────────────────────────────────────────────────────────────────────
# 5. Cross-modal similarity (cosine, since both are L2-normalized)
# ──────────────────────────────────────────────────────────────────────
similarity = seq_repr @ text_repr.t()  # (B, B)

print("=== Cross-modal similarity matrix ===")
print(f"  Shape: {tuple(similarity.shape)}")
for i in range(batch_size):
    row = ", ".join(f"{similarity[i, j]:.4f}" for j in range(batch_size))
    print(f"  seq[{i}] vs text[0..{batch_size-1}]: [{row}]")
print()


# ──────────────────────────────────────────────────────────────────────
# 6. Forward with both modalities (computes contrastive loss)
# ──────────────────────────────────────────────────────────────────────
with torch.no_grad():
    out = model(
        seq_embeddings=seq_emb,
        text_embeddings=text_emb,
        seq_lengths=seq_lengths,
    )

print("=== Forward with both modalities ===")
print(f"  seq_repr:  {tuple(out.seq_repr.shape)}")
print(f"  text_repr: {tuple(out.text_repr.shape)}")
print(f"  loss:      {out.loss.item():.4f}  (InfoNCE)")
print(f"  logits:    {tuple(out.logits.shape)}")
print()


# ──────────────────────────────────────────────────────────────────────
# 7. Pooled (mean) sequence input — already-pooled (3072,) vectors
# ──────────────────────────────────────────────────────────────────────
# If you already have gene-level pooled AlphaGenome embeddings (3072,),
# you can pass them directly (2D tensor). The model auto-expands to L=1.
seq_pooled = torch.randn(batch_size, 3072, device=device)
seq_pooled = F.normalize(seq_pooled, dim=-1)

with torch.no_grad():
    pooled_repr = model.encode_sequence(seq_pooled)  # no seq_lengths needed

print("=== Pooled (2D) sequence input ===")
print(f"  Input:  seq_pooled {tuple(seq_pooled.shape)}")
print(f"  Output: pooled_repr {tuple(pooled_repr.shape)}")
print()


# ──────────────────────────────────────────────────────────────────────
# 8. Cross-modal retrieval example
# ──────────────────────────────────────────────────────────────────────
print("=== Cross-modal retrieval ===")

# Simulate a database of 100 gene sequence embeddings
n_genes = 100
db_seq = F.normalize(torch.randn(n_genes, 3072, device=device), dim=-1)
with torch.no_grad():
    db_seq_repr = model.encode_sequence(db_seq)  # (100, 512)

# Query with a single text embedding
query_text = torch.randn(1, 3072, device=device)
with torch.no_grad():
    query_repr = model.encode_text(query_text)  # (1, 512)

# Find top-5 most similar genes
scores = (query_repr @ db_seq_repr.t()).squeeze(0)  # (100,)
top5_indices = scores.argsort(descending=True)[:5]
print(f"  Database: {n_genes} genes")
print(f"  Top-5 matching gene indices: {top5_indices.tolist()}")
print(f"  Top-5 scores: {scores[top5_indices].tolist()}")
print()

print("Done!")