import torch
import torch.nn as nn
from typing import List, Dict

class AlphaFoldEmbedTokenizer:
    def __init__(self):
        self.vocab = {aa: i for i, aa in enumerate("ACDEFGHIKLMNPQRSTVWY")}
        self.vocab["<UNK>"] = len(self.vocab)
        
    def tokenize(self, seq: str) -> List[int]:
        return [self.vocab.get(aa, self.vocab["<UNK>"]) for aa in seq]

class AlphaFoldEmbedInference(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Embedding(25, 512)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True),
            num_layers=4
        )
        self.plddt_head = nn.Linear(512, 1)
        self.domain_head = nn.Linear(512, 1)

    def forward(self, input_ids):
        x = self.embed(input_ids)
        x = self.transformer(x)
        plddt = torch.sigmoid(self.plddt_head(x)).squeeze(-1) * 100 # 0-100 scale
        domain_boundaries = torch.sigmoid(self.domain_head(x)).squeeze(-1)
        return {"embeddings": x, "plddt": plddt, "domain_boundaries": domain_boundaries}

def process_proteins(proteins: List[str]):
    tokenizer = AlphaFoldEmbedTokenizer()
    model = AlphaFoldEmbedInference()
    model.eval()
    
    results = []
    with torch.no_grad():
        for seq in proteins:
            tokens = tokenizer.tokenize(seq)
            input_ids = torch.tensor([tokens])
            out = model(input_ids)
            results.append(out)
    return results

if __name__ == "__main__":
    res = process_proteins(["MKALIVLGLVLLSVTVQG", "MTEYKLVVVGAGGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGETCLLDILDTAGQ"])
    print("Done processing proteins.")