import torch
from torch import nn
SEQ_TYPES=("mutation","protein","ptm")
ALL_TYPES=("disease","mutation","protein","ptm","pathway")
class MLP(nn.Module):
 def __init__(self):
  super().__init__();self.net=nn.Sequential(nn.Linear(64,256),nn.ReLU(),nn.Linear(256,640))
 def forward(self,x):return self.net(x)
class ResidualV2(nn.Module):
 def __init__(self,relations,beta_init=0.0):
  super().__init__();self.residual=nn.ModuleDict({k:MLP() for k in SEQ_TYPES});self.fallback=nn.ModuleDict({k:MLP() for k in SEQ_TYPES});self.project=nn.ModuleDict({k:MLP() for k in ("disease","pathway")});self.norm=nn.ModuleDict({k:nn.LayerNorm(640) for k in ALL_TYPES});self.beta=nn.Parameter(torch.tensor(float(beta_init)));self.relation_vectors=nn.Parameter(torch.empty(len(relations),640));nn.init.xavier_uniform_(self.relation_vectors);self.relations=list(relations)
 def encode(self,node_type,graph,esm=None,valid=None):
  if node_type in SEQ_TYPES:
   r=self.residual[node_type](graph)
   if esm is None:raise ValueError("ESM required")
   z=self.norm[node_type](esm+self.beta*r);fb=self.norm[node_type](self.fallback[node_type](graph))
   if valid is not None:z=torch.where(valid[:,None],z,fb)
   preserve_mask=valid if valid is not None else torch.ones(len(z),dtype=torch.bool,device=z.device)
   return z,r,preserve_mask
  z=self.norm[node_type](self.project[node_type](graph));return z,z,torch.zeros(len(z),dtype=torch.bool,device=z.device)
 def score(self,h,t,relation_index):return (h*self.relation_vectors[relation_index]*t).sum(-1)/(640**0.5)