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 description에서 accession 추출 (예: 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는 별도 소스에서 병합, 아래 Step 참조
            )

import sqlite3
from collections import defaultdict


def build_go_index(gaf_path: Path, db_path: Path) -> None:
    """GAF (GO Annotation File)을 accession → [GO IDs] 매핑으로 인덱싱."""
    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:
        """서열 하나의 시퀀스 트랙 last hidden state의 평균 pooling.
        
        Returns: (hidden_dim,) 벡터. 1.4B 모델 기준 hidden_dim=1536.
        """
        protein = ESMProtein(sequence=sequence)
        encoded = self.model.encode(protein)
        # 시퀀스 트랙 last hidden state 추출
        output = self.model.forward(sequence_tokens=encoded.sequence.unsqueeze(0).to(self.device))
        # (1, seq_len, hidden_dim) → 평균 pooling → (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 한계에 맞춰 batch_size 조정."""
        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 baseline. 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) → 평균 pooling → (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)  # inner product (임베딩 정규화 후 코사인 유사도와 동치)
    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) 근처가 경험적 최적. 50만 서열이면 nlist~700 권장."""
    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]:
    """top-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 term을 거리 가중 다수결로 집계.
    
    Returns: {go_id: confidence_score} 딕셔너리.
    """
    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 term만 반환
    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 → set of true homolog acc
    method: callable,                     # method(query_acc) → [(target_acc, score)]
) -> dict:
    """PR curve로 검출력 정량화."""
    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],  # index 순서 → accession
    go_db: Path,
    k: int = 20,
) -> dict:
    """새 서열 하나에 대해 검색 + function transfer."""
    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,
    }