main.py

"""
Main entry point for training ANN, GCN, and GAT models using scvi-tools BaseModelClass.

``main``          — trains each model independently and compares performance.
``ensemble_main`` — trains all three via VCDN-fused ensemble.
"""
import os
import matplotlib.pyplot as plt
from binn.learn import (
    Hyperparameters,
    plot_confusion_matrix,
    plot_pca_latent,
    plot_last_layer_weights,
)
from data_utils import load_and_merge_tables, build_reactome_network, prepare_graph_data
from binn.train import train_graph_model, train_ensemble_model, train_ensemble_model_e2e


def _load_data_and_config(data_dir: str = "train_data"):
    """Load data, build graph, and create default Hyperparameters."""
    reactome_dir = os.path.join(data_dir, "reactome")
    clin_vars = [
        'Purity',
        'Ploidy',
        'Tumor.Coverage',
        'Normal.Coverage',
        'Mutation.burden',
        'Fraction.genome.altered',
        'Mutation_count',
    ]
    obs_vars = clin_vars + ['response']

    print("Loading and preparing data...")
    merged = load_and_merge_tables(data_dir)
    reactome_net = build_reactome_network(reactome_dir)
    data, adata, map_df = prepare_graph_data(merged, obs_vars, reactome_net, data_dir)

    print(f"Graph data: {data}")
    print(f"Number of nodes: {data.x.shape[0]}")
    print(f"Number of features: {data.x.shape[1]}")
    print(f"Number of edges: {data.edge_index.shape[1]}")
    print(f"Number of classes: {len(data.y.unique())}")
    print(f"Pathway mapping shape: {map_df.shape}")

    config = Hyperparameters(
        num_node_features=data.num_node_features,
        num_classes=int(len(data.y.unique())),
        lr=1e-3,
        w_decay=1e-5,
        epochs=400,
        patience=200,
        bias=False,
        heads=1,
        save_dir="weights",
    )
    print(f"Using device: {config.device}")


    return data, adata, map_df, config


def evaluate_model(model, graph_data, mask):
    """Evaluate model accuracy on a given mask."""
    predictions = model.predict(mask=mask)
    pred_classes = predictions.argmax(dim=1)
    true_labels = graph_data.y[mask]
    correct = (pred_classes == true_labels).sum().item()
    total = mask.sum().item()
    return correct / total


def main() -> None:
    """Train and compare ANN, GCN, and GAT models independently."""
    data, adata, map_df, config = _load_data_and_config()

    results = {}
    for name in ["ANN", "GCN", "GAT"]:
        print("\n" + "="*60)
        print(f"Training {name} model...")
        print("="*60)
        results[name] = train_graph_model(name, data, map_df, config, adata=adata)
    
    # Compare performances on test set
    print("\n" + "="*60)
    print("Model Performance Comparison (Test Set)")
    print("="*60)
    
    if hasattr(data, 'test_mask'):
        for model_name, model in results.items():
            accuracy = evaluate_model(model, data, data.test_mask)
            print(f"{model_name}: Test Accuracy = {accuracy:.4f} ({accuracy*100:.2f}%)")
        
        # Find best model
        accuracies = {name: evaluate_model(model, data, data.test_mask) 
                      for name, model in results.items()}
        best_model = max(accuracies, key=accuracies.get)
        print(f"\nBest performing model: {best_model} with {accuracies[best_model]*100:.2f}% accuracy")
    else:
        print("Warning: No test_mask found in data. Cannot evaluate on test set.")
    
  
    print("\n" + "="*60)
    print("Interpretability visualizations")
    print("="*60)
    if hasattr(data, 'test_mask'):
        test_mask = data.test_mask
        for model_name, model in results.items():
            print(f"\n[{model_name}] Confusion matrix (test)")
            plot_confusion_matrix(model, data, test_mask, title=f"{model_name} - Confusion matrix (test)")
            plt.show()

    for model_name, model in results.items():
        print(f"[{model_name}] PCA of latent representation")
        plot_pca_latent(model, data, title=f"{model_name} - PCA (latent)")
        plt.show()

    for model_name, model in results.items():
        print(f"[{model_name}] Final layer weights")
        plot_last_layer_weights(model, title=model_name)
        plt.show()

    print("\nDone!")


def ensemble_main() -> None:
    """Train an ANN+GCN+GAT ensemble fused via VCDN (staged)."""
    data, adata, map_df, config = _load_data_and_config()

    

    arch_names = {"ANN": "ANN", "GCN": "GCN", "GAT": "GAT"}
    graph_datas = {name: data for name in arch_names}
    map_dfs = {name: map_df for name in arch_names}

    print("\n" + "=" * 60)
    print("Training ensemble (ANN + GCN + GAT → VCDN)")
    print("=" * 60)

    ensemble = train_ensemble_model(
        names=arch_names,
        graph_datas=graph_datas,
        map_dfs=map_dfs,
        config=config,
        adata=adata,
    )

    print("\n" + "=" * 60)
    print("Ensemble Performance (Test Set)")
    print("=" * 60)

    if hasattr(data, 'test_mask'):
        preds = ensemble.predict(mask=data.test_mask)
        pred_classes = preds.argmax(dim=1)
        true_labels = data.y[data.test_mask]
        accuracy = (pred_classes == true_labels).sum().item() / data.test_mask.sum().item()
        print(f"Ensemble: Test Accuracy = {accuracy:.4f} ({accuracy * 100:.2f}%)")

        for name, model in ensemble.sub_models.items():
            sub_acc = evaluate_model(model, data, data.test_mask)
            print(f"  {name} (sub-model): Test Accuracy = {sub_acc:.4f} ({sub_acc * 100:.2f}%)")
    else:
        print("Warning: No test_mask found in data. Cannot evaluate on test set.")

    print("\nDone!")


def e2e_main() -> None:
    """Train an ANN+GCN+GAT ensemble end-to-end via VCDN."""
    data, adata, map_df, config = _load_data_and_config()

    arch_names = {"ANN": "ANN", "GCN": "GCN", "GAT": "GAT"}
    graph_datas = {name: data for name in arch_names}
    map_dfs = {name: map_df for name in arch_names}

    print("\n" + "=" * 60)
    print("Training end-to-end ensemble (ANN + GCN + GAT → VCDN)")
    print("=" * 60)

    ensemble = train_ensemble_model_e2e(
        names=arch_names,
        graph_datas=graph_datas,
        map_dfs=map_dfs,
        config=config,
        adata=adata,
    )

    print("\n" + "=" * 60)
    print("E2E Ensemble Performance (Test Set)")
    print("=" * 60)

    if hasattr(data, 'test_mask'):
        preds = ensemble.predict(mask=data.test_mask)
        pred_classes = preds.argmax(dim=1)
        true_labels = data.y[data.test_mask]
        accuracy = (pred_classes == true_labels).sum().item() / data.test_mask.sum().item()
        print(f"E2E Ensemble: Test Accuracy = {accuracy:.4f} ({accuracy * 100:.2f}%)")

        for name, model in ensemble.sub_models.items():
            sub_acc = evaluate_model(model, data, data.test_mask)
            print(f"  {name} (sub-model): Test Accuracy = {sub_acc:.4f} ({sub_acc * 100:.2f}%)")
    else:
        print("Warning: No test_mask found in data. Cannot evaluate on test set.")
    print("\nDone!")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Train graph models")
    parser.add_argument(
        "mode",
        nargs="?",
        default="def",
        choices=["def", "ens", "e2e"],
        help="def: independent models, ens: staged ensemble, e2e: end-to-end ensemble",
    )
    args = parser.parse_args()

    {"def": main, "ens": ensemble_main, "e2e": e2e_main}[args.mode]()