src/explain.py

"""
Model Explainability Suite using GNNExplainer and RDKit 2D Substructure Rendering.
Supports both single-task and multi-task models, with edge-aware and non-edge GNNs.
"""
import os
import sys

# Ensure project root is in sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import logging
import argparse
from typing import Dict, List, Optional, Tuple, Any, Union

import numpy as np
import torch
import torch.nn as nn
from torch_geometric.explain import Explainer, GNNExplainer, ModelConfig
from rdkit import Chem
from rdkit.Chem.Draw import rdMolDraw2D
import matplotlib.pyplot as plt

from src.data import smiles_to_graph
from src.models.gnn import build_gnn_model

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)


class TaskSpecificWrapper(nn.Module):
    """Wraps a multi-task GNN to isolate a single task output for GNNExplainer."""
    def __init__(self, base_model: nn.Module, task_idx: int = 0):
        super().__init__()
        self.base_model = base_model
        self.task_idx = task_idx

    def forward(
        self,
        x: torch.Tensor,
        edge_index: torch.Tensor,
        edge_attr: Optional[torch.Tensor] = None,
        batch: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        out = self.base_model(x, edge_index, edge_attr=edge_attr, batch=batch)
        if out.ndim == 1:
            return out.view(-1, 1)
        return out[:, self.task_idx:self.task_idx + 1]


def explain_molecule(
    model: nn.Module,
    smiles: str,
    task_type: str = 'regression',
    task_idx: int = 0,
    num_tasks: int = 1,
    num_epochs: int = 50,
    device: str = 'cpu'
) -> Tuple[np.ndarray, np.ndarray, float]:
    """
    Explain a molecular graph prediction using GNNExplainer.
    Returns:
        atom_importances: np.ndarray of shape (num_atoms,)
        edge_importances: np.ndarray of shape (num_edges,)
        pred_value: float scalar prediction
    """
    model.eval()
    graph = smiles_to_graph(smiles)
    if graph is None:
        raise ValueError(f"Invalid SMILES string: '{smiles}'")

    graph = graph.to(device)

    # Determine prediction
    with torch.no_grad():
        raw_pred = model(graph.x, graph.edge_index, edge_attr=graph.edge_attr)
        if task_type != 'regression':
            raw_pred = torch.sigmoid(raw_pred)
        pred_scalar = float(raw_pred[0, task_idx].item()) if raw_pred.ndim > 1 else float(raw_pred[0].item())

    # Wrap model if multi-task
    if num_tasks > 1:
        wrapped_model = TaskSpecificWrapper(model, task_idx=task_idx)
    else:
        wrapped_model = model

    mode = 'regression' if task_type == 'regression' else 'binary_classification'
    explainer = Explainer(
        model=wrapped_model,
        algorithm=GNNExplainer(epochs=num_epochs),
        explanation_type='model',
        node_mask_type='attributes',
        edge_mask_type='object' if graph.edge_index.size(1) > 0 else None,
        model_config=ModelConfig(
            mode=mode,
            task_level='graph',
            return_type='raw'
        )
    )

    explanation = explainer(
        x=graph.x,
        edge_index=graph.edge_index,
        edge_attr=graph.edge_attr
    )

    # Atom importance: sum across feature attributes
    if explanation.node_mask is not None:
        atom_scores = explanation.node_mask.sum(dim=-1).detach().cpu().numpy()
    else:
        atom_scores = np.ones(graph.num_nodes, dtype=np.float32)

    # Edge importance
    if explanation.edge_mask is not None and graph.edge_index.size(1) > 0:
        edge_scores = explanation.edge_mask.detach().cpu().numpy()
    else:
        edge_scores = np.zeros(graph.edge_index.size(1), dtype=np.float32)

    return atom_scores, edge_scores, pred_scalar


def render_molecule_explanation(
    smiles: str,
    atom_importances: np.ndarray,
    width: int = 500,
    height: int = 400,
    colormap_name: str = 'Reds',
    output_path: Optional[str] = None
) -> bytes:
    """
    Render 2D molecular structure with atom-level importance heatmaps.
    Outputs high-resolution PNG bytes and optionally saves to disk.
    """
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Invalid SMILES string '{smiles}'")

    num_atoms = mol.GetNumAtoms()
    if len(atom_importances) != num_atoms:
        atom_scores = np.resize(atom_importances, num_atoms)
    else:
        atom_scores = np.array(atom_importances, dtype=np.float32)

    # Normalize scores between 0.15 and 0.95 for clear visual contrast
    min_s = float(atom_scores.min())
    max_s = float(atom_scores.max())
    range_s = max_s - min_s if max_s > min_s else 1.0
    norm_scores = 0.15 + 0.8 * ((atom_scores - min_s) / range_s)

    cmap = plt.colormaps[colormap_name]
    atom_colors = {}
    radii = {}

    for i in range(num_atoms):
        rgba = cmap(float(norm_scores[i]))
        atom_colors[i] = (float(rgba[0]), float(rgba[1]), float(rgba[2]))
        radii[i] = 0.35 + 0.25 * float(norm_scores[i])

    drawer = rdMolDraw2D.MolDraw2DCairo(width, height)
    draw_options = drawer.drawOptions()
    draw_options.useBWAtomPalette()
    draw_options.scaleBondWidth = True
    draw_options.bondLineWidth = 2

    drawer.DrawMolecule(
        mol,
        highlightAtoms=list(atom_colors.keys()),
        highlightAtomColors=atom_colors,
        highlightAtomRadii=radii
    )
    drawer.FinishDrawing()
    png_bytes = drawer.GetDrawingText()

    if output_path is not None:
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'wb') as f:
            f.write(png_bytes)
        logger.info(f"Saved explanation rendering -> {output_path}")

    return png_bytes


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Explain GNN molecular property prediction")
    parser.add_argument("--smiles", type=str, default="CC(=O)OC1=CC=CC=C1C(=O)O", help="SMILES to explain")
    parser.add_argument("--model", type=str, default="gcn", help="Model name (gcn, gat, gin, mpnn)")
    parser.add_argument("--dataset", type=str, default="esol", help="Dataset name (esol, bbbp, tox21)")
    parser.add_argument("--task_idx", type=int, default=0, help="Task index to explain (0-11 for Tox21)")
    parser.add_argument("--output", type=str, default="results/plots/explanation_demo.png", help="Output PNG path")
    args = parser.parse_args()

    from src.evaluate import load_trained_model
    model, payload = load_trained_model(args.model, args.dataset)
    num_tasks = payload['num_tasks'] if payload else 1
    task_type = 'regression' if args.dataset == 'esol' else ('multitask_classification' if args.dataset == 'tox21' else 'classification')

    atom_imp, edge_imp, pred = explain_molecule(
        model=model,
        smiles=args.smiles,
        task_type=task_type,
        task_idx=args.task_idx,
        num_tasks=num_tasks,
        num_epochs=30
    )

    print(f"Prediction for {args.smiles}: {pred:.4f}")
    print(f"Top 3 important atom indices: {np.argsort(atom_imp)[-3:][::-1]}")
    render_molecule_explanation(args.smiles, atom_imp, output_path=args.output)