inference_app.py

from __future__ import annotations
import time
import json 
import gradio as gr
from gradio_molecule3d import Molecule3D
import torch
from pinder.core import get_pinder_location
get_pinder_location()
from pytorch_lightning import LightningModule

import torch
import lightning.pytorch as pl
import torch.nn.functional as F

import torch.nn as nn
import torchmetrics
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.nn import global_mean_pool
from torch.nn import Sequential, Linear, BatchNorm1d, ReLU
from torch_scatter import scatter
from torch.nn import Module


import pinder.core as pinder
pinder.__version__
from torch_geometric.loader import DataLoader
from pinder.core.loader.dataset import get_geo_loader
from pinder.core import download_dataset
from pinder.core import get_index
from pinder.core import get_metadata
from pathlib import Path
import pandas as pd
from pinder.core import PinderSystem
import torch
from pinder.core.loader.dataset import PPIDataset
from pinder.core.loader.geodata import NodeRepresentation
import pickle
from pinder.core import get_index, PinderSystem
from torch_geometric.data import HeteroData
import os

from enum import Enum

import numpy as np
import torch
import lightning.pytorch as pl
from numpy.typing import NDArray
from torch_geometric.data import HeteroData

from pinder.core.index.system import PinderSystem
from pinder.core.loader.structure import Structure
from pinder.core.utils import constants as pc
from pinder.core.utils.log import setup_logger
from pinder.core.index.system import _align_monomers_with_mask
from pinder.core.loader.structure import Structure

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.nn import global_mean_pool
from torch.nn import Sequential, Linear, BatchNorm1d, ReLU
from torch_scatter import scatter
from torch.nn import Module
import time
from torch_geometric.nn import global_max_pool
import copy
import inspect
import warnings
from typing import Optional, Tuple, Union

import torch
from torch import Tensor

from torch_geometric.data import Data, Dataset, HeteroData
from torch_geometric.data.feature_store import FeatureStore
from torch_geometric.data.graph_store import GraphStore
from torch_geometric.loader import (
    LinkLoader,
    LinkNeighborLoader,
    NeighborLoader,
    NodeLoader,
)
from torch_geometric.loader.dataloader import DataLoader
from torch_geometric.loader.utils import get_edge_label_index, get_input_nodes
from torch_geometric.sampler import BaseSampler, NeighborSampler
from torch_geometric.typing import InputEdges, InputNodes

try:
    from lightning.pytorch import LightningDataModule as PLLightningDataModule
    no_pytorch_lightning = False
except (ImportError, ModuleNotFoundError):
    PLLightningDataModule = object
    no_pytorch_lightning = True

from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.loggers.tensorboard import TensorBoardLogger
from lightning.pytorch.callbacks.early_stopping import EarlyStopping
from torch_geometric.data.lightning.datamodule import LightningDataset
from pytorch_lightning.loggers.wandb import WandbLogger
def get_system(system_id: str) -> PinderSystem:
    return PinderSystem(system_id)
from Bio import PDB
from Bio.PDB.PDBIO import PDBIO
from pinder.core.structure.atoms import atom_array_from_pdb_file
from pathlib import Path
from pinder.eval.dockq.biotite_dockq import BiotiteDockQ

def extract_coordinates_from_pdb(filename, atom_name="CA"):
    """
    Extracts coordinates for specific atoms from a PDB file and returns them as a list of tuples.
    Each tuple contains (x, y, z) coordinates of the specified atom type.
    
    Parameters:
        filename (str): Path to the PDB file.
        atom_name (str): The name of the atom to filter by (e.g., "CA" for alpha carbon).
        
    Returns:
        list of tuple: List of coordinates as (x, y, z) tuples for the specified atom.
    """
    parser = PDB.PDBParser(QUIET=True)
    structure = parser.get_structure("structure", filename)

    coordinates = []

    # Loop through each model, chain, residue, and atom to collect coordinates of specified atom
    for model in structure:
        for chain in model:
            for residue in chain:
                for atom in residue:
                      # Filter for specific atom name
                        xyz = atom.coord  # Coordinates are in a numpy array
                        coordinates.append([xyz[0], xyz[1], xyz[2]])

    return coordinates

log = setup_logger(__name__)

try:
    from torch_cluster import knn_graph

    torch_cluster_installed = True
except ImportError as e:
    log.warning(
        "torch-cluster is not installed!"
        "Please install the appropriate library for your pytorch installation."
        "See https://github.com/rusty1s/pytorch_cluster/issues/185 for background."
    )
    torch_cluster_installed = False


def structure2tensor(
    atom_coordinates: NDArray[np.double] | None = None,
    atom_types: NDArray[np.str_] | None = None,
    element_types: NDArray[np.str_] | None = None,
    residue_coordinates: NDArray[np.double] | None = None,
    residue_ids: NDArray[np.int_] | None = None,
    residue_types: NDArray[np.str_] | None = None,
    chain_ids: NDArray[np.str_] | None = None,
    dtype: torch.dtype = torch.float32,
) -> dict[str, torch.Tensor]:
    property_dict = {}
    if atom_types is not None:
        unknown_name_idx = max(pc.ALL_ATOM_POSNS.values()) + 1
        types_array_at = np.zeros((len(atom_types), 1))
        for i, name in enumerate(atom_types):
            types_array_at[i] = pc.ALL_ATOM_POSNS.get(name, unknown_name_idx)
        property_dict["atom_types"] = torch.tensor(types_array_at).type(dtype)
    if element_types is not None:
        types_array_ele = np.zeros((len(element_types), 1))
        for i, name in enumerate(element_types):
            types_array_ele[i] = pc.ELE2NUM.get(name, pc.ELE2NUM["other"])
        property_dict["element_types"] = torch.tensor(types_array_ele).type(dtype)
    if residue_types is not None:
        unknown_name_idx = max(pc.AA_TO_INDEX.values()) + 1
        types_array_res = np.zeros((len(residue_types), 1))
        for i, name in enumerate(residue_types):
            types_array_res[i] = pc.AA_TO_INDEX.get(name, unknown_name_idx)
        property_dict["residue_types"] = torch.tensor(types_array_res).type(dtype)

    if atom_coordinates is not None:
        property_dict["atom_coordinates"] = torch.tensor(atom_coordinates, dtype=dtype)
        
    if residue_coordinates is not None:
        property_dict["residue_coordinates"] = torch.tensor(
            residue_coordinates, dtype=dtype
        )
    if residue_ids is not None:
        property_dict["residue_ids"] = torch.tensor(residue_ids, dtype=dtype)
    if chain_ids is not None:
        property_dict["chain_ids"] = torch.zeros(len(chain_ids), dtype=dtype)
        property_dict["chain_ids"][chain_ids == "L"] = 1
    return property_dict


class NodeRepresentation(Enum):
    Surface = "surface"
    Atom = "atom"
    Residue = "residue"


class PairedPDB(HeteroData):  # type: ignore
    @classmethod
    def from_tuple_system(
        cls,
    
        tupal: tuple = (Structure , Structure , Structure),
        
        add_edges: bool = True,
        k: int = 10,
        
    ) -> PairedPDB:
        return cls.from_structure_pair(
            
            holo=tupal[0],
            apo=tupal[1],
            add_edges=add_edges,
            k=k,
        )

    @classmethod
    def from_structure_pair(
        cls,

        holo: Structure,
        apo: Structure,
        
        add_edges: bool = True,
        k: int = 10,
    ) -> PairedPDB:
        graph = cls()
        holo_calpha = holo.filter("atom_name", mask=["CA"])
        apo_calpha = apo.filter("atom_name", mask=["CA"])
        r_h = (holo.dataframe['chain_id'] == 'R').sum()
        r_a = (apo.dataframe['chain_id'] == 'R').sum()
        
        holo_r_props = structure2tensor(
            atom_coordinates=holo.coords[:r_h],
            atom_types=holo.atom_array.atom_name[:r_h],
            element_types=holo.atom_array.element[:r_h],
            residue_coordinates=holo_calpha.coords[:r_h],
            residue_types=holo_calpha.atom_array.res_name[:r_h],
            residue_ids=holo_calpha.atom_array.res_id[:r_h],
        )
        holo_l_props = structure2tensor(
            atom_coordinates=holo.coords[r_h:],
            
            atom_types=holo.atom_array.atom_name[r_h:],
            element_types=holo.atom_array.element[r_h:],
            residue_coordinates=holo_calpha.coords[r_h:],
            residue_types=holo_calpha.atom_array.res_name[r_h:],
            residue_ids=holo_calpha.atom_array.res_id[r_h:],
        )
        apo_r_props = structure2tensor(
            atom_coordinates=apo.coords[:r_a],
            atom_types=apo.atom_array.atom_name[:r_a],
            element_types=apo.atom_array.element[:r_a],
            residue_coordinates=apo_calpha.coords[:r_a],
            residue_types=apo_calpha.atom_array.res_name[:r_a],
            residue_ids=apo_calpha.atom_array.res_id[:r_a],
        )
        apo_l_props = structure2tensor(
            atom_coordinates=apo.coords[r_a:],
            atom_types=apo.atom_array.atom_name[r_a:],
            element_types=apo.atom_array.element[r_a:],
            residue_coordinates=apo_calpha.coords[r_a:],
            residue_types=apo_calpha.atom_array.res_name[r_a:],
            residue_ids=apo_calpha.atom_array.res_id[r_a:],
        )
       
        
       
        graph["ligand"].x = apo_l_props["atom_types"]
        graph["ligand"].pos = apo_l_props["atom_coordinates"]
        graph["receptor"].x = apo_r_props["atom_types"]
        graph["receptor"].pos = apo_r_props["atom_coordinates"]
        graph["ligand"].y = holo_l_props["atom_coordinates"]
        # graph["ligand"].pos = holo_l_props["atom_coordinates"]
        graph["receptor"].y = holo_r_props["atom_coordinates"]
        # graph["receptor"].pos = holo_r_props["atom_coordinates"]
        if add_edges and torch_cluster_installed:
                graph["ligand"].edge_index = knn_graph(
                    graph["ligand"].pos, k=k
                )
                graph["receptor"].edge_index = knn_graph(
                    graph["receptor"].pos, k=k
                )
                # graph["ligand"].edge_index = knn_graph(
                #     graph["ligand"].pos, k=k
                # )
                # graph["receptor"].edge_index = knn_graph(
                #     graph["receptor"].pos, k=k
                # )

        return graph
      
#create_graph takes inputs apo_ligand, apo_residue and paired holo as pdb3(ground truth). 
def create_graph(pdb1, pdb2, k=5):
    r"""
    Create a heterogeneous graph from two PDB files, with the ligand and receptor
    as separate nodes, and their respective features and edges.
    
    Args:
        pdb1 (str): PDB file path for ligand.
        pdb2 (str): PDB file path for receptor.
        coords3 (list): List of coordinates used for `y` values (e.g., binding affinity, etc.).
        k (int): Number of nearest neighbors for constructing the knn graph.
    
    Returns:
        HeteroData: A PyG HeteroData object containing ligand and receptor data.
    """
    # Extract coordinates from PDB files
    coords1 = torch.tensor(extract_coordinates_from_pdb(pdb1),dtype=torch.float)  
    coords2 = torch.tensor(extract_coordinates_from_pdb(pdb2),dtype=torch.float)  
    # coords3 = torch.tensor(extract_coordinates_from_pdb(pdb3),dtype=torch.float)
    # Create the HeteroData object
    data = HeteroData()

    # Define ligand node features
    data["ligand"].x = torch.tensor(coords1, dtype=torch.float)
    data["ligand"].pos = coords1
    # data["ligand"].y = torch.tensor(coords3[:len(coords1)], dtype=torch.float)

    # Define receptor node features
    data["receptor"].x = torch.tensor(coords2, dtype=torch.float)
    data["receptor"].pos = coords2
    # data["receptor"].y = torch.tensor(coords3[len(coords1):], dtype=torch.float)

    # Construct k-NN graph for ligand
    ligand_edge_index = knn_graph(data["ligand"].pos, k=k)
    data["ligand"].edge_index = ligand_edge_index

    # Construct k-NN graph for receptor
    receptor_edge_index = knn_graph(data["receptor"].pos, k=k)
    data["receptor"].edge_index = receptor_edge_index

    # Convert edge index to SparseTensor for ligand
    data["ligand", "ligand"].edge_index = ligand_edge_index

    # Convert edge index to SparseTensor for receptor
    data["receptor", "receptor"].edge_index = receptor_edge_index

    return data


def update_pdb_coordinates_from_tensor(input_filename, output_filename, coordinates_tensor):
    r"""
    Updates atom coordinates in a PDB file with new transformed coordinates provided in a tensor.

    Parameters:
    - input_filename (str): Path to the original PDB file.
    - output_filename (str): Path to the new PDB file to save updated coordinates.
    - coordinates_tensor (torch.Tensor): Tensor of shape (1, N, 3) with transformed coordinates.
    """
    # Convert the tensor to a list of tuples
    new_coordinates = coordinates_tensor.squeeze(0).tolist()

    # Create a parser and parse the structure
    parser = PDB.PDBParser(QUIET=True)
    structure = parser.get_structure("structure", input_filename)

    # Flattened iterator for atoms to update coordinates
    atom_iterator = (atom for model in structure for chain in model for residue in chain for atom in residue)

    # Update each atom's coordinates
    for atom, (new_x, new_y, new_z) in zip(atom_iterator, new_coordinates):
        original_anisou = atom.get_anisou()
        original_uij = atom.get_siguij()
        original_tm= atom.get_sigatm()
        original_occupancy = atom.get_occupancy()
        original_bfactor = atom.get_bfactor()
        original_altloc = atom.get_altloc()
        original_fullname = atom.get_fullname()
        original_serial_number = atom.get_serial_number()
        original_element = atom.get_charge()
        original_id = atom.get_full_id()
        original_idx = atom.get_id()
        original_level = atom.get_level()
        original_name = atom.get_name()
        original_parent = atom.get_parent()
        original_radius = atom.get_radius()
        

        # Update only the atom coordinates, keep other fields intact
        atom.coord = np.array([new_x, new_y, new_z])

        # Reapply the preserved properties
        atom.set_anisou(original_anisou)
        atom.set_siguij(original_uij)
        atom.set_sigatm(original_tm)
        atom.set_occupancy(original_occupancy)
        atom.set_bfactor(original_bfactor)
        atom.set_altloc(original_altloc)
        # atom.set_fullname(original_fullname)
        atom.set_serial_number(original_serial_number)
        atom.set_charge(original_element)
        atom.set_radius(original_radius)
        atom.set_parent(original_parent)
        # atom.set_name(original_name)
        # atom.set_leve
    
    output_filename = "/tmp/" + output_filename
    # Save the updated structure to a new PDB file
    io = PDBIO()
    io.set_structure(structure)
    io.save(output_filename)

    # Return the path to the updated PDB file
    return output_filename
 
def merge_pdb_files(file1, file2, output_file):
    r"""
    Merges two PDB files by concatenating them without altering their contents.
    
    Parameters:
    - file1 (str): Path to the first PDB file (e.g., receptor).
    - file2 (str): Path to the second PDB file (e.g., ligand).
    - output_file (str): Path to the output file where the merged structure will be saved.
    """
    output_file = "/tmp/" + output_file
    with open(output_file, 'w') as outfile:
        # Copy the contents of the first file
        with open(file1, 'r') as f1:
            lines = f1.readlines()
            # Write all lines except the last 'END' line
            outfile.writelines(lines[:-1])
        # Copy the contents of the second file
        with open(file2, 'r') as f2:
            outfile.write(f2.read())
    
    print(f"Merged PDB saved to {output_file}")   
    return output_file

class MPNNLayer(MessagePassing):
    def __init__(self, emb_dim=64, edge_dim=4, aggr='add'):
        r"""Message Passing Neural Network Layer

        Args:
            emb_dim: (int) - hidden dimension d
            edge_dim: (int) - edge feature dimension d_e
            aggr: (str) - aggregation function \oplus (sum/mean/max)
        """
        # Set the aggregation function
        super().__init__(aggr=aggr)

        self.emb_dim = emb_dim
        self.edge_dim = edge_dim

        # MLP \psi for computing messages m_ij
        # Implemented as a stack of Linear->BN->ReLU->Linear->BN->ReLU
        # dims: (2d + d_e) -> d
        self.mlp_msg = Sequential(
            Linear(2*emb_dim + edge_dim, emb_dim), BatchNorm1d(emb_dim), ReLU(),
            Linear(emb_dim, emb_dim), BatchNorm1d(emb_dim), ReLU()
          )
        
        # MLP \phi for computing updated node features h_i^{l+1}
        # Implemented as a stack of Linear->BN->ReLU->Linear->BN->ReLU
        # dims: 2d -> d
        self.mlp_upd = Sequential(
            Linear(2*emb_dim, emb_dim), BatchNorm1d(emb_dim), ReLU(), 
            Linear(emb_dim, emb_dim), BatchNorm1d(emb_dim), ReLU()
          )

    def forward(self, h, edge_index, edge_attr):
        r"""
        The forward pass updates node features h via one round of message passing.

        As our MPNNLayer class inherits from the PyG MessagePassing parent class,
        we simply need to call the propagate() function which starts the 
        message passing procedure: message() -> aggregate() -> update().
        
        The MessagePassing class handles most of the logic for the implementation.
        To build custom GNNs, we only need to define our own message(), 
        aggregate(), and update() functions (defined subsequently).

        Args:
            h: (n, d) - initial node features
            edge_index: (e, 2) - pairs of edges (i, j)
            edge_attr: (e, d_e) - edge features

        Returns:
            out: (n, d) - updated node features
        """
        out = self.propagate(edge_index, h=h, edge_attr=edge_attr)
        return out

    def message(self, h_i, h_j, edge_attr):
        r"""Step (1) Message

        The message() function constructs messages from source nodes j 
        to destination nodes i for each edge (i, j) in edge_index.

        The arguments can be a bit tricky to understand: message() can take 
        any arguments that were initially passed to propagate. Additionally, 
        we can differentiate destination nodes and source nodes by appending 
        _i or _j to the variable name, e.g. for the node features h, we
        can use h_i and h_j. 
        
        This part is critical to understand as the message() function
        constructs messages for each edge in the graph. The indexing of the
        original node features h (or other node variables) is handled under
        the hood by PyG.

        Args:
            h_i: (e, d) - destination node features
            h_j: (e, d) - source node features
            edge_attr: (e, d_e) - edge features
        
        Returns:
            msg: (e, d) - messages m_ij passed through MLP \psi
        """
        msg = torch.cat([h_i, h_j, edge_attr], dim=-1)
        return self.mlp_msg(msg)
    
    def aggregate(self, inputs, index):
        r"""Step (2) Aggregate

        The aggregate function aggregates the messages from neighboring nodes,
        according to the chosen aggregation function ('sum' by default).

        Args:
            inputs: (e, d) - messages m_ij from destination to source nodes
            index: (e, 1) - list of source nodes for each edge/message in input

        Returns:
            aggr_out: (n, d) - aggregated messages m_i
        """
        return scatter(inputs, index, dim=self.node_dim, reduce=self.aggr)
    
    def update(self, aggr_out, h):
        r"""
        Step (3) Update

        The update() function computes the final node features by combining the 
        aggregated messages with the initial node features.

        update() takes the first argument aggr_out, the result of aggregate(), 
        as well as any optional arguments that were initially passed to 
        propagate(). E.g. in this case, we additionally pass h.

        Args:
            aggr_out: (n, d) - aggregated messages m_i
            h: (n, d) - initial node features

        Returns:
            upd_out: (n, d) - updated node features passed through MLP \phi
        """
        upd_out = torch.cat([h, aggr_out], dim=-1)
        return self.mlp_upd(upd_out)

    def __repr__(self) -> str:
        return (f'{self.__class__.__name__}(emb_dim={self.emb_dim}, aggr={self.aggr})')
class MPNNModel(Module):
    def __init__(self, num_layers=4, emb_dim=64, in_dim=11, edge_dim=4, out_dim=1):
        r"""Message Passing Neural Network model for graph property prediction

        Args:
            num_layers: (int) - number of message passing layers L
            emb_dim: (int) - hidden dimension d
            in_dim: (int) - initial node feature dimension d_n
            edge_dim: (int) - edge feature dimension d_e
            out_dim: (int) - output dimension (fixed to 1)
        """
        super().__init__()
        
        # Linear projection for initial node features
        # dim: d_n -> d
        self.lin_in = Linear(in_dim, emb_dim)
        
        # Stack of MPNN layers
        self.convs = torch.nn.ModuleList()
        for layer in range(num_layers):
            self.convs.append(MPNNLayer(emb_dim, edge_dim, aggr='add'))
        
        # Global pooling/readout function R (mean pooling)
        # PyG handles the underlying logic via global_mean_pool()
        self.pool = global_mean_pool

        # Linear prediction head
        # dim: d -> out_dim
        self.lin_pred = Linear(emb_dim, out_dim)
        
    def forward(self, data):
        r"""
        Args:
            data: (PyG.Data) - batch of PyG graphs

        Returns: 
            out: (batch_size, out_dim) - prediction for each graph
        """
        h = self.lin_in(data.x) # (n, d_n) -> (n, d)
        
        for conv in self.convs:
            h = h + conv(h, data.edge_index, data.edge_attr) # (n, d) -> (n, d)
            # Note that we add a residual connection after each MPNN layer

        h_graph = self.pool(h, data.batch) # (n, d) -> (batch_size, d)

        out = self.lin_pred(h_graph) # (batch_size, d) -> (batch_size, 1)

        return out.view(-1)
    

class EquivariantMPNNLayer(MessagePassing):
    def __init__(self, emb_dim=64,  aggr='add'):
        r"""Message Passing Neural Network Layer

        This layer is equivariant to 3D rotations and translations.

        Args:
            emb_dim: (int) - hidden dimension d
            edge_dim: (int) - edge feature dimension d_e
            aggr: (str) - aggregation function \oplus (sum/mean/max)
        """
        # Set the aggregation function
        super().__init__(aggr=aggr)

        self.emb_dim = emb_dim
       
 
        #
        self.mlp_msg =  Sequential(
                  Linear(2 * emb_dim  + 1, emb_dim),
                  BatchNorm1d(emb_dim),
                  ReLU(),
                  Linear(emb_dim, emb_dim),
                   BatchNorm1d(emb_dim),
                   ReLU()
                    )
                       
        
        self.mlp_pos = Sequential(
                 Linear(emb_dim, emb_dim),
                 BatchNorm1d(emb_dim),
                 ReLU(),
                 Linear(emb_dim,1)
        ) # MLP \psi
        self.mlp_upd = Sequential(
                       Linear(2*emb_dim, emb_dim), BatchNorm1d(emb_dim), ReLU(), Linear(emb_dim,emb_dim), BatchNorm1d(emb_dim), ReLU()
        )  # MLP \phi
        # ===========================================

    def forward(self, h, pos, edge_index):
        r"""
        The forward pass updates node features h via one round of message passing.

        Args:
            h: (n, d) - initial node features
            pos: (n, 3) - initial node coordinates
            edge_index: (e, 2) - pairs of edges (i, j)
            edge_attr: (e, d_e) - edge features

        Returns:
            out: [(n, d),(n,3)] - updated node features
        """
       
        #
        out = self.propagate(edge_index=edge_index, h=h, pos=pos)
        return out
        # ==========================================

    
    #
    def message(self, h_i,h_j,pos_i,pos_j):
        # Compute distance between nodes i and j (Euclidean distance)
        #distance_ij = torch.norm(pos_i - pos_j, dim=-1, keepdim=True)  # (e, 1)
        pos_diff = pos_i - pos_j
        dists = torch.norm(pos_diff,dim=-1).unsqueeze(1)

        # Concatenate node features, edge features, and distance
        msg = torch.cat([h_i , h_j, dists], dim=-1)
        msg = self.mlp_msg(msg)
        pos_diff = pos_diff * self.mlp_pos(msg)  # (e, 2d + d_e + 1)

      
  # (e, d)
        return msg , pos_diff
    #   ...  
    #
    def aggregate(self, inputs, index):
        r"""The aggregate function aggregates the messages from neighboring nodes,
        according to the chosen aggregation function ('sum' by default).

        Args:
            inputs: (e, d) - messages m_ij from destination to source nodes
            index: (e, 1) - list of source nodes for each edge/message in input

        Returns:
            aggr_out: (n, d) - aggregated messages m_i
        """
        msgs , pos_diffs = inputs

        msg_aggr = scatter(msgs, index , dim = self.node_dim , reduce = self.aggr)

        pos_aggr = scatter(pos_diffs, index, dim = self.node_dim , reduce = "mean")


        return msg_aggr , pos_aggr
    
    def update(self, aggr_out, h , pos):
        msg_aggr , pos_aggr = aggr_out

        upd_out = self.mlp_upd(torch.cat((h, msg_aggr), dim=-1))

        upd_pos = pos + pos_aggr

        return upd_out , upd_pos
    

    def __repr__(self) -> str:
        return (f'{self.__class__.__name__}(emb_dim={self.emb_dim}, aggr={self.aggr})')
    
class FinalMPNNModel(MPNNModel):
    def __init__(self, num_layers=4, emb_dim=64, in_dim=3,  num_heads = 2):
        r"""Message Passing Neural Network model for graph property prediction

        This model uses both node features and coordinates as inputs, and
        is invariant to 3D rotations and translations (the constituent MPNN layers
        are equivariant to 3D rotations and translations).

        Args:
            num_layers: (int) - number of message passing layers L
            emb_dim: (int) - hidden dimension d
            in_dim: (int) - initial node feature dimension d_n
            edge_dim: (int) - edge feature dimension d_e
            out_dim: (int) - output dimension (fixed to 1)
        """
        super().__init__()
        
        # Linear projection for initial node features
        # dim: d_n -> d
        self.lin_in = Linear(in_dim, emb_dim)
        self.equiv_layer = EquivariantMPNNLayer(emb_dim=emb_dim)
        # Stack of MPNN layers
        self.convs = torch.nn.ModuleList()
        for layer in range(num_layers):
            self.convs.append(EquivariantMPNNLayer(emb_dim, aggr='add'))
       
        
        self.cross_attention = nn.MultiheadAttention(emb_dim, num_heads, batch_first=True)
        self.fc_rotation = nn.Linear(emb_dim, 9)
        self.fc_translation = nn.Linear(emb_dim, 3)
        # Global pooling/readout function R (mean pooling)
        # PyG handles the underlying logic via global_mean_pool()
        # self.pool = global_mean_pool

    def naive_single(self, receptor, ligand , receptor_edge_index , ligand_edge_index):
        r"""
        Processes a single receptor-ligand pair.

        Args:
            receptor: Tensor of shape (1, num_receptor_atoms, 3) (receptor coordinates)
            ligand: Tensor of shape (1, num_ligand_atoms, 3) (ligand coordinates)

        Returns:
            rotation_matrix: Tensor of shape (1, 3, 3) predicted rotation matrix for the ligand.
            translation_vector: Tensor of shape (1, 3) predicted translation vector for the ligand.
        
        """

        
        # h_receptor = receptor  # Initial node features for the receptor
        # h_ligand = ligand 
        h_receptor = self.lin_in(receptor)
        h_ligand = self.lin_in(ligand)      # Initial node features for the ligand
        pos_receptor = receptor  # Initial positions
        pos_ligand = ligand 
        
        for layer in self.convs:
        # Apply the equivariant message-passing layer for both receptor and ligand
            h_receptor, pos_receptor = layer(h_receptor, pos_receptor,receptor_edge_index  )
            h_ligand, pos_ligand = layer(h_ligand, pos_ligand, ligand_edge_index)
            # print("Shape of h_receptor:", h_receptor.shape)
            # print("Shape of h_ligand:", h_ligand.shape)
        # Pass the layer outputs through MLPs for embeddings
            emb_features_receptor = h_receptor
            emb_features_ligand = h_ligand
            
        attn_output, _ = self.cross_attention(emb_features_receptor, emb_features_ligand, emb_features_ligand)
        rotation_matrix = self.fc_rotation(attn_output.mean(dim=0))
        rotation_matrix = rotation_matrix.view(-1, 3, 3)
        translation_vector = self.fc_translation(attn_output.mean(dim=0))
        return rotation_matrix, translation_vector
        
        
   

    def forward(self, data):
        r"""
        The main forward pass of the model.

        Args:
            batch: Same as in forward_rot_trans.

        Returns:
            transformed_ligands: List of tensors, each of shape (1, num_ligand_atoms, 3)
            representing the transformed ligand coordinates after applying the predicted
            rotation and translation.
        """
        receptor = data['receptor']['pos']
        ligand = data['ligand']['pos']
        receptor_edge_index = data['receptor']['edge_index']
        ligand_edge_index = data['ligand']['edge_index']
        
        rotation_matrix, translation_vector = self.naive_single(receptor, ligand,receptor_edge_index , ligand_edge_index)
        # for i in range(len(ligands)):
        #     ligands[i] = ligands[i] @ rotation_matrix[i] + translation_vector[i]
        ligands = data['ligand']['pos'] @ rotation_matrix + translation_vector
        return ligands 

class FinalMPNNModelight(pl.LightningModule):
    def __init__(self, num_layers=4, emb_dim=32, in_dim=3, num_heads=1, lr=1e-4):
        super().__init__()
        
        self.lin_in = nn.Linear(in_dim, emb_dim)
        self.convs = nn.ModuleList([EquivariantMPNNLayer(emb_dim, aggr='add') for _ in range(num_layers)])
        self.cross_attention = nn.MultiheadAttention(emb_dim, num_heads, batch_first=True)
        self.fc_rotation = nn.Linear(emb_dim, 9)
        self.fc_translation = nn.Linear(emb_dim, 3)
        self.lr = lr
        

    def naive_single(self, receptor, ligand, receptor_edge_index, ligand_edge_index):
        h_receptor = self.lin_in(receptor)
        h_ligand = self.lin_in(ligand)
        pos_receptor, pos_ligand = receptor, ligand

        for layer in self.convs:
            h_receptor, pos_receptor = layer(h_receptor, pos_receptor, receptor_edge_index)
            h_ligand, pos_ligand = layer(h_ligand, pos_ligand, ligand_edge_index)

        attn_output, _ = self.cross_attention(h_receptor, h_ligand, h_ligand)
        rotation_matrix = self.fc_rotation(attn_output.mean(dim=0)).view(-1, 3, 3)
        translation_vector = self.fc_translation(attn_output.mean(dim=0))
        return rotation_matrix, translation_vector

    def forward(self, data):
        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
        receptor = data['receptor']['pos'].to(device)
        ligand = data['ligand']['pos'].to(device)
        receptor_edge_index = data['receptor', 'receptor']['edge_index'].to(device)
        ligand_edge_index = data['ligand', 'ligand']['edge_index'].to(device)

        rotation_matrix, translation_vector = self.naive_single(receptor, ligand, receptor_edge_index, ligand_edge_index)
        # transformed_ligand = torch.matmul(ligand ,rotation_matrix) + translation_vector
        return rotation_matrix, translation_vector

   
    def training_step(self, batch, batch_idx):
        ligand_pred = self(batch)
        ligand_true = batch['ligand']['y']
        loss = F.mse_loss(ligand_pred.squeeze(0), ligand_true)
        self.log('train_loss', loss, batch_size=8)
        return loss

    
    def validation_step(self, batch, batch_idx):
        ligand_pred = self(batch)
        ligand_true = batch['ligand']['y']
        loss = F.l1_loss(ligand_pred.squeeze(0), ligand_true)
        
        self.log('val_loss', loss, prog_bar=True, batch_size=8)
        
        return loss

    
    def test_step(self, batch, batch_idx):
        ligand_pred = self(batch)
        ligand_true = batch['ligand']['y']
        loss = F.l1_loss(ligand_pred.squeeze(0), ligand_true)  
        self.log('test_loss', loss, prog_bar=True, batch_size=8)
        return loss
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode="min", factor=0.1, patience=5
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val_loss",  # Monitor validation loss to adjust the learning rate
            },
        }

model_path = "./EquiMPNN-epoch=413-val_loss=9.25-val_acc=0.00.ckpt"
model = FinalMPNNModelight.load_from_checkpoint(model_path)
trainer = pl.Trainer(
   
   
   fast_dev_run=False,
   accelerator="gpu" if torch.cuda.is_available() else "cpu",
   precision="bf16-mixed",
  
   devices=1,
)
model.eval()
def predict (input_seq_1, input_msa_1, input_protein_1, input_seq_2,input_msa_2,  input_protein_2):
    start_time = time.time()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    data = create_graph(input_protein_1, input_protein_2, k=10)
    R_chain, L_chain = ["R"], ["L"]
    with torch.no_grad():
        mat, vect = model(data)
    mat = mat.to(device)
    vect = vect.to(device)
    ligand1 = torch.tensor(extract_coordinates_from_pdb(input_protein_1),dtype=torch.float).to(device)
    # receptor1 = torch.tensor(extract_coordinates_from_pdb(input_protein_2),dtype=torch.float).to(device)
    transformed_ligand = torch.matmul(ligand1, mat) + vect
    # transformed_receptor = torch.matmul(receptor1, mat) + vect
    file1 = update_pdb_coordinates_from_tensor(input_protein_1, "holo_ligand.pdb", transformed_ligand)
    # file2 = update_pdb_coordinates_from_tensor(input_protein_2, "holo_receptor.pdb", transformed_receptor)
    out_pdb = merge_pdb_files(file1,input_protein_2,"output.pdb")
    
    # return an output pdb file with the protein and two chains A and B.  
    # also return a JSON with any metrics you want to report
    metrics = {"mean_plddt": 80, "binding_affinity": 2}
#     native = './test_out (1).pdb'
#     decoys = out_pdb
#     bdq = BiotiteDockQ(
#     native=native, decoys=decoys,
#     # These are optional and if not specified will be assigned based on number of atoms (receptor > ligand)
#     native_receptor_chain=R_chain,
#     native_ligand_chain=L_chain,
#     decoy_receptor_chain=R_chain,
#     decoy_ligand_chain=L_chain,
# )
#     dockq = bdq.calculate()
#     metrics['DockQ'] = dockq
    end_time = time.time()
    run_time = end_time - start_time
    
    return out_pdb,json.dumps(metrics), run_time

with gr.Blocks() as app:

    gr.Markdown("# Template for inference")

    gr.Markdown("EquiMPNN MOdel")   
    with gr.Row():
        with gr.Column():
            input_seq_1 = gr.Textbox(lines=3, label="Input Protein 1 sequence (FASTA)")
            input_msa_1 = gr.File(label="Input MSA Protein 1 (A3M)")
            input_protein_1 = gr.File(label="Input Protein 2 monomer (PDB)")
        with gr.Column():
            input_seq_2 = gr.Textbox(lines=3, label="Input Protein 2 sequence (FASTA)")
            input_msa_2 = gr.File(label="Input MSA Protein 2 (A3M)")
            input_protein_2 = gr.File(label="Input Protein 2 structure (PDB)")
        
        
    
    # define any options here

    # for automated inference the default options are used
    # slider_option = gr.Slider(0,10, label="Slider Option")
    # checkbox_option = gr.Checkbox(label="Checkbox Option")
    # dropdown_option = gr.Dropdown(["Option 1", "Option 2", "Option 3"], label="Radio Option")

    btn = gr.Button("Run Inference")

    gr.Examples(
        [
            [
                "GSGSPLAQQIKNIHSFIHQAKAAGRMDEVRTLQENLHQLMHEYFQQSD",
                "3v1c_A.pdb",
                "GSGSPLAQQIKNIHSFIHQAKAAGRMDEVRTLQENLHQLMHEYFQQSD",
                "3v1c_B.pdb",
                
            ],
        ],
        [input_seq_1, input_protein_1, input_seq_2,  input_protein_2],
    )
    reps =    [
    {
      "model": 0,
      "style": "cartoon",
      "chain": "A",
      "color": "whiteCarbon",
    },
    {
      "model": 0,
      "style": "cartoon",
       "chain": "B",
      "color": "greenCarbon",
    },
    {
      "model": 0,
      "chain": "A",
      "style": "stick",
      "sidechain": True,
      "color": "whiteCarbon",
    },
    {
      "model": 0,
      "chain": "B",
      "style": "stick",
      "sidechain": True,
      "color": "greenCarbon"
    }      
  ]
    # outputs 
    
    out = Molecule3D(reps=reps)
    metrics = gr.JSON(label="Metrics")
    run_time = gr.Textbox(label="Runtime")

    btn.click(predict, inputs=[input_seq_1, input_msa_1, input_protein_1, input_seq_2, input_msa_2,  input_protein_2], outputs=[out, metrics, run_time])

app.launch()