#!/usr/bin/env python3
"""
GeneSetCLIP Training Job — Self-contained script for HF Jobs.

Downloads MSigDB data from Hub, then trains the contrastive model.
"""

import os
import sys

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Step 1: Download data from Hub
print("=" * 70)
print("Step 1: Downloading MSigDB data from Hub...")
print("=" * 70)

from huggingface_hub import hf_hub_download

DATA_DIR = "/tmp/data"
os.makedirs(DATA_DIR, exist_ok=True)

for split in ["train", "val", "test"]:
    path = hf_hub_download(
        repo_id="AliSaadatV/msigdb-contrastive-data",
        filename=f"{split}.jsonl",
        repo_type="dataset",
        local_dir=DATA_DIR,
    )
    size_mb = os.path.getsize(path) / 1024 / 1024
    print(f"  {split}.jsonl: {size_mb:.1f} MB")

print("Data downloaded!\n")

# Step 2: Import and configure training
import json
import math
import random
import time
from collections import defaultdict
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
import numpy as np
from huggingface_hub import HfApi
import trackio


# ============================================================
# Configuration
# ============================================================

@dataclass
class Config:
    gene_model_id: str = "maayanlab/gsfm-rummagene"
    text_model_id: str = "FremyCompany/BioLORD-2023"
    shared_dim: int = 256
    gene_dim: int = 256
    text_dim: int = 768
    proj_hidden_dim: int = 512
    proj_dropout: float = 0.1
    batch_size: int = 256
    lr: float = 1e-4
    gene_encoder_lr: float = 1e-5
    weight_decay: float = 0.01
    warmup_steps: int = 500
    max_epochs: int = 50
    patience: int = 10
    temperature_init: float = 0.07
    learnable_temperature: bool = True
    gene_dropout_rate: float = 0.2
    max_gene_set_size: int = 512
    data_dir: str = DATA_DIR
    output_dir: str = "/tmp/output"
    hub_model_id: str = "AliSaadatV/GeneSetCLIP"
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
    num_workers: int = 4
    mixed_precision: bool = True
    log_every: int = 10
    eval_every: int = 1
    save_every: int = 5


# ============================================================
# Dataset
# ============================================================

class GeneSetTextDataset(Dataset):
    def __init__(self, jsonl_path, vocab, max_genes=512, gene_dropout=0.0, pad_idx=1):
        self.records = []
        with open(jsonl_path) as f:
            for line in f:
                self.records.append(json.loads(line))
        self.vocab = vocab
        self.max_genes = max_genes
        self.gene_dropout = gene_dropout
        self.pad_idx = pad_idx

    def __len__(self):
        return len(self.records)

    def __getitem__(self, idx):
        record = self.records[idx]
        text = record["text"]
        genes = record["genes"]
        token_ids = [self.vocab.get(g, 0) for g in genes]

        if self.gene_dropout > 0:
            n_keep = max(3, int(len(token_ids) * (1 - self.gene_dropout)))
            if n_keep < len(token_ids):
                token_ids = random.sample(token_ids, n_keep)

        if len(token_ids) > self.max_genes:
            token_ids = random.sample(token_ids, self.max_genes)

        n_genes = len(token_ids)
        if n_genes < self.max_genes:
            token_ids = token_ids + [self.pad_idx] * (self.max_genes - n_genes)

        return {
            "text": text,
            "gene_ids": torch.tensor(token_ids, dtype=torch.long),
            "n_genes": n_genes,
            "id": record["id"],
        }


def collate_fn(batch):
    return {
        "text": [item["text"] for item in batch],
        "gene_ids": torch.stack([item["gene_ids"] for item in batch]),
        "n_genes": torch.tensor([item["n_genes"] for item in batch]),
        "ids": [item["id"] for item in batch],
    }


# ============================================================
# Model
# ============================================================

class ProjectionHead(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, output_dim),
            nn.LayerNorm(output_dim),
        )

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


class GeneSetCLIP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.log_temperature = nn.Parameter(
            torch.log(torch.tensor(config.temperature_init)),
            requires_grad=config.learnable_temperature,
        )
        self.text_proj = ProjectionHead(config.text_dim, config.proj_hidden_dim,
                                        config.shared_dim, config.proj_dropout)
        self.gene_proj = ProjectionHead(config.gene_dim, config.shared_dim,
                                        config.shared_dim, config.proj_dropout)

    @property
    def temperature(self):
        return torch.clamp(self.log_temperature.exp(), min=0.01, max=1.0)

    def forward(self, gene_emb, text_emb):
        z_gene = F.normalize(self.gene_proj(gene_emb), dim=-1)
        z_text = F.normalize(self.text_proj(text_emb), dim=-1)
        tau = self.temperature
        logits = z_gene @ z_text.T / tau
        B = logits.size(0)
        labels = torch.arange(B, device=logits.device)
        loss_g2t = F.cross_entropy(logits, labels)
        loss_t2g = F.cross_entropy(logits.T, labels)
        loss = (loss_g2t + loss_t2g) / 2
        with torch.no_grad():
            g2t_acc = (logits.argmax(dim=1) == labels).float().mean()
            t2g_acc = (logits.T.argmax(dim=1) == labels).float().mean()
            avg_acc = (g2t_acc + t2g_acc) / 2
        metrics = {
            "loss": loss.item(), "g2t_acc": g2t_acc.item(),
            "t2g_acc": t2g_acc.item(), "avg_acc": avg_acc.item(),
            "temperature": tau.item(),
        }
        return loss, z_gene, z_text, metrics

    def get_embeddings(self, gene_emb=None, text_emb=None):
        z_gene = z_text = None
        if gene_emb is not None:
            z_gene = F.normalize(self.gene_proj(gene_emb), dim=-1)
        if text_emb is not None:
            z_text = F.normalize(self.text_proj(text_emb), dim=-1)
        return z_gene, z_text


# ============================================================
# Evaluation
# ============================================================

@torch.no_grad()
def evaluate_retrieval(model, gene_encoder, text_encoder, dataloader, device):
    model.eval()
    gene_encoder.eval()
    all_z_gene, all_z_text, all_ids = [], [], []
    total_loss, n_batches = 0, 0

    for batch in dataloader:
        gene_ids = batch["gene_ids"].to(device)
        texts = batch["text"]
        gene_emb = gene_encoder.encode(gene_ids)
        text_emb = text_encoder.encode(texts, convert_to_tensor=True, show_progress_bar=False)
        if text_emb.device != device:
            text_emb = text_emb.to(device)
        text_emb = text_emb.clone()
        loss, z_gene, z_text, _ = model(gene_emb, text_emb)
        total_loss += loss.item()
        n_batches += 1
        all_z_gene.append(z_gene.cpu())
        all_z_text.append(z_text.cpu())
        all_ids.extend(batch["ids"])

    all_z_gene = torch.cat(all_z_gene, dim=0)
    all_z_text = torch.cat(all_z_text, dim=0)
    N = len(all_z_gene)
    sim = all_z_gene @ all_z_text.T
    labels = torch.arange(N)

    def recall_at_k(sim_matrix, labels, k):
        topk = sim_matrix.topk(min(k, sim_matrix.size(1)), dim=1).indices
        return (topk == labels.unsqueeze(1)).any(dim=1).float().mean().item()

    def mrr(sim_matrix, labels):
        ranks = (sim_matrix.argsort(dim=1, descending=True) == labels.unsqueeze(1)).nonzero()[:, 1] + 1
        return (1.0 / ranks.float()).mean().item()

    results = {
        "loss": total_loss / max(n_batches, 1), "n_samples": N,
        "g2t_R@1": recall_at_k(sim, labels, 1),
        "g2t_R@5": recall_at_k(sim, labels, 5),
        "g2t_R@10": recall_at_k(sim, labels, 10),
        "g2t_MRR": mrr(sim, labels),
        "t2g_R@1": recall_at_k(sim.T, labels, 1),
        "t2g_R@5": recall_at_k(sim.T, labels, 5),
        "t2g_R@10": recall_at_k(sim.T, labels, 10),
        "t2g_MRR": mrr(sim.T, labels),
    }
    results["avg_R@1"] = (results["g2t_R@1"] + results["t2g_R@1"]) / 2
    results["avg_R@5"] = (results["g2t_R@5"] + results["t2g_R@5"]) / 2
    results["avg_R@10"] = (results["g2t_R@10"] + results["t2g_R@10"]) / 2
    results["avg_MRR"] = (results["g2t_MRR"] + results["t2g_MRR"]) / 2

    model.train()
    return results


# ============================================================
# Training
# ============================================================

def train(config):
    print("=" * 70)
    print("GeneSetCLIP Training")
    print("=" * 70)
    print(f"Device: {config.device}")
    print(f"Batch size: {config.batch_size}")
    print(f"Max epochs: {config.max_epochs}")

    os.makedirs(config.output_dir, exist_ok=True)

    # Load GSFM
    print("\nLoading GSFM gene encoder...")
    from gsfm import GSFM, Vocab
    vocab_obj = Vocab.from_pretrained(config.gene_model_id)
    gene_encoder = GSFM.from_pretrained(config.gene_model_id)
    gene_encoder.to(config.device)
    gene_encoder.train()
    vocab_dict = {token: i for i, token in enumerate(vocab_obj.vocab)}
    print(f"  GSFM vocab: {len(vocab_dict)} genes")

    # Load BioLORD (frozen)
    print("Loading BioLORD text encoder (frozen)...")
    from sentence_transformers import SentenceTransformer
    text_encoder = SentenceTransformer(config.text_model_id, device=config.device)
    for param in text_encoder.parameters():
        param.requires_grad = False
    text_encoder.eval()

    # Model
    print("Building GeneSetCLIP model...")
    model = GeneSetCLIP(config).to(config.device)
    print(f"  Projection params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
    print(f"  Gene encoder params: {sum(p.numel() for p in gene_encoder.parameters()):,}")

    # Data
    print("\nLoading datasets...")
    train_ds = GeneSetTextDataset(os.path.join(config.data_dir, "train.jsonl"),
                                   vocab_dict, config.max_gene_set_size, config.gene_dropout_rate)
    val_ds = GeneSetTextDataset(os.path.join(config.data_dir, "val.jsonl"),
                                 vocab_dict, config.max_gene_set_size, 0.0)
    test_ds = GeneSetTextDataset(os.path.join(config.data_dir, "test.jsonl"),
                                  vocab_dict, config.max_gene_set_size, 0.0)
    print(f"  Train: {len(train_ds)}, Val: {len(val_ds)}, Test: {len(test_ds)}")

    train_loader = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True,
                              collate_fn=collate_fn, num_workers=config.num_workers,
                              pin_memory=True, drop_last=True)
    val_loader = DataLoader(val_ds, batch_size=config.batch_size, shuffle=False,
                            collate_fn=collate_fn, num_workers=config.num_workers)
    test_loader = DataLoader(test_ds, batch_size=config.batch_size, shuffle=False,
                             collate_fn=collate_fn, num_workers=config.num_workers)

    steps_per_epoch = len(train_loader)
    total_steps = steps_per_epoch * config.max_epochs
    print(f"  Steps/epoch: {steps_per_epoch}, Total: {total_steps}")

    # Optimizer
    optimizer = torch.optim.AdamW([
        {"params": list(model.text_proj.parameters()) + list(model.gene_proj.parameters()) +
                   [model.log_temperature], "lr": config.lr, "weight_decay": config.weight_decay},
        {"params": gene_encoder.parameters(), "lr": config.gene_encoder_lr,
         "weight_decay": config.weight_decay},
    ])

    def lr_lambda(step):
        if step < config.warmup_steps:
            return step / max(config.warmup_steps, 1)
        progress = (step - config.warmup_steps) / max(total_steps - config.warmup_steps, 1)
        return 0.5 * (1 + math.cos(math.pi * progress))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    scaler = torch.amp.GradScaler('cuda') if config.mixed_precision and config.device == "cuda" else None

    # Tracking
    trackio.init(project="GeneSetCLIP",
                 name=f"bs{config.batch_size}_lr{config.lr}_temp{config.temperature_init}")

    # Training loop
    best_val_mrr = 0
    patience_counter = 0
    global_step = 0

    for epoch in range(1, config.max_epochs + 1):
        model.train()
        gene_encoder.train()
        epoch_loss, epoch_acc, n_batches = 0, 0, 0

        for batch in train_loader:
            gene_ids = batch["gene_ids"].to(config.device)
            texts = batch["text"]

            gene_emb = gene_encoder.encode(gene_ids)
            with torch.no_grad():
                text_emb = text_encoder.encode(texts, convert_to_tensor=True, show_progress_bar=False)
                if text_emb.device != torch.device(config.device):
                    text_emb = text_emb.to(config.device)
                text_emb = text_emb.clone()

            if scaler is not None:
                with torch.amp.autocast('cuda'):
                    loss, _, _, metrics = model(gene_emb, text_emb)
                optimizer.zero_grad()
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(
                    list(model.parameters()) + list(gene_encoder.parameters()), 1.0)
                scaler.step(optimizer)
                scaler.update()
            else:
                loss, _, _, metrics = model(gene_emb, text_emb)
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(
                    list(model.parameters()) + list(gene_encoder.parameters()), 1.0)
                optimizer.step()

            scheduler.step()
            global_step += 1
            epoch_loss += metrics["loss"]
            epoch_acc += metrics["avg_acc"]
            n_batches += 1

            if global_step % config.log_every == 0:
                lr_p = optimizer.param_groups[0]["lr"]
                lr_g = optimizer.param_groups[1]["lr"]
                print(f"  Step {global_step:5d} | Loss: {metrics['loss']:.4f} | "
                      f"Acc: {metrics['avg_acc']:.3f} | τ: {metrics['temperature']:.4f} | "
                      f"LR: {lr_p:.2e}/{lr_g:.2e}")
                trackio.log({
                    "train/loss": metrics["loss"], "train/avg_acc": metrics["avg_acc"],
                    "train/g2t_acc": metrics["g2t_acc"], "train/t2g_acc": metrics["t2g_acc"],
                    "train/temperature": metrics["temperature"],
                    "train/lr_proj": lr_p, "train/lr_gene": lr_g, "step": global_step,
                })

        avg_loss = epoch_loss / max(n_batches, 1)
        avg_acc = epoch_acc / max(n_batches, 1)
        print(f"\nEpoch {epoch}/{config.max_epochs} | Loss: {avg_loss:.4f} | Acc: {avg_acc:.3f}")

        # Validation
        if epoch % config.eval_every == 0:
            print("  Evaluating...")
            val_results = evaluate_retrieval(model, gene_encoder, text_encoder,
                                             val_loader, config.device)
            print(f"  Val R@1: {val_results['avg_R@1']:.3f} | R@5: {val_results['avg_R@5']:.3f} | "
                  f"MRR: {val_results['avg_MRR']:.3f}")
            trackio.log({
                "val/loss": val_results["loss"],
                "val/avg_R@1": val_results["avg_R@1"],
                "val/avg_R@5": val_results["avg_R@5"],
                "val/avg_R@10": val_results["avg_R@10"],
                "val/avg_MRR": val_results["avg_MRR"],
                "epoch": epoch,
            })

            if val_results["avg_MRR"] > best_val_mrr:
                best_val_mrr = val_results["avg_MRR"]
                patience_counter = 0
                save_dir = os.path.join(config.output_dir, "best_model")
                os.makedirs(save_dir, exist_ok=True)
                torch.save(model.state_dict(), os.path.join(save_dir, "clip_model.pt"))
                torch.save(gene_encoder.state_dict(), os.path.join(save_dir, "gene_encoder.pt"))
                with open(os.path.join(save_dir, "config.json"), "w") as f:
                    json.dump(vars(config), f, indent=2)
                print(f"  ✓ New best! MRR: {best_val_mrr:.4f}")
            else:
                patience_counter += 1
                print(f"  No improvement ({patience_counter}/{config.patience})")
                if patience_counter >= config.patience:
                    print(f"  Early stopping at epoch {epoch}")
                    break

    # Load best model
    print("\n" + "=" * 70)
    print("Final test evaluation...")
    best_path = os.path.join(config.output_dir, "best_model")
    if os.path.exists(best_path):
        model.load_state_dict(torch.load(os.path.join(best_path, "clip_model.pt"),
                                         map_location=config.device, weights_only=True))
        gene_encoder.load_state_dict(torch.load(os.path.join(best_path, "gene_encoder.pt"),
                                                map_location=config.device, weights_only=True))

    test_results = evaluate_retrieval(model, gene_encoder, text_encoder,
                                       test_loader, config.device)
    print(f"Test Results:")
    print(f"  G→T R@1: {test_results['g2t_R@1']:.3f}  R@5: {test_results['g2t_R@5']:.3f}  R@10: {test_results['g2t_R@10']:.3f}  MRR: {test_results['g2t_MRR']:.3f}")
    print(f"  T→G R@1: {test_results['t2g_R@1']:.3f}  R@5: {test_results['t2g_R@5']:.3f}  R@10: {test_results['t2g_R@10']:.3f}  MRR: {test_results['t2g_MRR']:.3f}")
    print(f"  Avg R@1: {test_results['avg_R@1']:.3f}  R@5: {test_results['avg_R@5']:.3f}  MRR: {test_results['avg_MRR']:.3f}")
    trackio.log({"test/" + k: v for k, v in test_results.items()})

    # Push to Hub
    print("\nPushing to Hub...")
    api = HfApi()
    try:
        api.create_repo(config.hub_model_id, exist_ok=True)
    except Exception as e:
        print(f"  Warning: {e}")

    upload_dir = os.path.join(config.output_dir, "hub_upload")
    os.makedirs(upload_dir, exist_ok=True)
    torch.save(model.state_dict(), os.path.join(upload_dir, "clip_model.pt"))
    torch.save(gene_encoder.state_dict(), os.path.join(upload_dir, "gene_encoder.pt"))
    with open(os.path.join(upload_dir, "config.json"), "w") as f:
        json.dump(vars(config), f, indent=2)
    with open(os.path.join(upload_dir, "vocab.json"), "w") as f:
        json.dump(vocab_dict, f)
    with open(os.path.join(upload_dir, "test_results.json"), "w") as f:
        json.dump(test_results, f, indent=2)

    readme = f"""# GeneSetCLIP

Contrastive model aligning gene-set embeddings (GSFM) with biomedical text descriptions (BioLORD-2023).

## Architecture
- **Gene encoder**: [GSFM](https://huggingface.co/maayanlab/gsfm-rummagene) (MLP autoencoder, 256-dim)
- **Text encoder**: [BioLORD-2023](https://huggingface.co/FremyCompany/BioLORD-2023) (768-dim, frozen)
- **Projection heads**: Maps both modalities to shared 256-dim space
- **Loss**: Symmetric InfoNCE with learnable temperature

## Training Data
- **MSigDB v2024.1** (Human + Mouse): ~50,000 gene set-text pairs
- Collections: H, C1-C8 (Human), MH, M1-M8 (Mouse)
- Train: C2/C5/C8/C1 | Val: C3/C4 | Test: H/C6/C7

## Test Results (H, C6, C7 — {test_results['n_samples']} gene sets)
| Metric | Gene→Text | Text→Gene | Average |
|--------|-----------|-----------|---------|
| R@1    | {test_results['g2t_R@1']:.3f}     | {test_results['t2g_R@1']:.3f}     | {test_results['avg_R@1']:.3f}   |
| R@5    | {test_results['g2t_R@5']:.3f}     | {test_results['t2g_R@5']:.3f}     | {test_results['avg_R@5']:.3f}   |
| R@10   | {test_results['g2t_R@10']:.3f}    | {test_results['t2g_R@10']:.3f}    | {test_results['avg_R@10']:.3f}  |
| MRR    | {test_results['g2t_MRR']:.3f}     | {test_results['t2g_MRR']:.3f}     | {test_results['avg_MRR']:.3f}   |

## Usage

```python
import torch
from gsfm import GSFM, Vocab
from sentence_transformers import SentenceTransformer
from huggingface_hub import hf_hub_download

# Load gene encoder + vocab
gene_encoder = GSFM.from_pretrained("maayanlab/gsfm-rummagene")
vocab = Vocab.from_pretrained("maayanlab/gsfm-rummagene")
gene_encoder.eval()

# Load text encoder
text_encoder = SentenceTransformer("FremyCompany/BioLORD-2023")

# Load GeneSetCLIP projection heads
clip_path = hf_hub_download("AliSaadatV/GeneSetCLIP", "clip_model.pt")
config_path = hf_hub_download("AliSaadatV/GeneSetCLIP", "config.json")

import json
with open(config_path) as f:
    cfg = json.load(f)

# Reconstruct model (small — just projection heads)
import torch.nn as nn, torch.nn.functional as F

class ProjectionHead(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.1):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.GELU(),
                                 nn.Dropout(dropout), nn.Linear(hidden_dim, output_dim),
                                 nn.LayerNorm(output_dim))
    def forward(self, x): return self.net(x)

class GeneSetCLIP(nn.Module):
    def __init__(self):
        super().__init__()
        self.log_temperature = nn.Parameter(torch.zeros(1))
        self.text_proj = ProjectionHead(768, 512, 256, 0.1)
        self.gene_proj = ProjectionHead(256, 256, 256, 0.1)

clip_model = GeneSetCLIP()
clip_model.load_state_dict(torch.load(clip_path, map_location="cpu", weights_only=True))
clip_model.eval()

# Encode a gene set
genes = ["TP53", "BRCA1", "EGFR", "MYC", "KRAS"]
gene_ids = torch.tensor([vocab(genes)])
with torch.no_grad():
    gene_emb = gene_encoder.encode(gene_ids)
    z_gene = F.normalize(clip_model.gene_proj(gene_emb), dim=-1)

# Encode text
text_emb = text_encoder.encode(["Tumor suppressor genes involved in cancer"],
                                convert_to_tensor=True)
with torch.no_grad():
    z_text = F.normalize(clip_model.text_proj(text_emb), dim=-1)

# Similarity
print(f"Similarity: {{(z_gene @ z_text.T).item():.3f}}")
```
"""
    with open(os.path.join(upload_dir, "README.md"), "w") as f:
        f.write(readme)

    api.upload_folder(folder_path=upload_dir, repo_id=config.hub_model_id,
                      commit_message="Upload GeneSetCLIP trained model")
    print(f"  Pushed to https://huggingface.co/{config.hub_model_id}")
    print("\nDone!")


if __name__ == "__main__":
    config = Config()
    train(config)