import scanpy as sc
from ccvgae import CCVGAE_agent

adata = sc.read_h5ad("your_data.h5ad")

agent = CCVGAE_agent(
    adata=adata,
    layer='counts',          # Use the 'counts' layer for training
    n_var=2000,              # Number of highly variable genes to use
    n_neighbors=15,          # Number of neighbors for graph construction
    latent_dim=10,           # Dimension of the latent space
    subgraph_size=512,       # Size of subgraphs for training
    num_subgraphs_per_epoch=10, # Number of subgraphs per epoch
    lr=1e-4,
    w_recon=1.0,             # Weight for reconstruction loss
    w_irecon=1.0,            # Weight for coupling reconstruction loss
    w_kl=1.0,                # Weight for KL divergence loss
)

# Train the model
agent.fit(epochs=300, update_steps=10)

latent_representation = agent.get_latent()

from sklearn.cluster import KMeans
clusters = KMeans(n_clusters=10).fit_predict(latent_representation)
adata.obs['ccvgae_clusters'] = clusters

# Example with custom preprocessing options
agent = CCVGAE_agent(
    adata=adata,
    layer='counts',           # Layer containing raw counts
    n_var=3000,               # Select top 3000 highly variable genes
    tech='PCA',               # Dimensionality reduction method
    n_neighbors=20,           # Number of neighbors for KNN graph
    batch_tech='harmony',     # Optional: use Harmony for batch correction
    all_feat=False,           # Use only highly variable genes
)

agent = CCVGAE_agent(
    adata=adata,
    # Architecture parameters
    encoder_type='graph',              # 'graph' or 'linear'
    graph_type='GAT',                  # 'GAT', 'GCN', 'SAGE', 'Cheb', etc.
    hidden_dim=128,                    # Hidden layer dimension
    latent_dim=10,                     # Latent space dimension
    i_dim=10,                          # Intermediate coupling dimension
    hidden_layers=2,                   # Number of hidden layers
    decoder_hidden_dim=128,            # Decoder hidden dimension
    dropout=0.05,                      # Dropout rate
    use_residual=True,                 # Use residual connections
    structure_decoder_type='mlp',      # 'mlp', 'bilinear', 'inner_product'
    feature_decoder_type='linear',     # Feature decoder type
)

agent = CCVGAE_agent(
    adata=adata,
    # Training parameters
    lr=1e-4,                          # Learning rate
    beta=1.0,                         # KL divergence regularization strength (β-VAE)
    graph=1.0,                        # Graph reconstruction loss weight
    w_recon=1.0,                      # Feature reconstruction loss weight
    w_irecon=1.0,                     # Coupling reconstruction loss weight
    w_kl=1.0,                         # KL divergence loss weight
    w_adj=1.0,                        # Adjacency reconstruction loss weight
    
    # Subgraph sampling for large datasets
    subgraph_size=512,                # Maximum nodes per subgraph
    num_subgraphs_per_epoch=10,       # Number of subgraphs per epoch
    sampling_method='random',          # Sampling strategy
    
    # Device configuration
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
    latent_type='q_m',                # 'q_m' (deterministic) or 'q_z' (stochastic)
)

# Train the model
agent.fit(epochs=300, update_steps=10, silent=False)

import scanpy as sc
from ccvgae import CCVGAE_agent

# Load data with multiple batches
adata = sc.read_h5ad("multi_batch_data.h5ad")

# Ensure batch information is in adata.obs
# adata.obs['batch'] should contain batch labels

# Initialize with batch correction
agent = CCVGAE_agent(
    adata=adata,
    layer='counts',
    n_var=2000,
    batch_tech='harmony',  # Use Harmony for batch correction
    n_neighbors=15,
    latent_dim=20,
)

agent.fit(epochs=500, update_steps=20)
latent = agent.get_latent()

# Add latent representation to adata for visualization
adata.obsm['X_ccvgae'] = latent
sc.pp.neighbors(adata, use_rep='X_ccvgae')
sc.tl.umap(adata)
sc.pl.umap(adata, color=['batch', 'cell_type'])

# For datasets with > 10,000 cells, use subgraph sampling
agent = CCVGAE_agent(
    adata=large_adata,
    layer='counts',
    n_var=2000,
    latent_dim=15,
    subgraph_size=1024,           # Larger subgraphs
    num_subgraphs_per_epoch=20,   # More subgraphs per epoch
)

agent.fit(epochs=300, update_steps=10)

# Fine-tune loss weights for specific analysis needs
agent = CCVGAE_agent(
    adata=adata,
    layer='counts',
    n_var=2000,
    latent_dim=10,
    w_recon=1.0,      # Standard reconstruction
    w_irecon=0.5,     # Reduce coupling weight
    w_kl=0.8,         # Reduce KL weight for more flexibility
    w_adj=1.2,        # Emphasize graph structure learning
    beta=0.5,         # Lighter regularization
)

agent.fit(epochs=300, update_steps=10)

# Use linear encoder when graph structure is less important
agent = CCVGAE_agent(
    adata=adata,
    layer='counts',
    n_var=2000,
    encoder_type='linear',  # Faster than graph encoder
    latent_dim=10,
    hidden_dim=256,
    hidden_layers=3,
)

agent.fit(epochs=200, update_steps=10)

# Get latent representations (deterministic posterior means)
latent = agent.get_latent()  # Shape: (n_cells, latent_dim)

# Add to AnnData for downstream analysis
adata.obsm['X_ccvgae'] = latent

# Access final scores after training
agent.score_final()
print(f"Final scores: {agent.final_score}")

# Clustering
from sklearn.cluster import KMeans
clusters = KMeans(n_clusters=8).fit_predict(latent)
adata.obs['ccvgae_clusters'] = clusters

# Visualization with UMAP
adata.obsm['X_ccvgae'] = latent
sc.pp.neighbors(adata, use_rep='X_ccvgae')
sc.tl.umap(adata)
sc.pl.umap(adata, color=['ccvgae_clusters', 'cell_type'])

# Trajectory inference
import scanpy as sc
sc.tl.diffmap(adata, use_rep='X_ccvgae')
sc.tl.dpt(adata)