import gzip
from pathlib import Path
from typing import Iterator, NamedTuple

import requests
from Bio import SeqIO


SWISSPROT_URL = "https://ftp.uniprot.org/pub/databases/uniprot/current_release/knowledgebase/complete/uniprot_sprot.fasta.gz"


class ProteinRecord(NamedTuple):
    """タンパク質配列のレコード。"""
    accession: str
    sequence: str
    description: str
    go_terms: list[str]


def download_swissprot(dest: Path) -> Path:
    """UniProt Swiss-Prot FASTAをダウンロードします。"""
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists():
        return dest
    with requests.get(SWISSPROT_URL, stream=True, timeout=300) as resp:
        resp.raise_for_status()
        with open(dest, "wb") as f:
            for chunk in resp.iter_content(chunk_size=8192):
                f.write(chunk)
    return dest


def parse_swissprot(fasta_gz: Path, max_len: int = 1024) -> Iterator[ProteinRecord]:
    """FASTAをストリーミングでパースします。max_lenよりも長いシーケンス（ESMコンテキスト制限）はスキップします。"""
    with gzip.open(fasta_gz, "rt") as f:
        for record in SeqIO.parse(f, "fasta"):
            seq = str(record.seq).upper()
            if len(seq) > max_len or len(seq) < 30:
                continue
            # UniProt FASTAの説明からアクセッションを抽出します（例：sp|P12345|GENE_ORG）
            parts = record.id.split("|")
            accession = parts[1] if len(parts) >= 2 else record.id
            yield ProteinRecord(
                accession=accession,
                sequence=seq,
                description=record.description,
                go_terms=[],  # GO用語は別のソースからマージされます。次のステップを参照してください。
            )

import sqlite3
from collections import defaultdict


def build_go_index(gaf_path: Path, db_path: Path) -> None:
    """GAF（GOアノテーションファイル）を、アクセッション→[GO ID]マッピングとしてインデックス化します。"""
    conn = sqlite3.connect(db_path)
    conn.execute("CREATE TABLE IF NOT EXISTS go (accession TEXT, go_id TEXT)")
    conn.execute("CREATE INDEX IF NOT EXISTS idx_acc ON go(accession)")
    with gzip.open(gaf_path, "rt") as f:
        batch = []
        for line in f:
            if line.startswith("!"):
                continue
            fields = line.split("\t")
            if len(fields) < 5:
                continue
            batch.append((fields[1], fields[4]))  # DB_Object_ID, GO_ID
            if len(batch) >= 10000:
                conn.executemany("INSERT INTO go VALUES (?, ?)", batch)
                batch.clear()
        if batch:
            conn.executemany("INSERT INTO go VALUES (?, ?)", batch)
    conn.commit()
    conn.close()


def lookup_go_terms(accession: str, db_path: Path) -> list[str]:
    conn = sqlite3.connect(db_path)
    cursor = conn.execute("SELECT go_id FROM go WHERE accession = ?", (accession,))
    terms = [row[0] for row in cursor]
    conn.close()
    return terms

import torch
from esm.models.esm3 import ESM3
from esm.sdk.api import ESMProtein


class ESM3Embedder:
    """ESM3シーケンス埋め込み抽出ラッパー。"""

    def __init__(self, device: str = "cuda"):
        self.device = device
        self.model = ESM3.from_pretrained("esm3-sm-open-v1").to(device).eval()

    @torch.no_grad()
    def embed(self, sequence: str) -> torch.Tensor:
        """単一のシーケンスに対して、シーケンストラックの最後の隠れ状態を抽出します。

        戻り値：形状(hidden_dim,)のベクトル。1.4Bモデルの場合、hidden_dim=1536。
        """
        protein = ESMProtein(sequence=sequence)
        encoded = self.model.encode(protein)
        # シーケンストラックの最後の隠れ状態を抽出します
        output = self.model.forward(sequence_tokens=encoded.sequence.unsqueeze(0).to(self.device))
        # (1, seq_len, hidden_dim) → 平均プーリング → (hidden_dim,)
        return output.embeddings.squeeze(0).mean(dim=0).cpu()

    def embed_batch(self, sequences: list[str], batch_size: int = 4) -> torch.Tensor:
        """シーケンスのバッチを処理します。バッチサイズを調整して、VRAM制限に合わせます。"""
        embeddings = []
        for i in range(0, len(sequences), batch_size):
            chunk = sequences[i:i + batch_size]
            for seq in chunk:
                embeddings.append(self.embed(seq))
        return torch.stack(embeddings)

from transformers import AutoTokenizer, AutoModel


class ESM2Embedder:
    """ESM2ベースライン。hidden_dim=1280（t33_650M）。"""

    MODEL_ID = "facebook/esm2_t33_650M_UR50D"

    def __init__(self, device: str = "cuda"):
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(self.MODEL_ID)
        self.model = AutoModel.from_pretrained(self.MODEL_ID).to(device).eval()

    @torch.no_grad()
    def embed(self, sequence: str) -> torch.Tensor:
        inputs = self.tokenizer(sequence, return_tensors="pt", truncation=True, max_length=1024).to(self.device)
        outputs = self.model(**inputs)
        # last_hidden_state: (1, seq_len, 1280) → 平均プーリング → (1280,)
        return outputs.last_hidden_state.squeeze(0).mean(dim=0).cpu()

import faiss
import numpy as np


def build_flat_index(embeddings: np.ndarray) -> faiss.Index:
    """精度100％、ベンチマーク用。"""
    d = embeddings.shape[1]
    index = faiss.IndexFlatIP(d)  # 内積（埋め込みの正規化後のコサイン類似度と同等）
    faiss.normalize_L2(embeddings)
    index.add(embeddings)
    return index


def build_ivf_index(embeddings: np.ndarray, nlist: int = 4096) -> faiss.Index:
    """IVF：nlist=sqrt（N）が経験的な最適値です。推奨nlist〜700（シーケンス数50万）。"""
    d = embeddings.shape[1]
    quantizer = faiss.IndexFlatIP(d)
    index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT)
    faiss.normalize_L2(embeddings)
    index.train(embeddings)
    index.add(embeddings)
    return index


def build_hnsw_index(embeddings: np.ndarray, M: int = 32) -> faiss.Index:
    """HNSW：M=16〜64、大きいほど精度↑、メモリ↑。efConstruction=200が標準。"""
    d = embeddings.shape[1]
    index = faiss.IndexHNSWFlat(d, M, faiss.METRIC_INNER_PRODUCT)
    index.hnsw.efConstruction = 200
    faiss.normalize_L2(embeddings)
    index.add(embeddings)
    return index

def search(
    index: faiss.Index,
    query_embedding: np.ndarray,
    k: int = 10,
    ef_search: int | None = None,
    nprobe: int | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """上位kの近傍検索。HNSWのefSearch、IVFのnprobeは、検索中に調整できます。"""
    if ef_search is not None and hasattr(index, "hnsw"):
        index.hnsw.efSearch = ef_search
    if nprobe is not None and hasattr(index, "nprobe"):
        index.nprobe = nprobe
    faiss.normalize_L2(query_embedding.reshape(1, -1))
    distances, indices = index.search(query_embedding.reshape(1, -1), k)
    return distances[0], indices[0]

from collections import Counter


def transfer_go_annotations(
    neighbor_accessions: list[str],
    neighbor_distances: list[float],
    db_path: Path,
    min_votes: int = 2,
    distance_threshold: float = 0.75,
) -> dict[str, float]:
    """近傍のGO用語を、距離重み付きの多数決で集約します。

    戻り値：{go_id: 信頼度スコア}の辞書。
    """
    weighted_votes: Counter = Counter()
    total_weight = 0.0
    for acc, dist in zip(neighbor_accessions, neighbor_distances):
        if dist < distance_threshold:
            continue
        weight = dist  # コサイン類似度自体を重みとして使用します
        total_weight += weight
        for go_id in lookup_go_terms(acc, db_path):
            weighted_votes[go_id] += weight
    if total_weight == 0:
        return {}
    # 少なくともmin_votes回出現したGO用語のみを返します
    return {
        go_id: score / total_weight
        for go_id, score in weighted_votes.items()
        if score >= min_votes * (total_weight / len(neighbor_accessions))
    }

from sklearn.metrics import precision_recall_curve, auc


def evaluate_pr_curve(
    query_pairs: list[tuple[str, str]],  # (query_acc, target_acc)
    ground_truth: dict[str, set[str]],   # query_acc → 実際の相同アクセッションのセット
    method: callable,                     # method(query_acc) → [(target_acc, score)]
) -> dict:
    """PR曲線を使用して、検出性能を定量化します。"""
    y_true, y_score = [], []
    for query_acc, target_acc in query_pairs:
        preds = method(query_acc)  # [(acc, score), ...]
        scored = {acc: score for acc, score in preds}
        y_true.append(1 if target_acc in ground_truth.get(query_acc, set()) else 0)
        y_score.append(scored.get(target_acc, 0.0))
    precision, recall, _ = precision_recall_curve(y_true, y_score)
    return {
        "auprc": auc(recall, precision),
        "recall_at_p90": max(
            (r for p, r in zip(precision, recall) if p >= 0.9), default=0.0
        ),
    }

def full_pipeline(
    query_sequence: str,
    embedder: ESM3Embedder,
    index: faiss.Index,
    accession_map: list[str],  # インデックスの順序 → アクセッション
    go_db: Path,
    k: int = 20,
) -> dict:
    """新しいシーケンスに対して、検索+機能転送を実行します。"""
    query_emb = embedder.embed(query_sequence).numpy()
    distances, indices = search(index, query_emb, k=k, ef_search=128)
    neighbor_accs = [accession_map[i] for i in indices]
    go_predictions = transfer_go_annotations(neighbor_accs, distances.tolist(), go_db)
    return {
        "query_length": len(query_sequence),
        "top_neighbors": list(zip(neighbor_accs, distances.tolist())),
        "predicted_go_terms": go_predictions,
    }

## パフォーマンス、コスト、および既知の失敗事例

### パフォーマンスの参照（公開ベンチマークを使用）

| アプローチ | データセット | Remote Homolog Recall@P=0.9 | Function Transfer F1 | ソース |
|---|---|:---:|:---:|---|
| BLAST (E ≤ 1e-5) | SCOPe 40 twilight | 0.32 | 0.51 | Rives et al., PNAS 2021 [1] |
| PSI-BLAST 3 iter | SCOPe 40 | 0.44 | 0.58 | Rost 1999 (Classic) |
| HHblits | SCOPe 40 | 0.53 | 0.61 | Remmert et al., Nat Methods 2012 |
| ESM2 650M + KNN | CATH 20 | 0.61 | 0.64 | Rao et al., ICML 2021 [2] |
| ESM3 1.4B + KNN (HNSW) | CATH 20 | ~0.67 (推定) | ~0.68 (推定) | EvolutionaryScale 2024 [5] |
| Foldseek (構造ベース) | CATH 20 | 0.78 | — | van Kempen et al., Nat Biotech 2024 [3] |

Foldseek は依然として強力ですが、事前に計算された構造予測（例：AlphaFold）が必要であり、ESM エンベディングは、シーケンスのみを使用してすぐに検索できるという利点があります。

### インデックスによる測定されたパフォーマンス（公式 FAISS ベンチマークに基づく）

| インデックス | 500k ベクトル (d=1536) 検索レイテンシー | 精度 (recall@10) | メモリ |
|---|:---:|:---:|:---:|
| IndexFlatIP | 150〜200ms | 100% | 3.0 GB |
| IndexIVFFlat (nlist=4096, nprobe=32) | 8〜15ms | 96〜98% | 3.1 GB |
| IndexHNSWFlat (M=32, efSearch=128) | 3〜6ms | 98〜99% | 4.5 GB |

**推奨:** リアルタイムサービスには HNSW、バッチ分析には IVF、ベンチマークには Flat。

### 学習者が再現するための推定コスト

- API コスト：0（完全にローカル）。
- エンベディング抽出：RTX 4060 8GB を使用して 500k シーケンスで約 2〜4 時間（1 秒あたり 40〜60 シーケンス）。
- FAISS インデックス構築：HNSW M=32 で約 20〜30 分（CPU）。
- 検索：1 クエリあたり 3〜6ms。

### 3 つの既知の失敗事例（コミュニティ/論文の収集）

1. **長いシーケンス（>1024 AA）の切り捨てによるドメイン情報の損失**
   症状：マルチドメインタンパク質の最初の 1024 AA のみをエンベディングすると、後続のドメインがエンベディングに反映されません。
   原因：ESM のコンテキスト長制限。
   解決策：(a) 各ドメインを個別にエンベディングし、平均または連結する、(b) スライディングウィンドウを使用して複数のエンベディングを生成し、最大プーリングを実行する。
   ソース：HuggingFace ESM Community — GitHub Issues [7]。

2. **IVF インデックスの不十分な nprobe によるリモートホモログの検出漏れ**
   症状：nprobe をデフォルト値（1）に設定すると、リコールが急激に低下します。これは、リモートホモログが別のクラスタにクラスタ化される可能性が高いためです。
   原因：IVF は、量子化器の近くのクラスタのみを訪問します。遠いホモログは、進化系統の違いにより、異なるクラスタに割り当てられます。
   解決策：nprobe を nlist の少なくとも 1%（50〜100）に設定します。本番環境のサービスでは、適応的な nprobe ロジックを推奨します。
   ソース：FAISS Wiki — "Guidelines for choosing an index" [8]。

3. **エンベディングの正規化の欠如によるコサイン類似度と L2 類似度の混同**
   症状：IndexFlatIP をエンベディングを正規化せずに使用すると、内積はコサイン類似度から逸脱し、異なる長さのタンパク質の比較が歪みます。
   原因：ESM エンベディングのノルムは、シーケンスの長さと組成によって異なります。
   解決策：`faiss.normalize_L2()` を呼び出して、インデックスに追加およびクエリを実行する前に、常にエンベディングを正規化します。
   ソース：FAISS FAQ — "Normalization for cosine similarity" [9]。

## 拡張アイデア

- **Chroma Vector DB の統合:** メタデータフィルタリングが必要な場合（例：特定のタクソンまたは特定の GO 用語のみを検索する場合）は、FAISS ではなく Chroma/Qdrant を使用します。
- **ハイブリッド検索:** エンベディング KNN + BLAST HSP スコアをアンサンブル化します。高い精度が求められる場合に有利です。
- **構造エンベディングの組み合わせ:** ESM3 の構造トラック + Foldseek 3Di アルファベットを組み合わせて、デュアルのシーケンス-構造検索を実行します。
- **韓国語 UI サービス:** FastAPI + Streamlit を使用して、社内サービスを構築します。新しいシーケンスを入力すると、システムはすぐに上位 10 件の類似シーケンス + GO 予測を返します。

## 次の章

- Chapter 08 `docking-hybrid-diffusion`: エンベディング検索を使用して候補となるターゲットタンパク質のリストを作成し、DiffDock-Glide で検証します。
- Chapter 11 `structure-affinity-boltz`: エンベディングを使用して類似のタンパク質を見つけ、Boltz-2 で結合親和性を予測します（ドラッグリポジショニングパイプライン）。
- Chapter 12 `single-cell-perturbation`: エンベディング類似性を使用して、シングルセル遺伝子攪乱のターゲット候補を選択します。
- Chapter 13 `protein-design-multimodal`: 本章のシーケンストラックに対応する ESM3 の構造/機能トラックを利用します。
- Chapter 14 `bio-mcp-agent`: 本章の検索パイプラインを、自律エージェントが使用できる MCP ツールとして公開します。

## 参考文献

1. Rives A, Meier J, Sercu T, et al. "2億5000万のタンパク質配列に対して、教師なし学習をスケールアップすることで、生物学的構造と機能が出現する。" PNAS 2021. `https://www.pnas.org/doi/10.1073/pnas.2016239118`
2. Rao R, Meier J, Sercu T, Ovchinnikov S, Rives A. "Transformerタンパク質言語モデルは、教師なしの構造学習を行う。" ICLR 2021. `https://openreview.net/forum?id=fylclEqgvgd`
3. van Kempen M, Kim S, Tumescheit C, et al. "Foldseekを用いた、高速かつ正確なタンパク質構造検索。" Nature Biotechnology 2024. `https://www.nature.com/articles/s41587-023-01773-0`
4. CAFA (Critical Assessment of Function Annotation): `https://www.biofunctionprediction.org/cafa/`
5. Hayes T, Rao R, Akin H, et al. "言語モデルを用いて、5億年間の進化をシミュレーションする。" bioRxiv 2024 (ESM3). `https://www.biorxiv.org/content/10.1101/2024.07.01.600583v1`
6. UniProt-GOA: `https://www.ebi.ac.uk/GOA/downloads`
7. HuggingFace ESMコミュニティディスカッション: `https://huggingface.co/facebook/esm2_t33_650M_UR50D/discussions`
8. FAISS Wiki — インデックスを選択するためのガイドライン: `https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index`
9. FAISS FAQ — 正規化: `https://github.com/facebookresearch/faiss/wiki/FAQ`
10. EvolutionaryScale ESM3 GitHub: `https://github.com/evolutionaryscale/esm`
11. FAISS GitHub: `https://github.com/facebookresearch/faiss`
12. UniProt Swiss-Prot: `https://www.uniprot.org/`
13. Biopython: `https://biopython.org/`
14. SCOPe (Structural Classification of Proteins extended): `https://scop.berkeley.edu/`
15. CATH (Class · Architecture · Topology · Homology) database: `https://www.cathdb.info/`