app/predictor.py

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

from app.config import (
    MODEL_PATH, GO_TERMS_PATH,
    ESM2_MODEL_NAME, ESM2_DIM, HIDDEN_DIM, NUM_LABELS,
    THRESHOLD, WINDOW_SIZE, WINDOW_STRIDE,
)
from app.model import ProteinGNN
from app.schemas import PredictionResponse, ProteinResult, GOTermPrediction
from utils.fasta_parser import parse_fasta, sliding_window


class ProteinPredictor:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.gnn: ProteinGNN = None
        self.tokenizer = None
        self.esm_model = None
        self.go_terms: List[str] = []
        self.is_ready = False

    # ── Load ──────────────────────────────────────────────────────────────────
    def load(self):
        # 1. Load GO terms list  (format: GO:ID \t name \t ia_weight)
        if not GO_TERMS_PATH.exists():
            raise FileNotFoundError(f"GO terms file not found at {GO_TERMS_PATH}")
        with open(GO_TERMS_PATH) as f:
            self.go_terms = [line.strip().split("\t")[0] for line in f if line.strip()]
        print(f"  ✅ GO terms loaded: {len(self.go_terms)}")

        # 2. Load GNN weights
        if not MODEL_PATH.exists():
            raise FileNotFoundError(f"Model file not found at {MODEL_PATH}")
        self.gnn = ProteinGNN(ESM2_DIM, HIDDEN_DIM, NUM_LABELS).to(self.device)
        checkpoint = torch.load(MODEL_PATH, map_location=self.device, weights_only=False)
        checkpoint = torch.load(MODEL_PATH, map_location=self.device, weights_only=False)
        state = (checkpoint.get("model_state_dict")
                         or checkpoint.get("model_state")
                         or checkpoint
                         )
        if any(k.startswith("gnn.") for k in state.keys()):
                state = {k[4:]: v for k, v in state.items() if k.startswith("gnn.")}

                self.gnn.load_state_dict(state)

        self.gnn.eval()
        print(f"  ✅ GNN model loaded from {MODEL_PATH}")

        # 3. Load ESM2 tokenizer + model
        print(f"  🔄 Loading ESM2 ({ESM2_MODEL_NAME}) — this may take a moment...")
        self.tokenizer = AutoTokenizer.from_pretrained(ESM2_MODEL_NAME)
        self.esm_model = AutoModel.from_pretrained(ESM2_MODEL_NAME).to(self.device)
        self.esm_model.eval()
        print("  ✅ ESM2 ready")

        self.is_ready = True

    # ── Public API ────────────────────────────────────────────────────────────
    def predict(self, fasta_text: str) -> PredictionResponse:
        proteins = parse_fasta(fasta_text)
        results = [self._predict_one(pid, seq) for pid, seq in proteins]
        return PredictionResponse(results=results, total_proteins=len(results))

    # ── Internal ──────────────────────────────────────────────────────────────
    def _predict_one(self, protein_id: str, sequence: str) -> ProteinResult:
        # 1. ESM2 embedding (with sliding window for long sequences)
        embedding = self._get_esm2_embedding(sequence)       # (1280,)

        # 2. Inductive GNN inference (classifier only, no graph)
        with torch.no_grad():
            x = embedding.unsqueeze(0).to(self.device)       # (1, 1280)
            logits = self.gnn.forward_inductive(x)            # (1, 4201)
            probs = torch.sigmoid(logits).squeeze(0).cpu().numpy()  # (4201,)

        # 3. Apply threshold
        predicted_indices = np.where(probs >= THRESHOLD)[0]
        go_predictions = [
            GOTermPrediction(
                go_term=self.go_terms[i],
                score=round(float(probs[i]), 4),
            )
            for i in predicted_indices
        ]
        # Sort by confidence descending
        go_predictions.sort(key=lambda x: x.score, reverse=True)

        return ProteinResult(
            protein_id=protein_id,
            sequence_length=len(sequence),
            predicted_go_terms=go_predictions,
            num_predictions=len(go_predictions),
        )

    def _get_esm2_embedding(self, sequence: str) -> torch.Tensor:
        """
        Get ESM2 mean-pooled embedding.
        Uses sliding window + averaging for sequences longer than WINDOW_SIZE.
        """
        windows = sliding_window(sequence, WINDOW_SIZE, WINDOW_STRIDE)
        window_embeddings = []

        for window_seq in windows:
            inputs = self.tokenizer(
                window_seq,
                return_tensors="pt",
                truncation=True,
                max_length=WINDOW_SIZE + 2,   # +2 for [CLS] and [EOS]
                padding=False,
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            with torch.no_grad():
                outputs = self.esm_model(**inputs)

            # Mean-pool over sequence tokens (exclude [CLS] and [EOS])
            hidden = outputs.last_hidden_state[0, 1:-1, :]   # (seq_len, 1280)
            emb = hidden.mean(dim=0)                          # (1280,)
            window_embeddings.append(emb)

        # Average all window embeddings
        final_embedding = torch.stack(window_embeddings).mean(dim=0)
        return final_embedding