"""
Inference for the Tox21 leaderboard Space.

Contract (from the leaderboard paper, Appendix A.1 and Section 4):
  - predict() takes a list of SMILES strings
  - returns {smiles: {task_name: score}} for ALL 12 tasks
  - scores are floats in [0, 1]
  - EVERY input molecule must get a prediction for EVERY task. The leaderboard
    validates all molecule-target pairs are present and errors out otherwise.
    Never drop a molecule, no matter what the preprocessing does.
"""

import os
import torch
import torch.nn as nn
import deepsmiles
from rdkit import Chem, RDLogger
from transformers import AutoTokenizer, AutoModel

RDLogger.DisableLog("rdApp.*")

TASKS = [
    "NR-AR", "NR-AR-LBD", "NR-AhR", "NR-Aromatase", "NR-ER", "NR-ER-LBD",
    "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53",
]

CHECKPOINT = os.environ.get("CHECKPOINT_DIR", "checkpoints/deepberta-tox21")
CONVERTER = deepsmiles.Converter(rings=True, branches=True)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Neutral score for molecules we cannot encode. 0.5 is uninformative for
# ROC-AUC rather than actively wrong. Any conversion failures MUST be
# reported in the model card -- silently scoring 0.5 on part of the test
# set is exactly the kind of thing reviewers need disclosed.
FALLBACK_SCORE = 0.5


class DeepBERTaForTox21(nn.Module):
    def __init__(self, path, n_tasks=12, dropout=0.1):
        super().__init__()
        self.encoder = AutoModel.from_pretrained(path)
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(self.encoder.config.hidden_size, n_tasks)

    def forward(self, input_ids, attention_mask):
        out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        mask = attention_mask.unsqueeze(-1).float()
        pooled = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
        return self.classifier(self.dropout(pooled))


_tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
_model = DeepBERTaForTox21(CHECKPOINT).to(DEVICE)
_model.load_state_dict(
    torch.load(os.path.join(CHECKPOINT, "model.pt"), map_location=DEVICE,
               weights_only=True)
)
_model.eval()


def _encode(smiles):
    """SMILES -> DeepSMILES. Returns None if unconvertible."""
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is None:
            return None
        return CONVERTER.encode(Chem.MolToSmiles(mol))
    except Exception:
        return None


def predict(smiles_list, batch_size=32):
    results = {s: {t: FALLBACK_SCORE for t in TASKS} for s in smiles_list}

    encodable = [(s, _encode(s)) for s in smiles_list]
    usable = [(s, d) for s, d in encodable if d is not None]
    n_failed = len(smiles_list) - len(usable)
    if n_failed:
        print(f"WARNING: {n_failed}/{len(smiles_list)} molecules failed "
              f"DeepSMILES conversion; returning {FALLBACK_SCORE} for these.")

    with torch.no_grad():
        for i in range(0, len(usable), batch_size):
            chunk = usable[i:i + batch_size]
            enc = _tokenizer(
                [d for _, d in chunk], padding=True, truncation=True,
                max_length=256, return_tensors="pt",
            ).to(DEVICE)
            # sigmoid: the leaderboard requires values in [0, 1], not logits
            probs = torch.sigmoid(_model(**enc)).cpu().numpy()
            for (smi, _), row in zip(chunk, probs):
                results[smi] = {t: float(p) for t, p in zip(TASKS, row)}

    # The leaderboard requires a prediction for every molecule-target pair it
    # sent. The test set is 647 molecules but only 645 unique structures, so
    # check against the unique inputs -- and verify none went missing.
    missing = set(smiles_list) - set(results)
    assert not missing, f"no prediction for {len(missing)} molecules: {list(missing)[:5]}"
    for s, d in results.items():
        assert set(d) == set(TASKS), f"incomplete predictions for {s}"
        assert all(0.0 <= v <= 1.0 for v in d.values()), f"out of range for {s}"

    print(f"returned predictions for {len(results)} unique structures "
          f"from {len(smiles_list)} inputs, {len(TASKS)} tasks each")

    return results


if __name__ == "__main__":
    # awkward cases worth checking before deploying: salts, stereochemistry,
    # charged species, very small molecules, and deliberate garbage
    tests = [
        "CCO",
        "CC(=O)Oc1ccccc1C(=O)O",
        "C[C@H](N)C(=O)O",
        "[Na+].[Cl-]",
        "c1ccc2c(c1)ccc1ccccc12",
        "not_a_molecule",
    ]
    for smi, preds in predict(tests).items():
        print(f"{smi:35s} {preds['NR-AR']:.3f}  {preds['SR-MMP']:.3f}")