"""
run_lmeval.py -- drive the OFFICIAL lm-eval harness with the verified per-task
LLaDA-8B-Base protocol (common.LMEVAL_TASKS is the single source of truth).

Builds and execs:
  accelerate launch --num_processes N eval/llada_harness.py --model llada_dist
      --tasks <lm_task> --num_fewshot <fs> --batch_size <bs>
      --model_args model_path=<M>,cfg=<c>,is_check_greedy=False,mc_num=<mc>[,gen...][,weights_path=<W>]
      --output_path results/eval/<arm>/lmeval/<task> --log_samples

REF -> no weights_path (dense). BASE/OURS -> weights_path=<compressed dir>.
Per-item correctness for McNemar is produced afterward by
analysis/lmeval_to_items.py from the --log_samples output.
"""
import os
import sys
import argparse
import subprocess

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, HERE)
import common as C


def build_cmd(task, arm, weights, model_path, num_processes, limit, out_root,
              family="base"):
    cfgd = C.tasks_for(family)[task]
    harness = os.path.join(HERE, "eval", "llada_harness.py")
    out_path = os.path.join(out_root, arm, "lmeval", task)
    os.makedirs(out_path, exist_ok=True)

    margs = [f"model_path={model_path}", "is_check_greedy=False"]
    if cfgd["gen"]:
        margs += [f"gen_length={cfgd['gen_length']}", f"steps={cfgd['steps']}",
                  f"block_length={cfgd['block_length']}"]
    else:
        margs += [f"cfg={cfgd['cfg']}", f"mc_num={cfgd['mc_num']}"]
    if family == "instruct":
        # Instruct is an SFT chat model: its prompts must carry the chat template, and
        # the two EOS switches are part of the official per-task protocol. Base passes
        # none of these, so its command line is byte-identical to the finished table's.
        margs.append("chat_template=True")
        margs.append(f"logits_eos_inf={cfgd['logits_eos_inf']}")
        margs.append(f"confidence_eos_eot_inf={cfgd['confidence_eos_eot_inf']}")
    if arm != "ref" and weights:
        margs.append(f"weights_path={weights}")

    # unique rendezvous port per Slurm job so co-located DP jobs don't collide
    # on accelerate's default 29500 (observed: 2 DP jobs on one node -> ConnectionError)
    port = 20000 + (int(os.environ.get("SLURM_JOB_ID", "0")) % 20000)
    cmd = [
        "accelerate", "launch", "--num_processes", str(num_processes),
        "--main_process_port", str(port),
        harness, "--model", "llada_dist",
        "--include_path", os.path.join(HERE, "eval", "lm_tasks"),
        "--tasks", task, "--num_fewshot", str(cfgd["fs"]),
        "--batch_size", str(cfgd["batch_size"]),
        "--model_args", ",".join(margs),
        "--output_path", out_path, "--log_samples",
    ]
    if cfgd.get("code"):
        # humaneval/mbpp execute generated code (unsafe_code:true tasks)
        cmd.append("--confirm_run_unsafe_code")
    if limit:
        cmd += ["--limit", str(limit)]
    return cmd, out_path


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--family", default="base", choices=list(C.FAMILIES),
                    help="which per-task protocol table to use (common.tasks_for)")
    ap.add_argument("--task", required=True,
                    choices=sorted(set(C.LMEVAL_TASKS) | set(C.LMEVAL_TASKS_INSTRUCT)))
    ap.add_argument("--arm", required=True, choices=["ref", "base", "ours", "asvd", "wanda", "sparsegpt", "gptq", "awq", "quant_dllm", "sink"])
    ap.add_argument("--weights", default=None)
    ap.add_argument("--model_path", default=None,
                    help="defaults to the family's model (Base or Instruct)")
    ap.add_argument("--num_processes", type=int, default=1)
    ap.add_argument("--limit", default=None, help="lm-eval --limit (first-N / fraction); "
                    "validation only, omit for full runs")
    ap.add_argument("--out_root", default=os.path.join(HERE, "results", "eval"))
    ap.add_argument("--dry", action="store_true", help="print the command, do not run")
    args = ap.parse_args()

    if args.model_path is None:
        args.model_path = (C.INSTRUCT_MODEL_PATH if args.family == "instruct"
                           else C.DEFAULT_MODEL_PATH)
    if args.task not in C.tasks_for(args.family):
        sys.exit(f"task {args.task!r} is not in the {args.family} protocol table; "
                 f"available: {sorted(C.tasks_for(args.family))}")
    if args.arm == "ref" and args.weights:
        sys.exit("ref is dense: do not pass --weights")
    if args.arm in ("base", "ours", "asvd", "quant_dllm", "sink") and not args.weights:
        sys.exit(f"{args.arm} requires --weights")

    cmd, out_path = build_cmd(args.task, args.arm, args.weights,
                              os.path.expanduser(args.model_path),
                              args.num_processes, args.limit, args.out_root,
                              args.family)
    print("[run_lmeval] family=%s task=%s arm=%s metric=%s" % (
        args.family, args.task, args.arm, C.primary_metric(args.task, args.family)))
    print("[run_lmeval] " + " ".join(cmd))
    if args.dry:
        return
    env = dict(os.environ)
    env.setdefault("HF_ALLOW_CODE_EVAL", "1")
    env.setdefault("HF_DATASETS_TRUST_REMOTE_CODE", "true")
    rc = subprocess.call(cmd, env=env)
    print(f"[run_lmeval] exit={rc}  out={out_path}")
    sys.exit(rc)


if __name__ == "__main__":
    main()