import json, re, warnings
from pathlib import Path
from collections import Counter

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import streamlit as st
import matplotlib.pyplot as plt
import requests as req
import streamlit.components.v1 as components
from transformers import EsmModel, EsmTokenizer

try:
    import py3Dmol
    STMOL_AVAILABLE = True
except ImportError:
    STMOL_AVAILABLE = False

warnings.filterwarnings("ignore")

st.set_page_config(page_title="ARG Detection", page_icon="🧬", layout="wide")

DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"
TOP_K       = 3
MAX_LEN     = 1024
MIN_LEN     = 30
N_BIO       = 26
AA_LIST     = sorted("ACDEFGHIKLMNPQRSTVWY")
STANDARD_AA = set(AA_LIST)
MAX_BATCH = 500  
GENERIC_ORGANISM = [
    "bacteria, viruses, fungi, and other genome sequence "
    "associated with antimicrobial resistance"
]

EXAMPLE_ARG = (
    "MSIQHFRVALIPFFAAFCLPVFAHPETLVKVKDAEDQLGARVGYIELDLNSGKILESFRPEERFPMMSTFKVLL"
    "CGAVLSRVDAGQEQLGRRIHYSQNDLVEYSPVTEKHLTDGMTVRELCSAAITMSDNTAANLLLTTIGGPKELTA"
    "FLHNMGDHVTRLDRWEPELNEAIPNDERDTTMPVAMATTLRKLLTGELLTLASRQQLIDWMEADKVAGPLLRSAL"
    "PAGWFIADKSGAGERGSRGIIAALGPDGKPSRIVVIYTTGSQATMDERNRQIAEIGASLIKHW"
)
EXAMPLE_NON = (
    "MSENIKQRALDSLTQAQLDELEKHIEHEMARLQRLEQQLQEEDDEDEDDEDGAAAATQAQIDMSEGVTSTIIPP"
    "HNLEELPKDGKTVFDSALKNAGLEVHFQGLQDQLRQLKEQNQELKKQLQQAKEDQERQKPEEDEEVEEDEDDDD"
    "EDDDEEEDEEDEDDEDEDDAAPRRNVRQLQALQQAQAENLEKRNAKLAAELNQTIDDLENAIQKLKEQMQKMNEE"
    "IDQLQRDNQTLESELNQLKNELNELNRDLNEKLNELQEQIKQLQSELDTLRKENA"
)

_REGION_LABEL = {
    "en": {
        "N-terminal (0-33%)":   "N-terminal (0–33%)",
        "Mid (33-66%)":         "Mid (33–66%)",
        "C-terminal (66-100%)": "C-terminal (66–100%)",
    },
    "id": {
        "N-terminal (0-33%)":   "N-terminal (0–33%)",
        "Mid (33-66%)":         "Tengah (33–66%)",
        "C-terminal (66-100%)": "C-terminal (66–100%)",
    },
}

# Bilingual dic
def _load_texts() -> dict:
    try:
        with open("lang.json", encoding="utf-8") as f:
            return json.load(f)
    except FileNotFoundError:
        import os
        here = os.path.dirname(os.path.abspath(__file__))
        with open(os.path.join(here, "lang.json"), encoding="utf-8") as f:
            return json.load(f)

TEXTS = _load_texts()

if "lang" not in st.session_state:
    st.session_state["lang"] = "en"
T = TEXTS[st.session_state["lang"]]

st.title(T["page_title"])
st.caption(T["page_caption"])
st.markdown("---")


# Model archi
class CNNAttention(nn.Module):
    def __init__(self, embed_dim=320, proj_dim=128, n_filters=128,
                 kernel_size=5, n_bio=26, dropout=0.3):
        super(CNNAttention, self).__init__()
        pad = kernel_size // 2
        self.projection = nn.Sequential(
            nn.Linear(embed_dim, proj_dim),
            nn.LayerNorm(proj_dim),
            nn.ReLU()
        )
        self.conv1 = nn.Sequential(
            nn.Conv1d(proj_dim, n_filters, kernel_size, padding=pad),
            nn.BatchNorm1d(n_filters),
            nn.ReLU(),
            nn.Dropout(dropout * 0.5)
        )
        self.conv2 = nn.Sequential(
            nn.Conv1d(n_filters, n_filters, kernel_size, padding=pad),
            nn.BatchNorm1d(n_filters),
            nn.ReLU()
        )
        self.attn_W = nn.Linear(n_filters, n_filters)
        self.attn_v = nn.Linear(n_filters, 1, bias=False)
        self.classifier = nn.Sequential(
            nn.Linear(n_filters + n_bio, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(dropout * 0.5),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

    def forward(self, x, features, lengths=None, return_attention=False):
        B, L, _ = x.shape
        x = self.projection(x).transpose(1, 2)
        x = self.conv2(self.conv1(x)).transpose(1, 2)
        scores = self.attn_v(torch.tanh(self.attn_W(x)))
        if lengths is not None:
            mask = torch.arange(L, device=x.device).unsqueeze(0) >= lengths.unsqueeze(1)
            scores = scores.masked_fill(mask.unsqueeze(-1), float("-inf"))
        alpha   = torch.softmax(scores, dim=1)
        context = (alpha * x).sum(dim=1)
        pred    = self.classifier(torch.cat([context, features], dim=1)).squeeze(-1)
        if return_attention:
            return {"pred_prob": pred, "attention": alpha.squeeze(-1)}
        return pred


# Cached loader
@st.cache_resource
def load_esm2():
    try:
        tok = EsmTokenizer.from_pretrained("facebook/esm2_t6_8M_UR50D")
        mdl = EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D").to(DEVICE).eval()
        return tok, mdl
    except Exception as e:
        st.error(f"Failed to load ESM-2: {e}")
        return None, None


@st.cache_resource
def load_cnn_model():
    path = Path("models/best_attention.pt")
    if not path.exists():
        st.error("Model not found: models/best_attention.pt")
        return None
    ckpt = torch.load(path, map_location=DEVICE)
    cfg  = ckpt["config"]
    m = CNNAttention(
        embed_dim=cfg.get("embed_dim", 320),
        proj_dim=cfg["proj_dim"],
        n_filters=cfg["n_filters"],
        kernel_size=cfg["kernel_size"],
        n_bio=cfg.get("n_bio", 26),
        dropout=cfg["dropout"],
    ).to(DEVICE)
    m.load_state_dict(ckpt["state_dict"])
    m.eval()
    return m


@st.cache_resource
def load_card_bank():
    try:
        emb  = np.load("data/embeddings/card_bank_global.npy")
        meta = pd.read_csv("data/embeddings/card_bank_meta.csv")
        with open("data/embeddings/card_bank_accessions.json") as f:
            accs = json.load(f)
        return emb, meta, accs
    except FileNotFoundError as e:
        st.error(f"CARD bank files not found: {e}. Check data/embeddings/")
        st.stop()


@st.cache_resource
def load_scaler():
    try:
        with open("models/scaler_params.json") as f:
            sp = json.load(f)
    except FileNotFoundError:
        st.error("Scaler file not found: models/scaler_params.json")
        st.stop()
    from sklearn.preprocessing import StandardScaler
    s = StandardScaler()
    s.mean_          = np.array(sp["mean"])
    s.scale_         = np.array(sp["scale"])
    s.var_           = s.scale_ ** 2
    s.n_features_in_ = len(sp["feature_names"])
    return s


# Sequence utility func
def validate_sequence(seq):
    seq     = seq.upper().strip()
    seq     = re.sub(r"[^A-Z]", "", seq)
    invalid = [c for c in seq if c not in STANDARD_AA]
    clean   = "".join(c for c in seq if c in STANDARD_AA)
    return clean, invalid


def truncate_sequence(seq):
    if len(seq) > MAX_LEN:
        half = MAX_LEN // 2
        return seq[:half] + seq[-half:], True
    return seq, False


def parse_fasta(content):
    seqs, cur = [], {"header": "", "sequence": ""}
    for line in content.split("\n"):
        line = line.strip()
        if line.startswith(">"):
            if cur["sequence"]:
                seqs.append(cur)
            cur = {"header": line[1:], "sequence": ""}
        elif line:
            cur["sequence"] += line
    if cur["sequence"]:
        seqs.append(cur)
    return seqs


def parse_fasta_header(header):
    """
    Extract accession, protein_name, and organism from a FASTA header
    """
    meta = {
        "accession":    "",
        "protein_name": "",
        "organism":     "",
        "raw_header":   header,
    }

    # CARD: gb|AAT45742.1|ARO:3000988|TEM-126 [Escherichia coli]
    card = re.match(
        r'^(?:gb|emb|ref|dbj|pir|prf|pdb|lcl)\|([^|]+)\|ARO:[^|]+\|(.+?)'
        r'\s*(?:\[([^\]]+)\])?\s*$',
        header
    )
    if card:
        meta["accession"]    = card.group(1)
        meta["protein_name"] = card.group(2).strip()
        meta["organism"]     = card.group(3).strip() if card.group(3) else ""
        return meta

    # UniProt with OS=: sp|A0A0H2VG78|GLCP_STAES Protein OS=Staphylococcus
    uniprot_os = re.match(
        r'^(?:sp|tr)\|([^|]+)\|\S+\s+(.+?)\s+OS=(.+?)'
        r'(?:\s+OX=|\s+GN=|\s+PE=|\s*$)',
        header
    )
    if uniprot_os:
        meta["accession"]    = uniprot_os.group(1)
        meta["protein_name"] = uniprot_os.group(2).strip()
        meta["organism"]     = uniprot_os.group(3).strip()
        return meta

    # UniProt with bracket: sp|P00811|BLAT_ECOLI Beta-lactamase [Escherichia coli]
    uniprot = re.match(
        r'^(?:sp|tr)\|([^|]+)\|\S+\s+(.+?)\s*(?:\[([^\]]+)\])?\s*$',
        header
    )
    if uniprot:
        meta["accession"]    = uniprot.group(1)
        meta["protein_name"] = uniprot.group(2).strip()
        meta["organism"]     = uniprot.group(3).strip() if uniprot.group(3) else ""
        return meta

    # NCBI with organism: WP_000027057.1 TEM-1 beta-lactamase [Escherichia coli]
    ncbi = re.match(r'^(\S+)\s+(.+?)\s*\[([^\]]+)\]', header)
    if ncbi:
        meta["accession"]    = ncbi.group(1)
        meta["protein_name"] = ncbi.group(2).strip()
        meta["organism"]     = ncbi.group(3).strip()
        return meta

    # NCBI without organism: WP_000027057.1 TEM-1 beta-lactamase
    ncbi_no_org = re.match(r'^(\S+)\s+(.+)', header)
    if ncbi_no_org:
        meta["accession"]    = ncbi_no_org.group(1)
        meta["protein_name"] = ncbi_no_org.group(2).strip()
        return meta

    meta["accession"] = header.strip()
    return meta


def parse_genbank(content):
    """Parse GenBank flat-file format without BioPython."""
    records = []
    raw_records = re.split(r'\n//\s*\n?', content.strip())

    for raw in raw_records:
        if not raw.strip():
            continue

        accession = protein_name = organism = sequence = ""

        ver = re.search(r'^VERSION\s+(\S+)', raw, re.MULTILINE)
        if ver:
            accession = ver.group(1)
        else:
            acc = re.search(r'^ACCESSION\s+(\S+)', raw, re.MULTILINE)
            if acc:
                accession = acc.group(1)

        prod = re.search(r'/product="([^"]+)"', raw)
        if prod:
            protein_name = prod.group(1).strip()
        else:
            defn = re.search(
                r'^DEFINITION\s+(.+?)(?=\n[A-Z]|\nFEATURES)',
                raw, re.DOTALL | re.MULTILINE
            )
            if defn:
                name = re.sub(r'\s+', ' ', defn.group(1)).strip()
                protein_name = re.sub(r'\s*\[[^\]]+\]\.?\s*$', '', name).strip()

        org_qual = re.search(r'/organism="([^"]+)"', raw)
        if org_qual:
            organism = org_qual.group(1).strip()
        else:
            org_field = re.search(r'^\s{2}ORGANISM\s+(.+)', raw, re.MULTILINE)
            if org_field:
                organism = org_field.group(1).strip()

        origin = re.search(
            r'^ORIGIN\s*\n(.*?)(?=\n//|\Z)',
            raw, re.DOTALL | re.MULTILINE
        )
        if origin:
            sequence = re.sub(r'[\d\s]', '', origin.group(1)).upper()

        if not sequence or len(sequence) < 10:
            continue

        header = accession
        if protein_name:
            header += f" {protein_name}"
        if organism:
            header += f" [{organism}]"

        records.append({"header": header[:100], "sequence": sequence})

    return records


def detect_file_format(content, filename=""):
    """Detect whether a file is FASTA or GenBank by extension then content"""
    fname = filename.lower()
    if fname.endswith(('.gb', '.gbk', '.genbank')):
        return 'genbank'
    if fname.endswith(('.fasta', '.fa', '.faa', '.fna', '.ffn', '.frn')):
        return 'fasta'
    first = content[:500]
    if re.search(r'^LOCUS\s+\S+', first, re.MULTILINE):
        return 'genbank'
    if re.search(r'^ACCESSION\s+\S+', first, re.MULTILINE):
        return 'genbank'
    return 'fasta'


def compute_26_features(seq):
    n = len(seq)
    if n == 0:
        return np.zeros(N_BIO)
    aa_count = {a: 0 for a in AA_LIST}
    for c in seq:
        if c in aa_count:
            aa_count[c] += 1
    feats = [n] + [aa_count[a] / n for a in AA_LIST]
    feats += [
        sum(c in set("AVILPFMW") for c in seq) / n,
        sum(c in set("KRH")      for c in seq) / n,
        sum(c in set("DE")       for c in seq) / n,
        sum(c in set("FWY")      for c in seq) / n,
        sum(c in set("STNQ")     for c in seq) / n,
    ]
    return np.array(feats)


@torch.no_grad()
def embed_sequence(seq, tokenizer, esm_model):
    enc = tokenizer([seq], return_tensors="pt", padding=True,
                    truncation=True, max_length=MAX_LEN + 2)
    enc = {k: v.to(DEVICE) for k, v in enc.items()}
    out = esm_model(**enc)
    h   = out.last_hidden_state[:, 1:-1, :]
    L   = min(len(seq), MAX_LEN)
    local      = h[0, :L, :].cpu().numpy()
    global_emb = local.mean(axis=0)
    return local, global_emb


def cosine_search(query_emb, bank_emb, top_k=3):
    if bank_emb.shape[0] == 0:
        return np.array([]), np.array([])
    q_n     = query_emb / (np.linalg.norm(query_emb) + 1e-10)
    b_n     = bank_emb / (np.linalg.norm(bank_emb, axis=1, keepdims=True) + 1e-10)
    sims    = b_n @ q_n
    top_k   = min(top_k, len(sims))
    top_idx = np.argsort(sims)[::-1][:top_k]
    return top_idx, sims[top_idx]


def clean_organism(org):
    if not org or str(org).strip() in ("", "nan", "—"):
        return "Unspecified"
    for g in GENERIC_ORGANISM:
        if g in str(org).lower():
            return "Unspecified"
    return str(org).strip()


def get_similarity_rows(top_idx, top_scores, bank_accs, df_meta):
    rows = []
    for rank, (idx, score) in enumerate(zip(top_idx, top_scores), 1):
        if idx >= len(bank_accs):
            continue
        acc  = bank_accs[idx]
        meta = df_meta[df_meta["accession"] == acc]
        m    = meta.iloc[0].to_dict() if len(meta) > 0 else {}
        rows.append({
            "Rank":       rank,
            "Accession":  acc,
            "Similarity": f"{score:.4f}",
            "Short Name": str(m.get("card_short_name", "—") or "—"),
            "Drug Class": str(m.get("drug_class", "—") or "—")[:50],
            "Mechanism":  str(m.get("resistance_mechanism", "—") or "—")[:50],
            "Organism":   clean_organism(m.get("organism", "")),
        })
    return rows


def interpret_attention(seq, w):
    """
    Divide sequence into three equal thirds and sum attention per region
    """
    L  = len(seq)
    n3 = max(1, L // 3)
    regions = {
        "N-terminal (0-33%)":   float(w[:n3].sum()),
        "Mid (33-66%)":         float(w[n3:2*n3].sum()),
        "C-terminal (66-100%)": float(w[2*n3:].sum()),
    }
    dominant     = max(regions, key=regions.get)
    dom_pct      = regions[dominant] * 100
    top5         = sorted(range(len(w)), key=lambda i: w[i], reverse=True)[:5]
    top5         = sorted(top5)
    top_residues = [f"{seq[p]}{p+1}" for p in top5 if p < len(seq)]
    return {
        "dominant":     dominant,
        "dom_pct":      dom_pct,
        "top_residues": top_residues,
        "top5_idx":     top5,
        "regions":      regions,
    }


def predict_one(seq_raw, tokenizer, esm_model, cnn_model,
                bank_emb, bank_accs, df_meta, scaler):
    try:
        seq, was_trunc = truncate_sequence(seq_raw)
        local_emb, global_emb = embed_sequence(seq, tokenizer, esm_model)
        bio_norm = scaler.transform(compute_26_features(seq).reshape(1, -1))[0]
        x_t = torch.tensor(local_emb, dtype=torch.float32).unsqueeze(0).to(DEVICE)
        f_t = torch.tensor(bio_norm,  dtype=torch.float32).unsqueeze(0).to(DEVICE)
        L_t = torch.tensor([len(seq)]).to(DEVICE)
        with torch.no_grad():
            out = cnn_model(x_t, f_t, lengths=L_t, return_attention=True)
        prob = float(out["pred_prob"].item())
        attn = out["attention"].squeeze(0).cpu().numpy()
        if attn.sum() > 0:
            attn = attn / attn.sum()
        top_idx, top_scores = cosine_search(global_emb, bank_emb, top_k=TOP_K)
        rows = get_similarity_rows(top_idx, top_scores, bank_accs, df_meta)
        return seq, was_trunc, prob, attn, rows
    except RuntimeError as e:
        # Catches CUDA OOM and other torch errors
        raise RuntimeError(f"Prediction failed (possibly out of memory): {e}") from e


# Display utility func
def plot_attention(seq, attn, max_display=200):
    L = min(len(seq), len(attn), max_display)
    w = attn[:L].copy()
    if w.sum() > 0:
        w = w / w.sum()
    fig, ax = plt.subplots(figsize=(15, 3))
    colors = plt.cm.Reds(w / (w.max() + 1e-10))
    ax.bar(range(L), w, color=colors, width=1.0)
    ax.set_xlabel(T["attn_xlabel"].format(L=L, total=len(seq)))
    ax.set_ylabel(T["attn_ylabel"])
    ax.set_title(T["attn_chart_title"])
    plt.tight_layout()
    return fig, w


def plot_similarity_bar(rows):
    names  = [r["Short Name"][:20] for r in rows]
    scores = [float(r["Similarity"]) for r in rows]
    fig, ax = plt.subplots(figsize=(6, 3))
    colors = ["#DC2626" if s >= 0.9 else "#F59E0B" if s >= 0.7 else "#1C7293"
              for s in scores]
    bars = ax.barh(names, scores, color=colors, alpha=0.85)
    ax.set_xlim(0, 1.05)
    ax.set_xlabel(T["sim_xlabel"])
    ax.set_title(T["sim_chart_title"])
    ax.axvline(0.7, color="gray", ls="--", lw=1, alpha=0.5)
    for bar, score in zip(bars, scores):
        ax.text(score + 0.01, bar.get_y() + bar.get_height() / 2,
                f"{score:.4f}", va="center", fontsize=9)
    plt.tight_layout()
    return fig


def show_sequence_stats(seq, bio_raw):
    st.markdown(T["stats_header"])
    stats = {
        T["stat_hydrophobic"]: bio_raw[21],
        T["stat_positive"]:    bio_raw[22],
        T["stat_negative"]:    bio_raw[23],
        T["stat_aromatic"]:    bio_raw[24],
        T["stat_polar"]:       bio_raw[25],
    }
    for name, val in stats.items():
        st.write(f"- **{name}:** {val:.3f}")
    aa_vals  = {aa: bio_raw[i + 1] for i, aa in enumerate(AA_LIST)}
    top3_aa  = sorted(aa_vals, key=aa_vals.get, reverse=True)[:3]
    top3_str = ", ".join(f"{aa}({aa_vals[aa]:.2f})" for aa in top3_aa)
    st.write(f"- **{T['stat_top3_aa']}:** {top3_str}")


def fetch_esmfold_structure(seq, timeout=120):
    """Predict 3D structure via ESMFold API"""
    try:
        resp = req.post(
            "https://api.esmatlas.com/foldSequence/v1/pdb/",
            data=seq,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            timeout=60
        )
        if resp.status_code == 200 and len(resp.text.strip()) > 100:
            return resp.text
    except Exception:
        pass

    try:
        resp = req.post(
            "https://api-inference.huggingface.co/models/facebook/esmfold_v1",
            json={"inputs": seq},
            headers={"Content-Type": "application/json"},
            timeout=120
        )
        if resp.status_code == 200 and len(resp.text.strip()) > 100:
            return resp.text
    except Exception:
        pass

    return None


def render_3d_structure(pdb_str, top5_idx, seq):
    view = py3Dmol.view(width="100%", height=550)
    view.addModel(pdb_str, "pdb")
    view.setStyle({"cartoon": {"color": "spectrum", "thickness": 0.4, "opacity": 0.85}})

    colors = ["#DC2626", "#B91C1C", "#991B1B", "#7F1D1D", "#450A0A"]
    for rank, pos in enumerate(top5_idx):
        resi  = pos + 1
        aa    = seq[pos] if pos < len(seq) else "?"
        color = colors[rank] if rank < len(colors) else "#DC2626"
        view.addStyle(
            {"resi": resi, "atom": "CA"},
            {"sphere": {"color": color, "radius": 1.8}}
        )
        view.addLabel(
            f"{aa}{resi}",
            {
                "backgroundColor": color, "backgroundOpacity": 0.9,
                "fontColor": "white", "fontSize": 13, "fontStyle": "bold",
                "inFront": True, "showBackground": True, "fixed": False,
            },
            {"resi": resi, "atom": "CA"}
        )

    for label, color, resi in [
        ("N-term", "#1C7293", 1),
        ("C-term", "#374151", len(seq)),
    ]:
        view.addLabel(
            label,
            {
                "backgroundColor": color, "backgroundOpacity": 0.9,
                "fontColor": "white", "fontSize": 11,
                "inFront": True, "showBackground": True,
            },
            {"resi": resi, "atom": "CA"}
        )

    view.setBackgroundColor("white")
    view.zoomTo()
    view.spin(False)
    return view


# Sidebar
with st.sidebar:
    # Lang selector
    lang_options  = ["English", "Bahasa Indonesia"]
    current_index = 0 if st.session_state["lang"] == "en" else 1
    selected_lang = st.radio(
        T["sidebar_lang_label"],
        lang_options,
        index=current_index,
        horizontal=True,
    )
    new_lang = "en" if selected_lang == "English" else "id"
    if new_lang != st.session_state["lang"]:
        st.session_state["lang"] = new_lang
        st.rerun()

    st.markdown("---")
    st.header(T["sidebar_header"])
    st.markdown(T["sidebar_pipeline"])
    st.markdown("---")
    show_3d = st.checkbox(T["sidebar_3d_toggle"], value=True)
    st.caption(T["sidebar_3d_caption"])


# Input
mode = st.radio(
    T["mode_label"],
    [T["mode_manual"], T["mode_upload"]],
    horizontal=True,
)
sequences_to_predict = []

if mode == T["mode_manual"]:
    col_ex1, col_ex2, _ = st.columns([1, 1, 6])
    with col_ex1:
        if st.button(T["btn_example_arg"]):
            st.session_state["example_seq"] = EXAMPLE_ARG
    with col_ex2:
        if st.button(T["btn_example_non"]):
            st.session_state["example_seq"] = EXAMPLE_NON

    raw = st.text_area(
        T["textarea_label"],
        value=st.session_state.get("example_seq", ""),
        height=120,
        placeholder=T["textarea_placeholder"],
    )
    if raw:
        clean, invalid = validate_sequence(raw)
        if invalid:
            st.warning(T["warning_invalid_chars"].format(n=len(set(invalid))))
        if len(clean) < MIN_LEN:
            st.error(T["error_too_short"].format(min=MIN_LEN))
        else:
            sequences_to_predict.append({
                "header":   T["input_manual_header"],
                "sequence": clean,
            })

else:
    uploaded = st.file_uploader(
        T["upload_label"],
        type=["fasta", "fa", "faa", "txt", "gb", "gbk"],
        help=T["upload_help"],
    )
    if uploaded:
        raw_bytes = uploaded.read()
        try:
            content = raw_bytes.decode("utf-8")
        except UnicodeDecodeError:
            content = raw_bytes.decode("latin-1")

        file_format = detect_file_format(content, uploaded.name)

        if file_format == "genbank":
            parsed       = parse_genbank(content)
            format_label = "GenBank"
        else:
            raw_parsed   = parse_fasta(content)
            parsed       = [{"header": p["header"], "sequence": p["sequence"]}
                            for p in raw_parsed]
            format_label = "FASTA"

        n_invalid = 0
        for p in parsed:
            clean, _ = validate_sequence(p["sequence"])
            if len(clean) >= MIN_LEN:
                sequences_to_predict.append({
                    "header":   p["header"][:80],
                    "sequence": clean,
                })
            else:
                n_invalid += 1

        n_valid = len(sequences_to_predict)
        if n_valid > MAX_BATCH:
            st.warning(T["batch_cap_warning"].format(n=n_valid, cap=MAX_BATCH))
            sequences_to_predict = sequences_to_predict[:MAX_BATCH]
            n_valid = MAX_BATCH
        if n_valid > 0:
            msg = T["upload_success"].format(n=n_valid, fmt=format_label)
            if n_invalid > 0:
                msg += T["upload_skipped"].format(skip=n_invalid, min=MIN_LEN)
            st.success(msg)
            st.caption(T["upload_caption"].format(fmt=format_label, name=uploaded.name))
        else:
            st.error(T["upload_error"].format(fmt=format_label, min=MIN_LEN))


# Predict
if st.button(T["btn_predict"], type="primary", disabled=not sequences_to_predict):
    tokenizer, esm_model = load_esm2()
    if tokenizer is None:
        st.stop()
    cnn_model = load_cnn_model()
    if cnn_model is None:
        st.stop()
    bank_emb, df_meta, bank_accs = load_card_bank()
    scaler = load_scaler()

    lang_key    = st.session_state["lang"]
    is_batch    = len(sequences_to_predict) > 1
    all_results = []

    if is_batch:
        st.markdown("---")
        st.subheader(T["processing_header"].format(n=len(sequences_to_predict)))
        progress_bar = st.progress(0)
        status_text  = st.empty()

    for i, seq_data in enumerate(sequences_to_predict):
        if is_batch:
            status_text.text(
                T["processing_status"].format(
                    i=i + 1,
                    n=len(sequences_to_predict),
                    header=seq_data["header"][:40],
                )
            )

        try:
            seq, was_trunc, prob, attn, rows = predict_one(
                seq_data["sequence"], tokenizer, esm_model,
                cnn_model, bank_emb, bank_accs, df_meta, scaler
            )
        except Exception as e:
            if is_batch:
                all_results.append({
                    "Header":           seq_data["header"],
                    "Accession":        "", "Protein Name": "",
                    "Source Organism":  "", "Length (aa)": len(seq_data["sequence"]),
                    "Classification":   "ERROR",
                    "Confidence":       "—",
                    "Hydrophobic": "—", "Positive Charge": "—",
                    "Dominant Region": "—", "Top-5 Positions": "—",
                    "Top-1 Match": "—", "Top-1 Similarity": "—",
                    "Drug Class": "—", "Mechanism": "—", "Organism": "—",
                })
                progress_bar.progress((i + 1) / len(sequences_to_predict))
            else:
                st.error(f"Prediction error: {e}")
            continue

        is_resistant       = prob > 0.5
        label_internal     = "RESISTANT" if is_resistant else "NON-RESISTANT"
        label_display      = T["label_resistant"] if is_resistant else T["label_non_resistant"]
        confidence_display = prob if is_resistant else 1 - prob
        interp             = interpret_attention(seq, attn)
        bio_raw            = compute_26_features(seq)
        region_labels      = _REGION_LABEL[lang_key]

        # Batch mode
        if is_batch:
            fasta_meta_b = parse_fasta_header(seq_data["header"])
            all_results.append({
                "Header":           seq_data["header"],
                "Accession":        fasta_meta_b["accession"],
                "Protein Name":     fasta_meta_b["protein_name"],
                "Source Organism":  fasta_meta_b["organism"],
                "Length (aa)":      len(seq),
                "Classification":   label_internal,   # English for CSV
                "Confidence":       f"{confidence_display:.4f}",
                "Hydrophobic":      f"{bio_raw[21]:.3f}",
                "Positive Charge":  f"{bio_raw[22]:.3f}",
                "Dominant Region":  interp["dominant"],
                "Top-5 Positions":  ", ".join(interp["top_residues"]),
                "Top-1 Match":      rows[0]["Short Name"] if rows else "",
                "Top-1 Similarity": rows[0]["Similarity"] if rows else "",
                "Drug Class":       rows[0]["Drug Class"] if rows else "",
                "Mechanism":        rows[0]["Mechanism"] if rows else "",
                "Organism":         rows[0]["Organism"] if rows else "",
            })
            progress_bar.progress((i + 1) / len(sequences_to_predict))

        # Single mode
        else:
            st.markdown("---")
            st.subheader(seq_data["header"])

            if was_trunc:
                st.info(T["truncated_msg"].format(
                    orig=len(seq_data["sequence"]), new=len(seq)
                ))

            col1, col2 = st.columns([1, 2])

            with col1:
                st.metric(T["metric_classification"], label_display)
                st.metric(T["metric_confidence"],     f"{confidence_display:.4f}")
                st.metric(T["metric_length"],         f"{len(seq)} aa")
                if is_resistant:
                    st.error(T["result_resistant_msg"])
                else:
                    st.success(T["result_non_resistant_msg"])
                show_sequence_stats(seq, bio_raw)

            with col2:
                fig, _ = plot_attention(seq, attn)
                st.pyplot(fig)
                plt.close()

                st.markdown(T["attn_interp_header"])
                dominant_display = region_labels.get(
                    interp["dominant"], interp["dominant"]
                )
                st.markdown(T["attn_dominant_text"].format(
                    region=dominant_display,
                    pct=interp["dom_pct"],
                ))
                st.markdown(T["attn_top5_text"].format(
                    positions=", ".join(interp["top_residues"])
                ))
                for region_key, val in interp["regions"].items():
                    region_disp = region_labels.get(region_key, region_key)
                    col_l, col_r = st.columns([2, 3])
                    with col_l:
                        st.write(f"{region_disp}: **{val * 100:.1f}%**")
                    with col_r:
                        st.progress(float(val))
                st.caption(T["attn_caption"])

            # 3D Visualization
            if show_3d:
                st.markdown("---")
                st.markdown(T["3d_section_header"])
                st.caption(T["3d_caption"])
                if not STMOL_AVAILABLE:
                    st.warning(T["3d_no_pkg"])
                elif len(seq) > 400:
                    st.info(T["3d_too_long"].format(n=len(seq)))
                else:
                    with st.spinner(T["3d_spinner"]):
                        pdb_str = fetch_esmfold_structure(seq)
                    if pdb_str is None:
                        st.warning(T["3d_no_response"])
                    else:
                        view      = render_3d_structure(pdb_str, interp["top5_idx"], seq)
                        view_html = view._make_html()
                        components.html(view_html, height=570, scrolling=False)
                        st.caption(T["3d_pos_caption"].format(
                            positions=", ".join(interp["top_residues"])
                        ))

            # Similarity results
            if rows:
                st.markdown(T["sim_header"].format(k=TOP_K))
                col_tbl, col_chart = st.columns([2, 1])
                with col_tbl:
                    with st.expander(T["sim_expander"], expanded=True):
                        st.dataframe(
                            pd.DataFrame(rows),
                            use_container_width=True,
                            hide_index=True,
                        )
                with col_chart:
                    fig_sim = plot_similarity_bar(rows)
                    st.pyplot(fig_sim)
                    plt.close()

                if is_resistant:
                    all_drugs = []
                    for r in rows:
                        for d in r["Drug Class"].split(";"):
                            d = d.strip()
                            if d and d != "—":
                                all_drugs.append(d)
                    if all_drugs:
                        st.markdown(T["drug_header"])
                        for drug, cnt in Counter(all_drugs).most_common(3):
                            st.markdown(f"- **{drug}** ({cnt} hits)")

            # Download button
            W = 18   # column width for labels
            stat_lines = [
                f"  {T['stat_hydrophobic'].ljust(W)}: {bio_raw[21]:.3f}",
                f"  {T['stat_positive'].ljust(W)}: {bio_raw[22]:.3f}",
                f"  {T['stat_negative'].ljust(W)}: {bio_raw[23]:.3f}",
                f"  {T['stat_aromatic'].ljust(W)}: {bio_raw[24]:.3f}",
                f"  {T['stat_polar'].ljust(W)}: {bio_raw[25]:.3f}",
            ]
            result_lines = [
                T["dl_title"],
                T["dl_divider"],
                f"{T['dl_header_lbl'].ljust(W)}: {seq_data['header']}",
                f"{T['dl_length_lbl'].ljust(W)}: {len(seq)} aa",
                f"{T['dl_class_lbl'].ljust(W)}: {label_display}",
                f"{T['dl_conf_lbl'].ljust(W)}: {confidence_display:.4f}",
                "",
                T["dl_stats_section"],
                *stat_lines,
                "",
                T["dl_attn_section"],
                f"  {T['dl_dominant'].ljust(W)}: {dominant_display} ({interp['dom_pct']:.1f}%)",
                f"  {T['dl_top5'].ljust(W)}: {', '.join(interp['top_residues'])}",
                "",
                T["dl_similar"].format(k=TOP_K),
            ]
            for r in rows:
                result_lines += [
                    f"  Rank {r['Rank']}: {r['Accession']} (sim={r['Similarity']})",
                    f"    Short Name: {r['Short Name']}",
                    f"    Drug Class: {r['Drug Class']}",
                    f"    Mechanism:  {r['Mechanism']}",
                    f"    Organism:   {r['Organism']}",
                ]
            result_lines += ["", T["dl_disclaimer"]]

            safe_name = re.sub(r'[\\/:*?"<>|]', '_', seq_data['header'][:20])
            st.download_button(
                label=T["btn_download_single"],
                data="\n".join(result_lines),
                file_name=f"ARG_result_{safe_name}.txt",
                mime="text/plain",
            )

    # Batch sum
    if is_batch and all_results:
        status_text.text(T["processing_done"])
        st.markdown("---")
        st.subheader(T["batch_header"])

        df_batch    = pd.DataFrame(all_results)
        n_resistant = (df_batch["Classification"] == "RESISTANT").sum()
        n_total     = len(df_batch)

        col_m1, col_m2, col_m3 = st.columns(3)
        with col_m1:
            st.metric(T["batch_metric_total"],         n_total)
        with col_m2:
            st.metric(T["batch_metric_resistant"],     int(n_resistant))
        with col_m3:
            st.metric(T["batch_metric_non_resistant"], int(n_total - n_resistant))

        st.markdown(T["batch_summary_text"].format(
            n_res=n_resistant,
            n_tot=n_total,
            pct=n_resistant / n_total * 100,
        ))

        col_pie, col_tbl2 = st.columns([1, 2])
        with col_pie:
            fig_pie, ax_pie = plt.subplots(figsize=(4, 4))
            ax_pie.pie(
                [n_resistant, n_total - n_resistant],
                labels=[T["label_resistant"], T["label_non_resistant"]],
                colors=["#DC2626", "#1C7293"],
                autopct="%1.1f%%",
                startangle=90,
            )
            ax_pie.set_title(T["batch_pie_title"])
            st.pyplot(fig_pie)
            plt.close()
        with col_tbl2:
            st.dataframe(
                df_batch[["Header", "Classification", "Confidence",
                           "Length (aa)", "Top-1 Match", "Drug Class"]],
                use_container_width=True,
                hide_index=True,
            )

        with st.expander(T["batch_expander"]):
            st.dataframe(df_batch, use_container_width=True, hide_index=True)

        st.download_button(
            label=T["btn_download_batch"],
            data=df_batch.to_csv(index=False),
            file_name="ARG_batch_results.csv",
            mime="text/csv",
        )


# Footer
st.markdown("---")
st.caption(T["footer"])