"""Emoji prediction with tokens the model almost never saw.

tweet_eval 'emoji' (SemEval-2018 task 2): 20 emoji classes; given a tweet, which emoji goes with it? Each emoji is a
single Gemma token that appeared only ~3-120 times in our 2.05B English training tokens (a rare row: in A nearly
untrained; in B/C the frozen code still carries Gemma's meaning). Zero-shot: score the 20 emoji as the next token after
"<tweet> " and take the best. Accuracy with 95% bootstrap range; baselines: uniform 5%, majority class.

  python emoji_test.py --arm B --ckpt hf/B/step_07820.pt [--n 2000]
"""
import argparse, json, pathlib
import numpy as np, pandas as pd, torch, torch.nn.functional as F
import fastc

# SemEval-2018 Task 2 (English) label order, as used by tweet_eval 'emoji'
EMOJI = ["❤", "😍", "😂", "💕", "🔥", "😊", "😎", "✨", "💙", "😘", "📷", "🇺🇸", "☀", "💜", "😉", "💯", "😁", "🎄", "📸", "😜"]


@torch.no_grad()
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True); ap.add_argument("--ckpt", required=True); ap.add_argument("--tag", default="")
    ap.add_argument("--n", type=int, default=2000); ap.add_argument("--data", default="data")
    a = ap.parse_args()
    from transformers import AutoTokenizer
    from huggingface_hub import hf_hub_download
    tok = AutoTokenizer.from_pretrained("google/gemma-3-270m"); sep = tok.eos_token_id
    counts = np.load(pathlib.Path(a.data) / "unigram_counts.npy")
    df = pd.read_parquet(hf_hub_download("cardiffnlp/tweet_eval", "emoji/test-00000-of-00001.parquet", repo_type="dataset"))
    df = df.sample(frac=1, random_state=0).head(a.n)
    cands = [tok(" " + e, add_special_tokens=False).input_ids for e in EMOJI]
    info = {e: {"pieces": len(c), "train_count": int(sum(counts[i] for i in c))} for e, c in zip(EMOJI, cands)}
    m, tag = fastc.load(a.arm, a.ckpt, counts); tag = a.tag or tag
    correct, lp_true = [], []
    for text, y in zip(df["text"], df["label"]):
        ids = [sep] + tok(str(text).strip(), add_special_tokens=False).input_ids
        ids = ids[-500:]
        scores = []
        for c in cands:
            x = torch.tensor([ids + c], device="cuda")
            with torch.autocast("cuda", dtype=torch.bfloat16):
                lg = m.logits(m.hidden(x[:, :-1])).float()
            lp = F.log_softmax(lg[0, -len(c):], -1)
            scores.append(float(lp.gather(-1, x[0, -len(c):, None]).sum()))
        pred = int(np.argmax(scores)); correct.append(int(pred == int(y))); lp_true.append(scores[int(y)])
    correct = np.array(correct); rng = np.random.default_rng(0)
    b = correct[rng.integers(0, len(correct), (10000, len(correct)))].mean(1)
    maj = df["label"].value_counts(normalize=True).max()
    R = {"tag": tag, "n": int(len(correct)), "accuracy": float(correct.mean()), "ci95": [float(np.percentile(b, 2.5)), float(np.percentile(b, 97.5))],
         "majority_baseline": float(maj), "uniform_baseline": 0.05, "mean_logprob_true": float(np.mean(lp_true)), "emoji": info,
         "per_label_acc": {EMOJI[k]: float(correct[df["label"].values == k].mean()) for k in range(20) if (df["label"].values == k).any()}}
    print(f"{tag}: accuracy {R['accuracy']*100:.1f}% ({R['ci95'][0]*100:.1f}-{R['ci95'][1]*100:.1f}) | majority {maj*100:.1f}% | uniform 5%", flush=True)
    pathlib.Path("reports").mkdir(exist_ok=True)
    json.dump(R, open(f"reports/emoji_{tag}.json", "w"), indent=1, ensure_ascii=False)


if __name__ == "__main__":
    main()