# -*- coding: utf-8 -*-
import os
os.environ["CUDA_VISIBLE_DEVICES"] = os.environ.get("GPU", "2")
os.environ.setdefault("HF_DATASETS_TRUST_REMOTE_CODE", "1")
import sys, json, importlib.util, torch

MODEL = os.environ.get("MODEL", "/workspace/aether-annealed")
MERGE_ADAPTER = os.environ.get("MERGE_ADAPTER", "")
ADAPTER = os.environ.get("ADAPTER", "")
TAG = os.environ.get("TAG", "annealed")
LIMIT = os.environ.get("LIMIT", "")
LIMIT = int(LIMIT) if LIMIT else None
TASKS = os.environ.get("TASKS", "arc_easy,arc_challenge,boolq,hellaswag,openbookqa,piqa,sciq,winogrande").split(",")

sys.path.insert(0, MODEL)
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained(MODEL)
spec = importlib.util.spec_from_file_location("mini_config", MODEL + "/mini_config.py")
mc = importlib.util.module_from_spec(spec); spec.loader.exec_module(mc)
cfg = mc.mini_cfg
cfg.use_cache = False
cfg.output_router_logits = False
from aether_pkg.modeling_aether_v2_7way import AETHERV27wayForCausalLM
from safetensors.torch import load_file
print("[load]", MODEL, flush=True)
model = AETHERV27wayForCausalLM(cfg).to("cuda", dtype=torch.bfloat16)
model.load_state_dict(load_file(MODEL + "/model.safetensors"), strict=False)
if MERGE_ADAPTER:
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, MERGE_ADAPTER).merge_and_unload(); print("[merged]", MERGE_ADAPTER, flush=True)
if ADAPTER:
    from peft import PeftModel
    model = PeftModel.from_pretrained(model, ADAPTER); print("[adapter]", ADAPTER, flush=True)
model.eval()

import lm_eval
from lm_eval.models.huggingface import HFLM
# batch_size=1 on purpose: the NSA branches ignore attention_mask, so any left-padding
# in a batch would let queries attend to pad tokens and silently corrupt the scores.
lm = HFLM(pretrained=model, tokenizer=tok, batch_size=1, max_length=2048, backend="causal")
print("[lm_eval] tasks=%s limit=%s" % (TASKS, LIMIT), flush=True)
res = lm_eval.simple_evaluate(model=lm, tasks=TASKS, num_fewshot=0, limit=LIMIT,
                              bootstrap_iters=1000, verbosity="ERROR")
out = {}
print("\n=== AETHER-7B-5Attn (%s) — lm-eval-harness 0-shot ===" % TAG, flush=True)
print("%-18s %8s %8s   %8s %8s" % ("TASK", "acc", "±", "acc_norm", "±"), flush=True)
print("-" * 60, flush=True)
for t in TASKS:
    r = res["results"].get(t, {})
    a = r.get("acc,none"); ae = r.get("acc_stderr,none")
    n = r.get("acc_norm,none"); ne = r.get("acc_norm_stderr,none")
    out[t] = {"acc": a, "acc_stderr": ae, "acc_norm": n, "acc_norm_stderr": ne}
    f = lambda v: ("%6.1f" % (100 * v)) if isinstance(v, float) else "     -"
    print("%-18s %8s %8s   %8s %8s" % (t, f(a), f(ae), f(n), f(ne)), flush=True)
print("-" * 60, flush=True)
json.dump(out, open("/workspace/lmeval_%s.json" % TAG, "w"), indent=2, ensure_ascii=False)
print("[saved] /workspace/lmeval_%s.json" % TAG, flush=True)
print("LMEVAL_DONE", flush=True)