from helical.models.tahoe import Tahoe, TahoeConfig
import hydra
from omegaconf import DictConfig
import anndata as ad

@hydra.main(version_base=None, config_path="configs", config_name="tahoe_config")
def run(cfg: DictConfig):
    tahoe_config = TahoeConfig(**cfg)
    tahoe = Tahoe(configurer=tahoe_config)

    print(f"Loaded Tahoe model with {tahoe.model.n_layers} layers")
    print(f"Embedding dimension: {tahoe.config['d_model']}")

    ann_data = ad.read_h5ad("./yolksac_human.h5ad")

    # Process data - returns a DataLoader
    dataloader = tahoe.process_data(ann_data[:10])

    # Get cell embeddings from the DataLoader
    cell_embeddings = tahoe.get_embeddings(dataloader)
    print(f"Cell embeddings shape: {cell_embeddings.shape}")

    # Get both cell and gene embeddings
    cell_embeddings, gene_embeddings = tahoe.get_embeddings(
        dataloader, return_gene_embeddings=True
    )

    # Get embeddings with attention weights
    # Note: This requires using attn_impl='torch' instead of the default 'flash'
    # Only the last transformer layer's attention is returned to save memory
    # tahoe_config_torch = TahoeConfig(**cfg, attn_impl='torch')
    # tahoe_torch = Tahoe(configurer=tahoe_config_torch)
    # dataloader_torch = tahoe_torch.process_data(ann_data[:10])
    # cell_embeddings, attentions = tahoe_torch.get_embeddings(
    #     dataloader_torch, output_attentions=True
    # )
    # print(f"Cell embeddings shape: {cell_embeddings.shape}")
    # print(f"Attention shape: {attentions.shape}")  # (n_batches, batch_size, n_heads, seq_len, seq_len)

if __name__ == "__main__":
    run()
