from pathlib import Path
import numpy as np,pandas as pd,torch
SEQ_TYPES=("mutation","protein","ptm")
COUNTS={"disease":19019,"mutation":1055821,"protein":14231,"ptm":166288,"pathway":2791}
class Inputs:
 def __init__(self,root):
  self.root=Path(root);g=self.root/"work/graph_embedding_assets_v1/embeddings";e=self.root/"work/esm2_only_v1/core_embeddings"
  self.graph={k:np.load(g/f"{k}_embeddings.npy",mmap_mode="r") for k in COUNTS};self.esm={k:np.load(e/f"{k}_esm2_embeddings.npy",mmap_mode="r") for k in SEQ_TYPES};self.valid={k:np.load(e/f"{k}_sequence_valid_mask.npy",mmap_mode="r") for k in SEQ_TYPES}
 def batch(self,typ,idx,device):
  g=torch.as_tensor(np.asarray(self.graph[typ][idx],dtype=np.float32),device=device)
  if typ in SEQ_TYPES:
   e=torch.as_tensor(np.asarray(self.esm[typ][idx],dtype=np.float32),device=device);v=torch.as_tensor(np.asarray(self.valid[typ][idx],dtype=bool),device=device);return g,e,v
  return g,None,None
def load_edges(root,sanity=False):
 root=Path(root);ed=root/"work/graph_core/edges";base=[]
 for p in sorted(ed.glob("*.tsv.gz")):
  src,rel,dst=p.name.replace(".tsv.gz","").split("__");x=pd.read_csv(p,sep="\t")[["src_index","dst_index"]].to_numpy(np.int64);base.append([src,rel,dst,x])
 if sanity:
  mp=next(x for s,r,d,x in base if s=="mutation" and d=="protein");mset=set(mp[mp[:,1]<500,0].tolist())
  pp=next(x for s,r,d,x in base if s=="protein" and d=="ptm");pset=set(range(500));ptmset=set(pp[pp[:,0]<500,1].tolist())
  small=[]
  for s,r,d,x in base:
   keep=np.ones(len(x),bool)
   if s=="mutation" and d=="protein":keep=np.isin(x[:,0],list(mset))&np.isin(x[:,1],list(pset))
   elif s=="protein":keep=np.isin(x[:,0],list(pset))
   elif d=="mutation":keep=np.isin(x[:,1],list(mset))
   elif d=="protein":keep=np.isin(x[:,1],list(pset))
   elif d=="ptm":keep=np.isin(x[:,1],list(ptmset))
   x=x[keep][:5000];small.append([s,r,d,x])
  base=small
 out=[]
 for s,r,d,x in base:
  if not len(x):raise RuntimeError("empty sanity relation "+r)
  out.append((s,r,d,x));out.append((d,"reverse_"+r,s,x[:,[1,0]]))
 return out