扎克伯格 Biohub 蛋白质生物学"世界模型"：AI 颠覆药物发现的全景解析

# ESMC 表示提取伪代码
from esmc import ESMCModel, ESMCTokenizer

# 初始化模型和分词器
model = ESMCModel.from_pretrained("biohub/esmc-600m")
tokenizer = ESMCTokenizer.from_pretrained("biohub/esmc-600m")

# 输入蛋白质序列
sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDWERVMGDGERQFSTLKSTVEAIWAGIKATEAAVSEEFGLAPFLPDQIHFVHSQELLSRYPDLDAKGRERAIAKDLGAVFLVGIGGKLSDGHRHDVRAPDYDDWSTPSELGHAGLNGDILVWNPVLEDAFELSSMGIRVDADTLKHQLALTGDEDRLELEWHQALLRGEMPQTIGGGIGQSRLTMLLLQLPHIGQEQFGKQISELMTAVIAEIEKRISSLREKHKAEKRPg"
inputs = tokenizer(sequence, return_tensors="pt")

# 获取序列表示
with torch.no_grad():
    outputs = model(**inputs)
    # last_hidden_state 包含每个位置的表示
    sequence_embedding = outputs.last_hidden_state.mean(dim=1)

"""
ESMC 蛋白质语言模型推理模块
功能：提取蛋白质序列的语义表示，用于下游任务
"""

import torch
from transformers import AutoTokenizer, AutoModel
from typing import List, Dict, Tuple
import numpy as np


class ProteinRepresentationExtractor:
    """
    蛋白质序列表示提取器
    基于 ESMC 模型生成固定维度的蛋白质嵌入向量
    """
    
    def __init__(self, model_name: str = "biohub/esmc-600m", device: str = "cuda"):
        """
        初始化模型和分词器
        
        Args:
            model_name: HuggingFace 模型名称
            device: 计算设备 ('cuda' 或 'cpu')
        """
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        print(f"Loading model: {model_name} on {self.device}")
        
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()
    
    @torch.no_grad()
    def extract_representation(
        self, 
        sequence: str, 
        pooling: str = "mean"
    ) -> np.ndarray:
        """
        提取单个蛋白质序列的向量表示
        
        Args:
            sequence: 氨基酸序列字符串
            pooling: 池化策略 ('mean', 'cls', 'max')
        
        Returns:
            蛋白质嵌入向量 (hidden_size,)
        """
        # Tokenize
        inputs = self.tokenizer(
            sequence,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=1024
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # 前向传播
        outputs = self.model(**inputs)
        hidden_states = outputs.last_hidden_state  # (batch, seq_len, hidden)
        
        # 池化
        if pooling == "mean":
            embeddings = hidden_states.mean(dim=1)
        elif pooling == "cls":
            embeddings = hidden_states[:, 0, :]
        elif pooling == "max":
            embeddings = hidden_states.max(dim=1)[0]
        else:
            raise ValueError(f"Unknown pooling strategy: {pooling}")
        
        return embeddings.cpu().numpy()[0]
    
    @torch.no_grad()
    def batch_extract(
        self,
        sequences: List[str],
        batch_size: int = 8,
        pooling: str = "mean"
    ) -> np.ndarray:
        """
        批量提取多个蛋白质序列的表示
        
        Args:
            sequences: 蛋白质序列列表
            batch_size: 批处理大小
            pooling: 池化策略
        
        Returns:
            嵌入矩阵 (n_sequences, hidden_size)
        """
        all_embeddings = []
        
        for i in range(0, len(sequences), batch_size):
            batch = sequences[i:i + batch_size]
            inputs = self.tokenizer(
                batch,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=1024
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            outputs = self.model(**inputs)
            hidden_states = outputs.last_hidden_state
            
            if pooling == "mean":
                embeddings = hidden_states.mean(dim=1)
            elif pooling == "cls":
                embeddings = hidden_states[:, 0, :]
            else:
                embeddings = hidden_states.max(dim=1)[0]
            
            all_embeddings.append(embeddings.cpu().numpy())
            
            if (i // batch_size + 1) % 10 == 0:
                print(f"Processed {min(i + batch_size, len(sequences))}/{len(sequences)} sequences")
        
        return np.vstack(all_embeddings)
    
    def compute_similarity(
        self,
        seq1: str,
        seq2: str,
        metric: str = "cosine"
    ) -> float:
        """
        计算两个蛋白质序列的相似度
        
        Args:
            seq1: 第一个序列
            seq2: 第二个序列
            metric: 相似度度量 ('cosine', 'euclidean')
        
        Returns:
            相似度分数
        """
        emb1 = self.extract_representation(seq1)
        emb2 = self.extract_representation(seq2)
        
        if metric == "cosine":
            return np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
        elif metric == "euclidean":
            return -np.linalg.norm(emb1 - emb2)  # 负距离，越大越相似
        else:
            raise ValueError(f"Unknown metric: {metric}")


def find_similar_sequences(
    query_sequence: str,
    database_sequences: List[str],
    top_k: int = 10,
    model_name: str = "biohub/esmc-600m"
) -> List[Tuple[int, float]]:
    """
    在数据库中查找与查询序列最相似的蛋白质
    
    Args:
        query_sequence: 查询蛋白质序列
        database_sequences: 数据库中的蛋白质序列列表
        top_k: 返回前 k 个最相似结果
        model_name: 模型名称
    
    Returns:
        (index, similarity_score) 列表
    """
    extractor = ProteinRepresentationExtractor(model_name)
    
    # 提取查询表示
    query_emb = extractor.extract_representation(query_sequence)
    
    # 批量提取数据库表示
    db_embeddings = extractor.batch_extract(database_sequences)
    
    # 计算余弦相似度
    similarities = np.dot(db_embeddings, query_emb) / (
        np.linalg.norm(db_embeddings, axis=1) * np.linalg.norm(query_emb)
    )
    
    # 返回 top-k
    top_indices = np.argsort(similarities)[::-1][:top_k]
    return [(idx, similarities[idx]) for idx in top_indices]


# 使用示例
if __name__ == "__main__":
    # 示例蛋白质序列（人类 EGFR 部分序列）
    egfr_sequence = (
        "MRPSGTAGAALLALLAALCPASRALEEKKVCQGTSNKLTQLGTFEDHFLSLQRMFNNCEVVLGNLEITYVQRNYD"
        "LSFLKTIQEVAGYVILIALHSGRHRRTVTNGGALLEEDTTTNSATLQLSQVFGGFGGGEGGSGGSGGSVGAASED"
    )
    
    # 初始化提取器
    extractor = ProteinRepresentationExtractor("biohub/esmc-600m")
    
    # 提取表示
    embedding = extractor.extract_representation(egfr_sequence)
    print(f"Embedding shape: {embedding.shape}")
    print(f"Embedding sample (first 10 dims): {embedding[:10]}")
    
    # 批量处理示例
    sequences = [egfr_sequence, "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVK", "MLPGLIMLLHLHAYHVAGHNGIFWSIHGPGVPQAFEYGAQYVVQDLEGKVFSPQNFTSLQPGQ"]
    embeddings = extractor.batch_extract(sequences)
    print(f"Batch embeddings shape: {embeddings.shape}")

"""
ESMFold2 蛋白质结构预测模块
功能：基于蛋白质序列预测三维结构
支持抗体-抗原复合物预测
"""

import torch
import numpy as np
from typing import Optional, Dict, List, Tuple
import json
from dataclasses import dataclass


@dataclass
class StructureResult:
    """结构预测结果数据类"""
    sequence: str
    atom_positions: np.ndarray  # (L, 37, 3) 原子坐标
    confidence: np.ndarray      # (L,) pLDDT 置信度分数
    pae_matrix: np.ndarray      # 预测对齐误差矩阵
    mean_plddt: float


class ESMFold2Predictor:
    """
    ESMFold2 结构预测器
    
    特点：
    - 无需 MSA，直接从序列推理
    - 支持单链和复合物预测
    - 输出原子级精度结构
    """
    
    def __init__(self, model_name: str = "biohub/esmfold2", device: str = "cuda"):
        """
        初始化 ESMFold2 模型
        
        Args:
            model_name: 模型名称
            device: 计算设备
        """
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        print(f"Loading ESMFold2 on {self.device}")
        
        from transformers import EsmForProteinFolding
        
        self.model = EsmForProteinFolding.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()
    
    @torch.no_grad()
    def predict_structure(
        self,
        sequence: str,
        num_recycles: int = 3
    ) -> StructureResult:
        """
        预测单链蛋白质的三维结构
        
        Args:
            sequence: 氨基酸序列
            num_recycles: 循环次数（越多精度越高但越慢）
        
        Returns:
            StructureResult 对象
        """
        # Tokenize
        from transformers.models.esmfold_v2.configuration_esmfold_v2 import EsmFoldV2Config
        from transformers.models.esmfold_v2.tokenization_esmfast import EsmFastTokenizer
        
        tokenizer = EsmFastTokenizer.from_pretrained("facebook/esmfold_v2")
        tokens = tokenizer(
            [sequence],
            return_tensors="pt",
            padding=True
        )["input_ids"].to(self.device)
        
        # 推理
        outputs = self.model.model(
            esm_embeddings=None,
            esm_attention_mask=None,
            xs=tokens,
            num_recycles=num_recycles
        )
        
        # 提取结果
        atom_positions = outputs["positions"][0].cpu().numpy()  # (L, 37, 3)
        confidence = outputs["plddt"][0].cpu().numpy()         # (L,)
        pae_matrix = outputs["pae"][0].cpu().numpy()           # (L, L)
        
        return StructureResult(
            sequence=sequence,
            atom_positions=atom_positions,
            confidence=confidence,
            pae_matrix=pae_matrix,
            mean_plddt=float(confidence.mean())
        )
    
    @torch.no_grad()
    def predict_complex(
        self,
        chain_sequences: List[str],
        chain_ids: Optional[List[str]] = None
    ) -> Dict:
        """
        预测蛋白质复合物结构（如抗体-抗原）
        
        Args:
            chain_sequences: 各链的氨基酸序列列表
            chain_ids: 各链的标识符列表
        
        Returns:
            包含所有链结构的字典
        """
        if chain_ids is None:
            chain_ids = [f"chain_{i}" for i in range(len(chain_sequences))]
        
        # 合并序列（使用 <eos>` 分隔不同链）
        complex_sequence = "`".join(chain_sequences)
        complex_chain_ids = "`".join(chain_ids) + "`"
        
        from transformers.models.esmfold_v2.tokenization_esmfast import EsmFastTokenizer
        
        tokenizer = EsmFastTokenizer.from_pretrained("facebook/esmfold_v2")
        tokens = tokenizer(
            [complex_sequence],
            return_tensors="pt",
            padding=True
        )["input_ids"].to(self.device)
        
        # 推理
        outputs = self.model.model(
            esm_embeddings=None,
            esm_attention_mask=None,
            xs=tokens,
            num_recycles=3
        )
        
        # 解析结果
        results = {}
        offset = 0
        for i, (seq, cid) in enumerate(zip(chain_sequences, chain_ids)):
            chain_len = len(seq)
            results[cid] = {
                "sequence": seq,
                "atom_positions": outputs["positions"][0, offset:offset+chain_len].cpu().numpy(),
                "confidence": outputs["plddt"][0, offset:offset+chain_len].cpu().numpy(),
            }
            offset += chain_len
        
        return results
    
    def batch_predict(
        self,
        sequences: List[str],
        batch_size: int = 4
    ) -> List[StructureResult]:
        """
        批量预测多个蛋白质结构
        
        Args:
            sequences: 序列列表
            batch_size: 批大小
        
        Returns:
            StructureResult 列表
        """
        results = []
        for i in range(0, len(sequences), batch_size):
            batch = sequences[i:i + batch_size]
            print(f"Processing batch {i//batch_size + 1}/{(len(sequences)-1)//batch_size + 1}")
            
            for seq in batch:
                result = self.predict_structure(seq)
                results.append(result)
        
        return results
    
    def save_pdb(
        self,
        result: StructureResult,
        output_path: str
    ) -> None:
        """
        将预测结构保存为 PDB 格式
        
        Args:
            result: StructureResult 对象
            output_path: 输出文件路径
        """
        from Bio.PDB import PDBParser, Structure, Model, Chain, Residue, Atom
        
        # 原子类型和对应的残基中位置
        atom_names = ['N', 'CA', 'C', 'O', 'CB', 'CG', 'CD', 'CE', 'NZ', 'OG', 'OG1', 
                      'SD', 'OD1', 'OD2', 'OE1', 'OE2', 'NH1', 'NH2', 'OH', 'NE', 'ND1', 
                      'ND2', 'NE1', 'NE2', 'NH1', 'NH2', 'OXT']
        
        parser = PDBParser(QUIET=True)
        structure = parser.create_structure("model")
        model = Structure.Model(0)
        chain = Chain.Chain("A")
        
        seq = result.sequence
        positions = result.atom_positions
        
        for i, residue in enumerate(seq):
            if i >= len(positions):
                break
                
            res_id = (" ", i + 1, " ")
            residue_obj = Residue.Residue(res_id, residue, " ")
            
            for j, atom_name in enumerate(atom_names):
                if j >= len(positions[i]):
                    break
                    
                coords = positions[i][j]
                if np.isnan(coords).any():
                    continue
                    
                atom = Atom.Atom(atom_name, coords, 0, 0, " ", atom_name, i + 1)
                residue_obj.add(atom)
            
            chain.add(residue_obj)
        
        model.add(chain)
        structure.add(model)
        
        io = PDBIO()
        io.set_structure(structure)
        io.save(output_path)


def design_binder_for_target(
    target_sequence: str,
    target_binding_site: List[int],
    binder_length: int = 50,
    num_candidates: int = 100
) -> List[str]:
    """
    设计针对特定靶点结合位点的蛋白质结合物
    
    Args:
        target_sequence: 靶点蛋白质序列
        target_binding_site: 结合位点残基索引列表
        binder_length: 结合物长度
        num_candidates: 生成候选数量
    
    Returns:
        候选结合物序列列表
    """
    predictor = ESMFold2Predictor()
    
    # 简化的结合物设计流程（实际需要更复杂的优化）
    # 这里演示如何使用 ESMFold2 评估设计的序列
    
    candidate_sequences = []
    
    for i in range(num_candidates):
        # 随机生成候选序列（实际应用需要基于语言模型生成）
        candidate = generate_random_sequence(binder_length)
        candidate_sequences.append(candidate)
        
        if (i + 1) % 10 == 0:
            print(f"Generated {i + 1}/{num_candidates} candidates")
    
    return candidate_sequences


def generate_random_sequence(length: int) -> str:
    """生成随机氨基酸序列（仅用于演示）"""
    amino_acids = "ACDEFGHIKLMNPQRSTVWY"
    return "".join(np.random.choice(list(amino_acids)) for _ in range(length))


if __name__ == "__main__":
    # 示例：预测单链蛋白质结构
    sequence = "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSPLHSVYVDQWDWERVMGDGERQFSTLKSTVEAIWAGIKATEAAVSEEFGLAPFLPDQIHFVHSQELLSRYPDLDAKGRERAIAKDLGAVFLVGIGGKLSDGHRHDVRAPDYDDWSTPSELGHAGLNGDILVWNPVLEDAFELSSMGIRVDADTLKHQLALTGDEDRLELEWHQALLRGEMPQTIGGGIGQSRLTMLLLQLPHIGQEQFGKQISELMTAVIAEIEKRISSLREKHKAEKRP"
    
    predictor = ESMFold2Predictor()
    result = predictor.predict_structure(sequence)
    
    print(f"Sequence length: {len(sequence)}")
    print(f"Mean pLDDT (confidence): {result.mean_plddt:.2f}")
    print(f"Atom positions shape: {result.atom_positions.shape}")
    
    # 保存为 PDB
    predictor.save_pdb(result, "predicted_structure.pdb")
    print("Structure saved to predicted_structure.pdb")

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/elastic/go-elasticsearch/v8"
	"github.com/elastic/go-elasticsearch/v8/esapi"
)

/*
ESM Atlas 数据库查询模块
功能：查询68亿蛋白质序列和11亿预测结构
语言：Go
*/

// ProteinHit 蛋白质搜索结果
type ProteinHit struct {
	ID         string  `json:"id"`
	Sequence   string  `json:"sequence"`
	Taxonomy   string  `json:"taxonomy"`
	Source     string  `json:"source"`
	Similarity float64 `json:"similarity_score"`
}

// AtlasQueryResult Atlas查询结果
type AtlasQueryResult struct {
	Total    int64          `json:"total"`
	Hits     []ProteinHit   `json:"hits"`
	ScrollID string         `json:"scroll_id,omitempty"`
}

// ESMAtlasClient ESM Atlas 客户端
type ESMAtlasClient struct {
	client *elasticsearch.Client
}

// NewESMAtlasClient 创建新的 Atlas 客户端
func NewESMAtlasClient(addresses []string) (*ESMAtlasClient, error) {
	cfg := elasticsearch.Config{
		Addresses: addresses,
		Transport: &http.Transport{
			MaxIdleConnsPerHost:   10,
			ResponseHeaderTimeout: time.Second * 30,
		},
	}

	client, err := elasticsearch.NewClient(cfg)
	if err != nil {
		return nil, fmt.Errorf("failed to create client: %w", err)
	}

	// 检查连接
	res, err := client.Info()
	if err != nil {
		return nil, fmt.Errorf("failed to connect: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		return nil, fmt.Errorf("error response: %s", res.String())
	}

	log.Println("Connected to ESM Atlas successfully")
	return &ESMAtlasClient{client: client}, nil
}

// SearchBySequence 根据序列相似性搜索
func (c *ESMAtlasClient) SearchBySequence(
	ctx context.Context,
	querySequence string,
	limit int,
) (*AtlasQueryResult, error) {
	// 构建查询
	query := map[string]interface{}{
		"size": limit,
		"query": map[string]interface{}{
			"script_score": map[string]interface{}{
				"query": map[string]interface{}{
					"match_all": map[string]interface{}{},
				},
				"script": map[string]interface{}{
					"source": "esm_similarity(params.query_sequence, doc['sequence'].value)",
					"params": map[string]interface{}{
						"query_sequence": querySequence,
					},
				},
			},
		},
	}

	queryJSON, err := json.Marshal(query)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal query: %w", err)
	}

	// 执行搜索
	req := esapi.SearchRequest{
		Index: []string{"protein_sequences_v2"},
		Body:  queryJSON,
	}

	res, err := req.Do(ctx, c.client)
	if err != nil {
		return nil, fmt.Errorf("search failed: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		return nil, fmt.Errorf("search error: %s", res.String())
	}

	// 解析结果
	var response struct {
		Hits struct {
			Total struct {
				Value int64 `json:"value"`
			} `json:"total"`
			Hits []struct {
				ID        string `json:"_id"`
				Score     float64 `json:"_score"`
				Source struct {
					Sequence string `json:"sequence"`
					Taxonomy string `json:"taxonomy,omitempty"`
					Source   string `json:"source,omitempty"`
				} `json:"_source"`
			} `json:"hits"`
		} `json:"hits"`
	}

	if err := json.NewDecoder(res.Body).Decode(&response); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}

	// 转换结果
	hits := make([]ProteinHit, 0, len(response.Hits.Hits))
	for _, hit := range response.Hits.Hits {
		hits = append(hits, ProteinHit{
			ID:         hit.ID,
			Sequence:   hit.Source.Sequence,
			Taxonomy:   hit.Source.Taxonomy,
			Source:     hit.Source.Source,
			Similarity: hit.Score,
		})
	}

	return &AtlasQueryResult{
		Total: response.Hits.Total.Value,
		Hits:  hits,
	}, nil
}

// SearchByFunction 根据功能关键词搜索
func (c *ESMAtlasClient) SearchByFunction(
	ctx context.Context,
	functionKeyword string,
	limit int,
) (*AtlasQueryResult, error) {
	query := map[string]interface{}{
		"size": limit,
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"should": []map[string]interface{}{
					{
						"match": map[string]interface{}{
							"function_annotation": functionKeyword,
						},
					},
					{
						"match": map[string]interface{}{
							"go_terms": functionKeyword,
						},
					},
				},
			},
		},
	}

	queryJSON, err := json.Marshal(query)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal query: %w", err)
	}

	req := esapi.SearchRequest{
		Index: []string{"protein_sequences_v2"},
		Body:  queryJSON,
	}

	res, err := req.Do(ctx, c.client)
	if err != nil {
		return nil, fmt.Errorf("search failed: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		return nil, fmt.Errorf("search error: %s", res.String())
	}

	var response struct {
		Hits struct {
			Total struct {
				Value int64 `json:"value"`
			} `json:"total"`
			Hits []struct {
				ID        string `json:"_id"`
				Score     float64 `json:"_score"`
				Source struct {
					Sequence string `json:"sequence"`
					Taxonomy string `json:"taxonomy,omitempty"`
				} `json:"_source"`
			} `json:"hits"`
		} `json:"hits"`
	}

	if err := json.NewDecoder(res.Body).Decode(&response); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}

	hits := make([]ProteinHit, 0, len(response.Hits.Hits))
	for _, hit := range response.Hits.Hits {
		hits = append(hits, ProteinHit{
			ID:         hit.ID,
			Sequence:   hit.Source.Sequence,
			Taxonomy:   hit.Source.Taxonomy,
			Similarity: hit.Score,
		})
	}

	return &AtlasQueryResult{
		Total: response.Hits.Total.Value,
		Hits:  hits,
	}, nil
}

// GetProteinStructure 获取蛋白质预测结构
func (c *ESMAtlasClient) GetProteinStructure(
	ctx context.Context,
	proteinID string,
) ([]float64, error) {
	req := esapi.GetRequest{
		Index:      "protein_structures_v2",
		DocumentID: proteinID,
	}

	res, err := req.Do(ctx, c.client)
	if err != nil {
		return nil, fmt.Errorf("failed to get structure: %w", err)
	}
	defer res.Body.Close()

	if res.IsError() {
		if res.StatusCode == 404 {
			return nil, fmt.Errorf("structure not found for protein: %s", proteinID)
		}
		return nil, fmt.Errorf("get error: %s", res.String())
	}

	var response struct {
		Source struct {
			AtomCoordinates []float64 `json:"atom_coordinates"`
		} `json:"_source"`
	}

	if err := json.NewDecoder(res.Body).Decode(&response); err != nil {
		return nil, fmt.Errorf("failed to decode response: %w", err)
	}

	return response.Source.AtomCoordinates, nil
}

// BatchSearch 批量搜索多个序列
func (c *ESMAtlasClient) BatchSearch(
	ctx context.Context,
	sequences []string,
	limitPerQuery int,
) ([]*AtlasQueryResult, error) {
	results := make([]*AtlasQueryResult, 0, len(sequences))

	for i, seq := range sequences {
		result, err := c.SearchBySequence(ctx, seq, limitPerQuery)
		if err != nil {
			log.Printf("Warning: failed to search sequence %d: %v", i, err)
			continue
		}
		results = append(results, result)

		if (i+1)%10 == 0 {
			log.Printf("Processed %d/%d sequences", i+1, len(sequences))
		}
	}

	return results, nil
}

func main() {
	ctx := context.Background()

	// 连接到 ESM Atlas（示例地址）
	client, err := NewESMAtlasClient([]string{
		"https://api.esmatlas.com:9200",
	})
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
	}

	// 示例：搜索与 EGFR 相似的蛋白质
	egfrPartialSequence := "MRPSGTAGAALLALLAALCPASRALEEKKVCQGTSNKLTQLGTFEDHFLSLQRMFNNCEVVLGNLEITYVQRNYDLSFLKTIQEVAGYVILIALHSGRHRRTVTNGGALLEEDTTTNSATLQLSQVFGG"

	log.Println("Searching for EGFR-like proteins...")
	results, err := client.SearchBySequence(ctx, egfrPartialSequence, 20)
	if err != nil {
		log.Fatalf("Search failed: %v", err)
	}

	fmt.Printf("Found %d results:\n", results.Total)
	for i, hit := range results.Hits {
		fmt.Printf("%d. ID: %s, Similarity: %.4f, Taxonomy: %s\n",
			i+1, hit.ID, hit.Similarity, hit.Taxonomy)
	}

	// 示例：搜索功能关键词
	log.Println("\nSearching for kinases...")
	funcResults, err := client.SearchByFunction(ctx, "kinase activity", 10)
	if err != nil {
		log.Fatalf("Function search failed: %v", err)
	}

	fmt.Printf("Found %d kinase proteins:\n", funcResults.Total)
	for i, hit := range funcResults.Hits {
		fmt.Printf("%d. ID: %s, Taxonomy: %s\n", i+1, hit.ID, hit.Taxonomy)
	}
}

"""
靶点-结合物亲和力预测模块
功能：使用语言模型表示预测蛋白质-配体结合亲和力
"""

import torch
import torch.nn as nn
import numpy as np
from typing import List, Tuple, Optional, Dict
from dataclasses import dataclass


@dataclass
class BindingPrediction:
    """结合预测结果"""
    target_id: str
    binder_id: str
    predicted_affinity: float  # nM
    confidence: float
    binding_interface_residues: List[int]
    notes: str


class BindingAffinityPredictor(nn.Module):
    """
    基于Transformer的蛋白质结合亲和力预测模型
    
    架构：
    1. 双塔编码器分别处理靶点和结合物
    2. 交互层建模界面残基间的相互作用
    3. 回归头输出亲和力预测
    """
    
    def __init__(
        self,
        embedding_dim: int = 1280,
        num_heads: int = 8,
        num_layers: int = 4,
        dropout: float = 0.1
    ):
        super().__init__()
        
        # 靶点和结合物编码器（共享权重）
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=embedding_dim,
            nhead=num_heads,
            dim_feedforward=embedding_dim * 4,
            dropout=dropout,
            batch_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        # 位置编码
        self.pos_encoder = nn.Parameter(torch.randn(1, 2048, embedding_dim) * 0.02)
        
        # 交互注意力
        self.cross_attention = nn.MultiheadAttention(
            embed_dim=embedding_dim,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True
        )
        
        # 亲和力预测头
        self.affinity_head = nn.Sequential(
            nn.Linear(embedding_dim * 3, embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim, embedding_dim // 2),
            nn.ReLU(),
            nn.Linear(embedding_dim // 2, 1)
        )
        
        # 界面预测头
        self.interface_head = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim // 2),
            nn.ReLU(),
            nn.Linear(embedding_dim // 2, 1),
            nn.Sigmoid()
        )
    
    def forward(
        self,
        target_embeddings: torch.Tensor,  # (batch, L1, dim)
        binder_embeddings: torch.Tensor,   # (batch, L2, dim)
        target_mask: Optional[torch.Tensor] = None,
        binder_mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        前向传播
        
        Args:
            target_embeddings: 靶点蛋白质表示 (batch, L1, dim)
            binder_embeddings: 结合物蛋白质表示 (batch, L2, dim)
            target_mask: 靶点注意力掩码
            binder_mask: 结合物注意力掩码
        
        Returns:
            (predicted_affinity, interface_probabilities)
        """
        batch_size = target_embeddings.size(0)
        
        # 添加位置编码
        target_seq_len = target_embeddings.size(1)
        binder_seq_len = binder_embeddings.size(1)
        
        target_embeddings = target_embeddings + self.pos_encoder[:, :target_seq_len, :]
        binder_embeddings = binder_embeddings + self.pos_encoder[:, :binder_seq_len, :]
        
        # 编码靶点和结合物
        target_encoded = self.encoder(target_embeddings, src_key_padding_mask=target_mask)
        binder_encoded = self.encoder(binder_embeddings, src_key_padding_mask=binder_mask)
        
        # 交叉注意力计算相互作用
        # (batch, L1, L2) 注意力矩阵
        interaction_attn, _ = self.cross_attention(
            target_encoded, binder_encoded, binder_encoded,
            key_padding_mask=binder_mask
        )
        
        # 聚合交互特征
        # 对结合物维度加权求和
        attn_weights = torch.softmax(interaction_attn.mean(dim=-1), dim=-1)  # (batch, L1)
        
        # 靶点-结合物交互特征
        interaction_feature = torch.einsum('bl,brd->brd', attn_weights, binder_encoded)
        
        # 拼接靶点特征、交互特征和界面特征
        combined = torch.cat([
            target_encoded,
            interaction_feature,
            target_encoded * interaction_feature  # 逐元素乘积
        ], dim=-1)  # (batch, L1, dim * 3)
        
        # 预测界面概率
        interface_logits = self.interface_head(combined).squeeze(-1)  # (batch, L1)
        
        # 使用界面特征预测整体亲和力
        # 对界面残基特征加权平均
        if target_mask is not None:
            mask_expanded = target_mask.float().unsqueeze(-1)
        else:
            mask_expanded = torch.ones_like(target_encoded)
        
        interface_feature_pooled = (
            (combined * mask_expanded).sum(dim=1) / 
            (mask_expanded.sum(dim=1) + 1e-8)
        )
        
        affinity = self.affinity_head(interface_feature_pooled).squeeze(-1)
        
        return affinity, interface_logits


class BindingAffinityPredictorWrapper:
    """结合亲和力预测器的完整工作流封装"""
    
    def __init__(
        self,
        model_path: str = "biohub/binding_affinity_model.pt",
        device: str = "cuda"
    ):
        """
        初始化预测器
        
        Args:
            model_path: 预训练模型路径
            device: 计算设备
        """
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        
        # 加载蛋白质表示提取器
        from esmc import ESMCModel, ESMCTokenizer
        self.esmc = ESMCModel.from_pretrained("biohub/esmc-600m")
        self.esmc_tokenizer = ESMCTokenizer.from_pretrained("biohub/esmc-600m")
        self.esmc.to(self.device)
        self.esmc.eval()
        
        # 加载结合预测模型
        self.model = BindingAffinityPredictor()
        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
        self.model.to(self.device)
        self.model.eval()
    
    @torch.no_grad()
    def get_protein_embeddings(
        self,
        sequence: str,
        max_length: int = 2048
    ) -> torch.Tensor:
        """
        获取蛋白质序列的嵌入表示
        
        Args:
            sequence: 氨基酸序列
            max_length: 最大序列长度
        
        Returns:
            蛋白质嵌入 (L, embedding_dim)
        """
        inputs = self.esmc_tokenizer(
            sequence,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=max_length
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        outputs = self.esmc(**inputs)
        embeddings = outputs.last_hidden_state  # (1, L, dim)
        
        # 创建掩码（排除 padding）
        mask = inputs["attention_mask"].bool()
        
        return embeddings, mask
    
    @torch.no_grad()
    def predict_binding(
        self,
        target_sequence: str,
        binder_sequence: str,
        target_id: str = "target",
        binder_id: str = "binder"
    ) -> BindingPrediction:
        """
        预测靶点-结合物的结合亲和力
        
        Args:
            target_sequence: 靶点蛋白质序列
            binder_sequence: 结合物蛋白质序列
            target_id: 靶点标识符
            binder_id: 结合物标识符
        
        Returns:
            BindingPrediction 结果
        """
        # 获取嵌入
        target_emb, target_mask = self.get_protein_embeddings(target_sequence)
        binder_emb, binder_mask = self.get_protein_embeddings(binder_sequence)
        
        # 预测
        affinity, interface_logits = self.model(
            target_emb, binder_emb, 
            ~target_mask, ~binder_mask
        )
        
        # 转换亲和力（从对数空间）
        predicted_affinity_nM = 10 ** affinity.item()
        
        # 识别界面残基
        interface_probs = torch.sigmoid(interface_logits[0]).cpu().numpy()
        interface_residues = np.where(interface_probs > 0.5)[0].tolist()
        
        # 置信度评估
        confidence = 1.0 - np.std(interface_probs)
        
        # 生成注释
        if predicted_affinity_nM < 1:
            notes = "极强结合 (亚纳摩尔级)"
        elif predicted_affinity_nM < 100:
            notes = "强结合 (纳摩尔级)"
        elif predicted_affinity_nM < 1000:
            notes = "中等结合 (亚微摩尔级)"
        else:
            notes = "弱结合 (微摩尔级)"
        
        return BindingPrediction(
            target_id=target_id,
            binder_id=binder_id,
            predicted_affinity=predicted_affinity_nM,
            confidence=confidence,
            binding_interface_residues=interface_residues,
            notes=notes
        )
    
    @torch.no_grad()
    def screen_binders(
        self,
        target_sequence: str,
        binder_sequences: List[str],
        threshold_affinity: float = 500.0  # nM
    ) -> List[BindingPrediction]:
        """
        批量筛选结合物
        
        Args:
            target_sequence: 靶点序列
            binder_sequences: 候选结合物序列列表
            threshold_affinity: 亲和力阈值（nM）
        
        Returns:
            通过筛选的结合物列表
        """
        results = []
        
        for i, binder_seq in enumerate(binder_sequences):
            try:
                result = self.predict_binding(
                    target_sequence,
                    binder_seq,
                    target_id=f"target_{i}",
                    binder_id=f"binder_{i}"
                )
                
                if result.predicted_affinity < threshold_affinity:
                    results.append(result)
                
                if (i + 1) % 100 == 0:
                    print(f"Screened {i + 1}/{len(binder_sequences)} binders")
                    
            except Exception as e:
                print(f"Error screening binder {i}: {e}")
                continue
        
        # 按亲和力排序
        results.sort(key=lambda x: x.predicted_affinity)
        
        return results


def main():
    """演示程序"""
    
    # 初始化预测器
    predictor = BindingAffinityPredictorWrapper(
        model_path="biohub/binding_affinity_model.pt"
    )
    
    # 示例：PD-L1 靶点和设计的结合物
    pd_l1_sequence = (
        "MLAE IOPLLILIIFSWLIGEIEKFLHPDQIEGWYTSDPSSVISGMLITVQAQFRTQQMGN"
        "YTQLVWYLMEGNYQKPQLAIITQDPSDQPPEMITDEDRDQAKTPVTSGLEVPTTSGVPSA"
    )
    
    designed_binder = (
        "MKTHIIMLVLLSVLLAQDGVTQDGQWTTQDGILTVSTKTFLSQPFRELAYDTPTSVFTQEV"
        "QPEGSPAASQPQAEAGKPTPPRNPPSPSPPGHNQELPRNPHQALEA"
    )
    
    # 预测结合
    result = predictor.predict_binding(
        pd_l1_sequence,
        designed_binder,
        target_id="PD-L1",
        binder_id="Designed_Binder"
    )
    
    print(f"Target: {result.target_id}")
    print(f"Binder: {result.binder_id}")
    print(f"Predicted Affinity: {result.predicted_affinity:.2f} nM")
    print(f"Confidence: {result.confidence:.2%}")
    print(f"Binding Interface: {result.binding_interface_residues[:10]}...")
    print(f"Assessment: {result.notes}")
    
    # 批量筛选示例
    print("\n--- Batch Screening ---")
    candidate_binders = [
        "MKTHIIMLVLLSVLLAQDGVTQDGQWTTQDGILTVSTKTFLSQPFRELAYDTPTSVFTQEV",
        "MKTHIIMLVLLSVLLAQDGVTQDGQWTTQDGILTVSTKTFLSQPFRELAYDTPTSVFTQEVQPEGSPAASQ",
        "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVK",
        # ... 更多候选序列
    ]
    
    hits = predictor.screen_binders(pd_l1_sequence, candidate_binders, threshold_affinity=500.0)
    
    print(f"Found {len(hits)} binders with affinity < 500 nM")
    for i, hit in enumerate(hits[:10]):
        print(f"{i+1}. {hit.binder_id}: {hit.predicted_affinity:.2f} nM")


if __name__ == "__main__":
    main()