main.py

"""
Main entry point / reference example for BINN's two-persona public API.

``facade_main``   — bio-informed persona: the ``binn.BINN(...)`` one-liner path.
``main``          — computational persona: trains ANN/GCN/GAT independently and compares.
``ensemble_main`` — computational persona: trains all three via a staged VCDN-fused ensemble.
``e2e_main``      — computational persona: trains all three via an end-to-end VCDN-fused ensemble.

All four run on the same prostate-cancer dataset (``train_data/``), so the
bio-informed and computational paths are directly comparable.
"""
import os
import matplotlib.pyplot as plt

from binn import BINN
from binn.data import build_reactome_network
from binn.learn import (
    Hyperparameters,
    plot_confusion_matrix,
    plot_pca_latent,
    plot_last_layer_weights,
)
from data_utils import load_and_merge_tables, prepare_graph_data
from binn.train import train_graph_model, train_ensemble_model, train_ensemble_model_e2e

CLIN_VARS = [
    'Purity',
    'Ploidy',
    'Tumor.Coverage',
    'Normal.Coverage',
    'Mutation.burden',
    'Fraction.genome.altered',
    'Mutation_count',
]


def _load_data_and_config(data_dir: str = "train_data"):
    """Load data, build graph, and create default Hyperparameters (computational persona)."""
    obs_vars = CLIN_VARS + ['response']

    print("Loading and preparing data...")
    merged = load_and_merge_tables(data_dir)
    reactome_net = build_reactome_network()
    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 facade_main(data_dir: str = "train_data") -> None:
    """Bio-informed persona: ``binn.BINN(...)``, no manual architecture assembly.

    Same prostate-cancer dataset as ``main()``/``ensemble_main()``, but the
    caller never touches ``Hyperparameters``, ``GraphModelFactory``, or the
    Reactome network directly — just data + a pathway source name.
    """
    print("Loading data...")
    merged = load_and_merge_tables(data_dir)

    print("\n" + "=" * 60)
    print("binn.BINN — bio-informed persona (GCN, Reactome)")
    print("=" * 60)
    model = BINN(merged, pathway="reactome", label_col="response", architecture="GCN")
    model.fit(epochs=200, patience=50)
    print(model.summary())

    print("\nScoring pathway importance...")
    report = model.interpret(top_k=10, max_nodes_per_level=25)
    print(f"  {len(report['pathway_importance'])} layer transitions scored.")
    if "figure" in report:
        out_path = "binn_facade_sankey.html"
        report["figure"].write_html(out_path)
        print(f"  Wrote pathway Sankey diagram to {out_path}")

    print("\nDone!")


def main() -> None:
    """Computational persona: 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}%)")
    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}%)")
    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", "facade"],
        help="def: independent models, ens: staged ensemble, e2e: end-to-end ensemble, "
             "facade: bio-informed binn.BINN(...) one-liner",
    )
    args = parser.parse_args()

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