model.py

"""
model.py - Inference-only PoreGCN architecture for HuggingFace Space.

Stripped from PoreGCN_unified/model.py:
  - Training methods removed (no forward_with_porosity, no training utilities).
  - Porosity head kept because checkpoints may include those weights.
  - create_inference_model() is the single entry point used by xai_engine.py.

Architecture summary:
  - BondEmbedding: learnable edge-type embeddings (5 types)
  - AtomConvLayer: CGCNN-style atom-atom message passing
  - AtomPoreConvLayer: heterogeneous atom<->pore message passing (KEY INNOVATION)
  - PoreGCN: stacks the above, pools atoms and pores, predicts MTL targets

Frontend contract: this module is not imported directly by app.py.
xai_engine.py calls create_inference_model() internally.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

# Properties that receive a Softplus non-negativity constraint.
# log10_CO2_N2_selectivity is intentionally absent (log10 values can be negative).
DEFAULT_NONNEG_PROPERTIES = {
    'ASA', 'GSA', 'VF', 'LCD', 'PLD', 'density',
    'CO2_0.01bar_298K', 'CO2_0.1bar_298K', 'CO2_0.15bar_298K',
    'CO2_0.5bar_298K', 'CO2_1bar_298K',
    'CH4_0.9bar_298K', 'CH4_2.5bar_298K', 'CH4_35bar_298K',
    'N2_0.09bar_298K', 'N2_0.9bar_298K',
    'H2_2bar_77K', 'H2_100bar_77K',
    'Xe_1bar_273K', 'Xe_10bar_273K',
    'thermal_stability',
}

POROSITY_GATED_PROPERTIES = {'ASA', 'GSA'}


# =============================================================================
# BondEmbedding
# =============================================================================

class BondEmbedding(nn.Module):
    """Learnable embedding for 5 bond/edge types including ATOM_PORE (type 4)."""

    def __init__(self, num_edge_types: int = 5, embedding_dim: int = 16):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.embedding = nn.Embedding(num_edge_types, embedding_dim, padding_idx=0)

    def forward(self, edge_types: torch.LongTensor) -> torch.Tensor:
        return self.embedding(edge_types)


# =============================================================================
# AtomConvLayer
# =============================================================================

class AtomConvLayer(nn.Module):
    """CGCNN-style graph convolutional layer for atom-atom interactions."""

    def __init__(self, atom_fea_len: int, bond_fea_len: int):
        super().__init__()
        self.atom_fea_len = atom_fea_len
        self.bond_fea_len = bond_fea_len

        self.fc_full = nn.Linear(2 * atom_fea_len + bond_fea_len, 2 * atom_fea_len)
        self.sigmoid = nn.Sigmoid()
        self.softplus1 = nn.Softplus()
        self.bn1 = nn.BatchNorm1d(2 * atom_fea_len)
        self.bn2 = nn.BatchNorm1d(atom_fea_len)
        self.softplus2 = nn.Softplus()

    def forward(
        self,
        atom_fea: torch.Tensor,
        bond_fea: torch.Tensor,
        nbr_fea_idx: torch.LongTensor,
    ) -> torch.Tensor:
        """
        Args:
            atom_fea:    [N_atoms, atom_fea_len]
            bond_fea:    [N_atoms, M_neighbors, bond_fea_len]
            nbr_fea_idx: [N_atoms, M_neighbors]
        Returns:
            Updated atom features [N_atoms, atom_fea_len]
        """
        N, M = nbr_fea_idx.shape
        atom_nbr_fea = atom_fea[nbr_fea_idx, :]  # [N, M, atom_fea_len]

        total_nbr_fea = torch.cat([
            atom_fea.unsqueeze(1).expand(N, M, self.atom_fea_len),
            atom_nbr_fea,
            bond_fea,
        ], dim=2)  # [N, M, 2*atom_fea_len + bond_fea_len]

        total_gated_fea = self.fc_full(total_nbr_fea)
        total_gated_fea = self.bn1(
            total_gated_fea.view(-1, self.atom_fea_len * 2)
        ).view(N, M, self.atom_fea_len * 2)

        nbr_filter, nbr_core = total_gated_fea.chunk(2, dim=2)
        nbr_filter = self.sigmoid(nbr_filter)
        nbr_core = self.softplus1(nbr_core)

        nbr_sumed = torch.sum(nbr_filter * nbr_core, dim=1)
        nbr_sumed = self.bn2(nbr_sumed)

        return self.softplus2(atom_fea + nbr_sumed)


# =============================================================================
# AtomPoreConvLayer
# =============================================================================

class AtomPoreConvLayer(nn.Module):
    """
    Heterogeneous convolution layer for atom-pore interactions.

    Key innovation of PoreGCN: atoms receive messages from adjacent pore nodes,
    and pore nodes receive messages from surrounding atoms. Attention aggregation.
    """

    def __init__(self, atom_fea_len: int, pore_fea_len: int, edge_fea_len: int = 16):
        super().__init__()
        self.atom_fea_len = atom_fea_len
        self.pore_fea_len = pore_fea_len
        self.edge_fea_len = edge_fea_len

        self.atom_from_pore = nn.Sequential(
            nn.Linear(atom_fea_len + pore_fea_len + edge_fea_len, atom_fea_len * 2),
            nn.LayerNorm(atom_fea_len * 2),
            nn.ReLU(),
            nn.Linear(atom_fea_len * 2, atom_fea_len),
        )
        self.pore_from_atom = nn.Sequential(
            nn.Linear(pore_fea_len + atom_fea_len + edge_fea_len, pore_fea_len * 2),
            nn.LayerNorm(pore_fea_len * 2),
            nn.ReLU(),
            nn.Linear(pore_fea_len * 2, pore_fea_len),
        )
        self.attn_atom = nn.Linear(atom_fea_len + pore_fea_len, 1)
        self.attn_pore = nn.Linear(pore_fea_len + atom_fea_len, 1)

    def forward(
        self,
        atom_fea: torch.Tensor,
        pore_fea: torch.Tensor,
        atom_pore_edges: torch.LongTensor,
        edge_fea: torch.Tensor,
        atom_pore_batch: Optional[torch.LongTensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            atom_fea:        [N_atoms, atom_fea_len]
            pore_fea:        [N_pores, pore_fea_len]
            atom_pore_edges: [2, N_edges] where [0]=atom_idx, [1]=pore_idx
            edge_fea:        [N_edges, edge_fea_len]
        Returns:
            (updated_atom_fea, updated_pore_fea)
        """
        if atom_pore_edges.numel() == 0:
            return atom_fea, pore_fea

        atom_idx = atom_pore_edges[0]
        pore_idx = atom_pore_edges[1]

        atom_on_edge = atom_fea[atom_idx]
        pore_on_edge = pore_fea[pore_idx]

        # Atoms receive messages from pores
        atom_pore_cat = torch.cat([atom_on_edge, pore_on_edge, edge_fea], dim=1)
        atom_messages = self.atom_from_pore(atom_pore_cat)
        attn_a = torch.softmax(self.attn_atom(torch.cat([atom_on_edge, pore_on_edge], dim=1)), dim=0)
        atom_messages = atom_messages * attn_a

        atom_update = torch.zeros_like(atom_fea)
        atom_update.scatter_add_(0, atom_idx.unsqueeze(1).expand(-1, self.atom_fea_len), atom_messages)
        atom_fea_new = F.relu(atom_fea + atom_update)

        # Pores receive messages from atoms
        pore_atom_cat = torch.cat([pore_on_edge, atom_on_edge, edge_fea], dim=1)
        pore_messages = self.pore_from_atom(pore_atom_cat)
        attn_p = torch.softmax(self.attn_pore(torch.cat([pore_on_edge, atom_on_edge], dim=1)), dim=0)
        pore_messages = pore_messages * attn_p

        pore_update = torch.zeros_like(pore_fea)
        pore_update.scatter_add_(0, pore_idx.unsqueeze(1).expand(-1, self.pore_fea_len), pore_messages)
        pore_fea_new = F.relu(pore_fea + pore_update)

        return atom_fea_new, pore_fea_new


# =============================================================================
# PoreGCN
# =============================================================================

class PoreGCN(nn.Module):
    """
    Pore-Aware Graph Convolutional Network for MOF property prediction.

    Forward signature matches what xai_engine.py and build_graph.py produce
    (single-graph, not batched via DataLoader).
    """

    def __init__(
        self,
        orig_atom_fea_len: int,
        orig_pore_fea_len: int = 8,
        num_bond_types: int = 5,
        bond_embedding_dim: int = 16,
        atom_fea_len: int = 64,
        pore_fea_len: int = 32,
        n_conv: int = 3,
        h_fea_len: int = 128,
        n_h: int = 1,
        dropout: float = 0.0,
        num_targets: int = 1,
        property_names: Optional[List[str]] = None,
        use_nonneg_output: bool = True,
        use_porosity_head: bool = False,
        porosity_gate_weight: float = 1.0,
    ):
        super().__init__()
        self.atom_fea_len = atom_fea_len
        self.pore_fea_len = pore_fea_len
        self.n_conv = n_conv
        self.num_targets = num_targets
        self.use_nonneg_output = use_nonneg_output
        self.use_porosity_head = use_porosity_head
        self.porosity_gate_weight = porosity_gate_weight

        self.property_names = property_names or [f'prop_{i}' for i in range(num_targets)]
        self.nonneg_mask = [name in DEFAULT_NONNEG_PROPERTIES for name in self.property_names]
        self.porosity_gate_mask = [name in POROSITY_GATED_PROPERTIES for name in self.property_names]

        self.atom_embedding = nn.Linear(orig_atom_fea_len, atom_fea_len)
        self.pore_embedding = nn.Linear(orig_pore_fea_len, pore_fea_len)
        self.bond_embedding = BondEmbedding(num_bond_types, bond_embedding_dim)

        self.atom_convs = nn.ModuleList([
            AtomConvLayer(atom_fea_len, bond_embedding_dim)
            for _ in range(n_conv)
        ])
        self.atom_pore_convs = nn.ModuleList([
            AtomPoreConvLayer(atom_fea_len, pore_fea_len, bond_embedding_dim)
            for _ in range(n_conv)
        ])

        combined_fea_len = atom_fea_len + pore_fea_len
        if n_h > 0:
            layers: list = [nn.Linear(combined_fea_len, h_fea_len), nn.ReLU()]
            if dropout > 0:
                layers.append(nn.Dropout(dropout))
            for _ in range(n_h - 1):
                layers.extend([nn.Linear(h_fea_len, h_fea_len), nn.ReLU()])
                if dropout > 0:
                    layers.append(nn.Dropout(dropout))
            layers.append(nn.Linear(h_fea_len, num_targets))
            self.fc_out = nn.Sequential(*layers)
        else:
            self.fc_out = nn.Linear(combined_fea_len, num_targets)

        self.softplus = nn.Softplus(beta=1.0)

        if use_porosity_head:
            self.porosity_head = nn.Sequential(
                nn.Linear(combined_fea_len, h_fea_len),
                nn.ReLU(),
                nn.Dropout(dropout) if dropout > 0 else nn.Identity(),
                nn.Linear(h_fea_len, 1),
                nn.Sigmoid(),
            )
        else:
            self.porosity_head = None

    def forward(
        self,
        atom_fea: torch.Tensor,
        bond_types: torch.LongTensor,
        nbr_fea_idx: torch.LongTensor,
        crystal_atom_idx: List[torch.LongTensor],
        pore_fea: torch.Tensor,
        atom_pore_edges: torch.LongTensor,
        crystal_pore_idx: List[torch.LongTensor],
    ) -> torch.Tensor:
        """
        Args:
            atom_fea:        [N_atoms, orig_atom_fea_len]
            bond_types:      [N_atoms, M_neighbors]
            nbr_fea_idx:     [N_atoms, M_neighbors]
            crystal_atom_idx: list of LongTensors, one per crystal in batch
            pore_fea:        [N_pores, orig_pore_fea_len] (may be empty)
            atom_pore_edges: [2, N_ap_edges]
            crystal_pore_idx: list of LongTensors, one per crystal
        Returns:
            Predictions [batch_size, num_targets]
        """
        atom_fea = self.atom_embedding(atom_fea)
        pore_fea = self.pore_embedding(pore_fea)
        bond_fea = self.bond_embedding(bond_types)

        n_ap_edges = atom_pore_edges.shape[1] if atom_pore_edges.numel() > 0 else 0
        if n_ap_edges > 0:
            ap_types = torch.full((n_ap_edges,), 4, dtype=torch.long, device=atom_fea.device)
            ap_fea = self.bond_embedding(ap_types)
        else:
            ap_fea = torch.zeros((0, self.bond_embedding.embedding_dim), device=atom_fea.device)

        for i in range(self.n_conv):
            atom_fea = self.atom_convs[i](atom_fea, bond_fea, nbr_fea_idx)
            if pore_fea.shape[0] > 0:
                atom_fea, pore_fea = self.atom_pore_convs[i](
                    atom_fea, pore_fea, atom_pore_edges, ap_fea
                )

        atom_pool = self._pool_crystals(atom_fea, crystal_atom_idx)

        if pore_fea.shape[0] > 0:
            pore_pool = self._pool_crystals(pore_fea, crystal_pore_idx)
        else:
            batch_size = len(crystal_atom_idx)
            pore_pool = torch.zeros((batch_size, self.pore_fea_len), device=atom_fea.device)

        combined = torch.cat([atom_pool, pore_pool], dim=1)
        out = self.fc_out(combined)

        if self.use_nonneg_output and any(self.nonneg_mask):
            cols = []
            for i, is_nonneg in enumerate(self.nonneg_mask):
                cols.append(self.softplus(out[:, i:i+1]) if is_nonneg else out[:, i:i+1])
            out = torch.cat(cols, dim=1)

        if self.porosity_head is not None:
            porosity_prob = self.porosity_head(combined)
            gate = 1.0 - self.porosity_gate_weight + self.porosity_gate_weight * porosity_prob.squeeze(-1)
            cols = []
            for i, is_gated in enumerate(self.porosity_gate_mask):
                cols.append((out[:, i] * gate).unsqueeze(1) if is_gated else out[:, i:i+1])
            out = torch.cat(cols, dim=1)

        return out

    def _pool_crystals(
        self,
        fea: torch.Tensor,
        crystal_idx: List[torch.LongTensor],
    ) -> torch.Tensor:
        batch_size = len(crystal_idx)
        fea_dim = fea.shape[1]
        pooled = torch.zeros((batch_size, fea_dim), device=fea.device)
        for i, idx in enumerate(crystal_idx):
            if len(idx) > 0:
                pooled[i] = fea[idx].mean(dim=0)
        return pooled


# =============================================================================
# Public entry point
# =============================================================================

def _infer_arch(state_dict: Dict) -> Dict:
    """
    Infer all PoreGCN hyperparameters from a saved state_dict.

    This avoids needing to pass a config object at inference time, which is
    important because the checkpoint was saved with a training config that
    may not be importable from the HF Space environment.
    """
    orig_atom_fea_len = state_dict['atom_embedding.weight'].shape[1]
    orig_pore_fea_len = state_dict['pore_embedding.weight'].shape[1]
    atom_fea_len = state_dict['atom_embedding.weight'].shape[0]
    pore_fea_len = state_dict['pore_embedding.weight'].shape[0]

    n_conv = sum(
        1 for k in state_dict
        if k.startswith('atom_convs.') and k.endswith('.fc_full.weight')
    )

    fc_indices = sorted(
        int(k.split('.')[1])
        for k in state_dict
        if k.startswith('fc_out.') and k.endswith('.weight')
    )
    if not fc_indices:
        raise ValueError('Cannot infer architecture: fc_out weights missing from state_dict')

    h_fea_len = state_dict['fc_out.0.weight'].shape[0]
    last_idx = fc_indices[-1]
    num_targets = state_dict[f'fc_out.{last_idx}.weight'].shape[0]

    n_h = len(fc_indices) - 1
    gap = fc_indices[1] - fc_indices[0] if len(fc_indices) > 1 else 1
    dropout = 0.1 if gap >= 3 else 0.0

    has_porosity_head = any('porosity_head' in k for k in state_dict)

    # Bond embedding dim
    bond_embedding_dim = state_dict['bond_embedding.embedding.weight'].shape[1]

    return {
        'orig_atom_fea_len': orig_atom_fea_len,
        'orig_pore_fea_len': orig_pore_fea_len,
        'atom_fea_len': atom_fea_len,
        'pore_fea_len': pore_fea_len,
        'n_conv': n_conv,
        'h_fea_len': h_fea_len,
        'n_h': n_h,
        'dropout': dropout,  # Must match training value to reproduce fc_out Sequential structure
        'num_targets': num_targets,
        'bond_embedding_dim': bond_embedding_dim,
        'has_porosity_head': has_porosity_head,
    }


def create_inference_model(checkpoint_dict: Dict, device: str = 'cpu') -> PoreGCN:
    """
    Instantiate an eval-mode PoreGCN from a loaded checkpoint dict.

    Accepts two checkpoint formats produced by PoreGCN_unified training:
      - CV checkpoint format: keys 'model_state', 'model_config'
      - Final model format:   keys 'model_state_dict', 'config'

    Args:
        checkpoint_dict: dict loaded by torch.load()
        device: target device string (always 'cpu' on HF free tier)

    Returns:
        PoreGCN in eval mode, weights loaded, moved to device
    """
    # Resolve state_dict key (two checkpoint formats exist)
    if 'model_state' in checkpoint_dict:
        state_dict = checkpoint_dict['model_state']
        model_cfg = checkpoint_dict.get('model_config', {})
    elif 'model_state_dict' in checkpoint_dict:
        state_dict = checkpoint_dict['model_state_dict']
        model_cfg = checkpoint_dict.get('config', {})
    else:
        raise KeyError(
            'Checkpoint must have key "model_state" or "model_state_dict". '
            f'Got: {list(checkpoint_dict.keys())}'
        )

    arch = _infer_arch(state_dict)
    property_names = model_cfg.get('property_names', None)
    use_nonneg = model_cfg.get('use_nonneg_output', True)

    model = PoreGCN(
        orig_atom_fea_len=arch['orig_atom_fea_len'],
        orig_pore_fea_len=arch['orig_pore_fea_len'],
        num_bond_types=5,
        bond_embedding_dim=arch['bond_embedding_dim'],
        atom_fea_len=arch['atom_fea_len'],
        pore_fea_len=arch['pore_fea_len'],
        n_conv=arch['n_conv'],
        h_fea_len=arch['h_fea_len'],
        n_h=arch['n_h'],
        dropout=arch['dropout'],
        num_targets=arch['num_targets'],
        property_names=property_names,
        use_nonneg_output=use_nonneg,
        use_porosity_head=arch['has_porosity_head'],
    )

    model.load_state_dict(state_dict, strict=True)
    model.to(device)
    model.eval()
    return model