# ==============================================================================
# VENDORED, UNMODIFIED (except this header) from:
#   Sink-Aware-Pruning/eval_llada.py
# which is the OFFICIAL LLaDA lm-eval wrapper (ML-GSAI/LLaDA evaluation/eval_llada.py)
# plus a compressed-weight injection (replace_with_traject_svd) and a single-GPU
# null-guard on accelerator.wait_for_everyone(). Verified vs the official file:
# core likelihood/generation math is byte-identical (diff -bw shows only additions).
#
# It registers the lm-eval model "llada_dist" and delegates to lm_eval.cli_evaluate,
# so ALL standard lm-eval flags work (--tasks --num_fewshot --batch_size
# --model_args ...,weights_path=DIR --output_path --log_samples --limit).
#
# The A/B .pt weight format it loads is exactly what trajmc_main/compress/compress.py
# writes ({name.replace('.','_')}_A.pt / _B.pt), and its ".blocks." guard matches
# common.py's lm_head defense. This is the authoritative eval engine for REF/BASE/OURS
# (README section 0: evaluation protocol follows LLaDA official).
# ==============================================================================
'''
This file is inspired by the code from https://github.com/ML-GSAI/SMDM
'''
import os
import accelerate
import torch
import re
from pathlib import Path
import random
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import datasets
datasets.disable_caching()
from datasets import Dataset
from lm_eval.__main__ import cli_evaluate
from lm_eval.api.instance import Instance
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
from tqdm import tqdm

from transformers import AutoTokenizer, AutoModel
# [trajmc_main] ensure the vendored generate.py (this dir) is importable no matter
# how the script is launched (accelerate launch does not reliably add script dir).
import sys as _sys
_sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from generate import generate


class StaticLowRank(nn.Module):
    def __init__(self, A, B, bias=None):
        super().__init__()
        self.A = nn.Parameter(A.to(torch.bfloat16), requires_grad=False)
        self.B = nn.Parameter(B.to(torch.bfloat16), requires_grad=False)
        self.bias = nn.Parameter(bias.to(torch.bfloat16), requires_grad=False) if bias is not None else None

    def forward(self, x):
        out = F.linear(F.linear(x, self.B), self.A)
        if self.bias is not None:
            out = out + self.bias
        return out


def replace_with_traject_svd(model, weights_path, device):
    n_replaced = n_skipped = 0
    for name, mod in list(model.named_modules()):
        if not isinstance(mod, nn.Linear):
            continue
        # POLICY: only the 7 linears inside the 32 transformer blocks are ever compressed
        # (224 layers). Embedding and lm_head are NEVER compressed -- standard practice
        # (SVD-LLM / ASVD do the same). The top-level lm_head is named
        # `model.transformer.ff_out`, i.e. it shares the `ff_out` suffix with each block's MLP
        # down-projection, so a bare `ff_out$` regex silently caught it and a stale A/B left in
        # a weights dir then got loaded here -- the model evaluated no longer matched its own
        # metadata. This guard makes that class of bug impossible regardless of stray files.
        if ".blocks." not in name:
            continue
        prefix = name.replace(".", "_")
        path_a = os.path.join(weights_path, f"{prefix}_A.pt")
        path_b = os.path.join(weights_path, f"{prefix}_B.pt")
        if not (os.path.exists(path_a) and os.path.exists(path_b)):
            n_skipped += 1
            continue
        parent = model
        parts = name.split(".")
        for part in parts[:-1]:
            parent = getattr(parent, part)
        bias = mod.bias.detach().to(device) if mod.bias is not None else None
        A = torch.load(path_a, map_location=device).to(torch.bfloat16)
        B = torch.load(path_b, map_location=device).to(torch.bfloat16)
        setattr(parent, parts[-1], StaticLowRank(A, B, bias).to(device))
        # [trajmc] free the old dense weight now (list(named_modules) still refs `mod`,
        # so without this the caching allocator holds ALL old weights -> ~23.5GB peak,
        # OOMs 24GB GPUs). Freeing per-layer + periodic empty_cache bounds peak ~16GB.
        mod.weight = None
        del A, B
        if n_replaced % 16 == 0:
            torch.cuda.empty_cache()
        n_replaced += 1
    torch.cuda.empty_cache()
    print(f"TrajectSVD layers replaced: {n_replaced}  skipped: {n_skipped}")
    return model


def replace_with_decoded_weights(model, weights_path, device):
    """Load DAQ/SVDQuant decoded _weight.pt files into model parameters."""
    n_replaced = n_skipped = 0
    for name, mod in model.named_modules():
        if not isinstance(mod, nn.Linear):
            continue
        prefix = name.replace(".", "_")
        path_w = os.path.join(weights_path, f"{prefix}_weight.pt")
        if not os.path.exists(path_w):
            n_skipped += 1
            continue
        W = torch.load(path_w, map_location=device).to(torch.bfloat16)
        mod.weight.data = W
        n_replaced += 1
    print(f"Decoded weight layers replaced: {n_replaced}  skipped: {n_skipped}")
    return model


def replace_with_channel_masks(model, weights_path, device):
    """Apply Week-2 ff_out channel masks: zero pruned input channels (columns).
    Loads {prefix}_chanmask.pt (bool keep-mask, in_features long) per nn.Linear."""
    n_replaced = n_skipped = 0
    for name, mod in model.named_modules():
        if not isinstance(mod, nn.Linear):
            continue
        prefix = name.replace(".", "_")
        path_m = os.path.join(weights_path, f"{prefix}_chanmask.pt")
        if not os.path.exists(path_m):
            n_skipped += 1
            continue
        keep = torch.load(path_m, map_location=device)
        mod.weight.data = (mod.weight.data * keep.to(mod.weight.dtype).unsqueeze(0))
        n_replaced += 1
    print(f"Channel-mask layers pruned: {n_replaced}  skipped: {n_skipped}")
    return model


def set_seed(seed):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


@register_model("llada_dist")
class LLaDAEvalHarness(LM):
    def __init__(
        self,
        model_path='',
        mask_id=126336,
        max_length=4096,
        batch_size=32,
        mc_num=128,
        is_check_greedy=True,
        cfg=0.,
        weights_path=None,
        steps=1024,
        gen_length=1024,
        block_length=1024,
        remasking='low_confidence',
        device="cuda",
        chat_template=False,
        logits_eos_inf=False,
        confidence_eos_eot_inf=False,
        **kwargs,
    ):
        '''
        Args:
            model_path: LLaDA-8B-Base model path.
            mask_id: The token id of [MASK] is 126336.
            max_length: the max sequence length.
            batch_size: mini batch size.
            mc_num: Monte Carlo estimation iterations
            is_check_greedy: For certain metrics like LAMBADA, the evaluation requires the model to verify whether the answer 
                             is generated through greedy sampling conditioned on the prompt (note that this differs from conditional
                             generation). We implement this verification through the suffix_greedy_prediction() function, which 
                             returns a True/False judgment used for accuracy calculation. 
                             When is_check_greedy is set to True, the lm-evaluation-harness library automatically invokes this function. 
                             However, since none of the metrics in the LLaDA paper (https://arxiv.org/abs/2502.09992) require this functionality, 
                             we recommend setting is_check_greedy to False. This configuration causes suffix_greedy_prediction() to return False 
                             by default, significantly accelerating the evaluation process.
            cfg_scale: Unsupervised classifier-free guidance scale.
            chat_template: wrap each generate_until prompt in the model's chat template
                (LLaDA-8B-Instruct only). Base is a plain continuation model whose
                few-shot prompts already carry their own "Q:/A:" framing, so it must stay
                False there -- turning it on for Base would be an out-of-distribution input.
            logits_eos_inf / confidence_eos_eot_inf: Instruct EOS-suppression switches,
                forwarded to generate(); see eval/generate.py for the semantics and the
                caveat that they are our reading of the official EVAL.md table.
        '''
        super().__init__()

        accelerator = accelerate.Accelerator()
        if accelerator.num_processes > 1:
            self.accelerator = accelerator
        else:
            self.accelerator = None
        
        model_kwargs = {}
        if self.accelerator is not None:
            model_kwargs.update({'device_map': {'': f'{self.accelerator.device}'}})

        self.model = AutoModel.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16, **model_kwargs)
        self.model.eval()

        self.device = torch.device(device)
        if self.accelerator is not None:
            # [trajmc_main FIX] Inference-only data parallelism: do NOT call
            # accelerator.prepare() -- it wraps the model in DistributedDataParallel,
            # which allocates a SECOND full copy of the params (gradient reduction
            # buckets, needed only for training). For an 8B model that is +16GB and
            # OOMs a 24GB GPU (observed on A5000: "Tried to allocate 14.93 GiB").
            # lm-eval already shards docs across ranks and gathers results itself, so
            # the model only needs to live on each rank's device. Numerically identical
            # to single-process; just avoids the training-only DDP overhead.
            self.model = self.model.to(self.accelerator.device)
            self.device = torch.device(f'{self.accelerator.device}')
            self._rank = self.accelerator.local_process_index
            self._world_size = self.accelerator.num_processes
        else:
            self.model = self.model.to(device)

        if weights_path and str(weights_path).lower() not in ("", "none", "null"):
            wdir = Path(weights_path)
            if any(wdir.glob("*_A.pt")):
                self.model = replace_with_traject_svd(self.model, weights_path, self.device)
            elif any(wdir.glob("*_chanmask.pt")):
                self.model = replace_with_channel_masks(self.model, weights_path, self.device)
            elif any(wdir.glob("block_*.pt")):
                # structured-pruned model (head/channel): swap in smaller Linears + head-aware attn
                import sys as _sys
                _sys.path.insert(0, "/zpool-00/home/tl356/LLaDA/supplement")
                from structured_prune_lib import apply_structured_prune
                self.model = apply_structured_prune(self.model, weights_path, self.device)
            else:
                self.model = replace_with_decoded_weights(self.model, weights_path, self.device)

        self.mask_id = mask_id
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

        # lm-eval turns model_args into python via handle_arg_string ("True" -> True),
        # but coerce anyway so a hand-built string like chat_template=1 cannot silently
        # evaluate truthy-but-wrong.
        def _as_bool(v):
            if isinstance(v, str):
                return v.strip().lower() in ("1", "true", "yes", "y")
            return bool(v)

        self.chat_template = _as_bool(chat_template)
        self.logits_eos_inf = _as_bool(logits_eos_inf)
        self.confidence_eos_eot_inf = _as_bool(confidence_eos_eot_inf)

        # <|eot_id|> exists only in the Instruct vocab; convert_tokens_to_ids returns the
        # unk id (not None) for an absent token, so compare against the vocab explicitly.
        vocab = self.tokenizer.get_vocab()
        self.eos_id = vocab.get("<|endoftext|>")
        self.eot_id = vocab.get("<|eot_id|>")

        if self.chat_template and self.tokenizer.chat_template is None:
            raise ValueError(f"chat_template=True but {model_path} has no chat_template")
        if self.logits_eos_inf and self.eos_id is None:
            raise ValueError("logits_eos_inf=True but <|endoftext|> not in vocab")
        if self.confidence_eos_eot_inf and self.eos_id is None and self.eot_id is None:
            raise ValueError("confidence_eos_eot_inf=True but neither EOS nor EOT in vocab")
        if self.rank == 0:
            print(f"[llada_dist] chat_template={self.chat_template} "
                  f"logits_eos_inf={self.logits_eos_inf} "
                  f"confidence_eos_eot_inf={self.confidence_eos_eot_inf} "
                  f"eos_id={self.eos_id} eot_id={self.eot_id}", flush=True)

        self.mc_num = mc_num
        self.batch_size = int(batch_size)
        assert mc_num % self.batch_size == 0
        self.sampling_eps = 0.
        self.max_length = max_length
        self.is_check_greedy = is_check_greedy

        self.cfg = cfg
        self.steps = steps
        self.gen_length = gen_length
        self.block_length = block_length
        self.remasking = remasking    
    @property
    def rank(self):
        return self._rank
    
    @property
    def world_size(self):
        return self._world_size

    def _forward_process(self, batch, prompt_index):
        b, l = batch.shape

        target_len = (l - prompt_index.sum()).item()
        k = torch.randint(1, target_len + 1, (), device=batch.device)

        x = torch.round(torch.linspace(float(k), k + (b - 1) * (target_len / b), steps=b, device=batch.device)).long()
        x = ((x - 1) % target_len) + 1
        assert x.min() >= 1 and x.max() <= target_len

        indices = torch.arange(target_len, device=batch.device).repeat(b, 1)
        is_mask = indices < x.unsqueeze(1)

        for i in range(b):
            is_mask[i] = is_mask[i][torch.randperm(target_len)]

        is_mask = torch.cat((torch.zeros(b, prompt_index.sum(), dtype=torch.bool, device=batch.device), is_mask), dim=1)

        noisy_batch = torch.where(is_mask, self.mask_id, batch)

        return noisy_batch, (x / target_len).unsqueeze(1).repeat(1, l)

    @torch.no_grad()
    def get_logits(self, batch, prompt_index):
        if self.cfg > 0.:
            assert len(prompt_index) == batch.shape[1]
            prompt_index = prompt_index.unsqueeze(0).repeat(batch.shape[0], 1)
            un_batch = batch.clone()
            un_batch[prompt_index] = self.mask_id
            batch = torch.cat([batch, un_batch])

        logits = self.model(batch).logits

        if self.cfg > 0.:
            logits, un_logits = torch.chunk(logits, 2, dim=0)
            logits = un_logits + (self.cfg + 1) * (logits - un_logits)
        return logits[:, :batch.shape[1]]

    @torch.no_grad()
    def get_loglikelihood(self, prefix, target):
        seq = torch.concatenate([prefix, target])[None, :]
        seq = seq.repeat((self.batch_size, 1)).to(self.device)

        prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix)

        loss_acc = []
        for _ in range(self.mc_num // self.batch_size):
            perturbed_seq, p_mask = self._forward_process(seq, prompt_index)

            mask_indices = perturbed_seq == self.mask_id

            logits = self.get_logits(perturbed_seq, prompt_index)

            loss = F.cross_entropy(logits[mask_indices], seq[mask_indices], reduction='none') / p_mask[mask_indices]
            loss = loss.sum() / self.batch_size
            loss_acc.append(loss.item())

        return - sum(loss_acc) / len(loss_acc)

    @torch.no_grad()
    def suffix_greedy_prediction(self, prefix, target):
        if not self.is_check_greedy:
            return False

        seq = torch.full((1, len(prefix) + len(target)), self.mask_id, device=self.device)
        prompt_index = torch.arange(seq.shape[1], device=self.device) < len(prefix)
        prefix, target = prefix.to(self.device), target.to(self.device)
        seq[0, :len(prefix)] = prefix

        for i in range(len(target)):
            mask_index = (seq == self.mask_id)
            logits = self.get_logits(seq, prompt_index)[mask_index]
            x0 = torch.argmax(logits, dim=-1)

            p = torch.softmax(logits.to(torch.float32), dim=-1)
            confidence = torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)).squeeze(dim=-1)
            _, index = torch.sort(confidence, descending=True)
            x0[index[1:]] = self.mask_id
            seq[mask_index] = x0.clone()
        correct = target == seq[0, len(prefix):]
        correct = torch.all(correct)
        return correct

    def _encode_pair(self, context, continuation):
        n_spaces = len(context) - len(context.rstrip())
        if n_spaces > 0:
            continuation = context[-n_spaces:] + continuation
            context = context[:-n_spaces]

        whole_enc = self.tokenizer(context + continuation)["input_ids"]
        context_enc = self.tokenizer(context)["input_ids"]

        context_enc_len = len(context_enc)
        continuation_enc = whole_enc[context_enc_len:]

        return context_enc, continuation_enc

    def loglikelihood(self, requests):
        def _tokenize(e):
            prefix, target = self._encode_pair(e["prefix"], e["target"])
            return {
                "prefix_text": e["prefix"],
                "target_text": e["target"],
                "prefix": prefix,
                "target": target,
            }

        ds = []
        ds = [{"prefix": req.args[0], "target": req.args[1]} for req in requests]
        ds = Dataset.from_list(ds)
        ds = ds.map(_tokenize)
        ds = ds.with_format("torch")
        prompt_len = [len(x["prefix"]) + len(x["target"]) for x in ds]

        assert max(prompt_len) <= 4096

        out = []
        with torch.no_grad():
            for elem in tqdm(ds, desc="Computing likelihood..."):
                prefix = elem["prefix"]
                target = elem["target"]

                ll = self.get_loglikelihood(prefix, target)

                is_target_greedy_dec = self.suffix_greedy_prediction(prefix, target)

                out.append((ll, 1.0 if is_target_greedy_dec else 0.0))
        torch.cuda.empty_cache()
        return out

    def loglikelihood_rolling(self, requests):
        raise NotImplementedError

    def _wrap_prompt(self, text):
        """Instruct: wrap lm-eval's prompt (few-shot block included) as one user turn.

        Re-tokenizing the tokenize=False string is safe here -- verified byte-identical to
        apply_chat_template(tokenize=True), i.e. no duplicated <|startoftext|>.
        """
        if not self.chat_template:
            return text
        return self.tokenizer.apply_chat_template(
            [{"role": "user", "content": text}],
            add_generation_prompt=True,
            tokenize=False,
        )

    def generate_until(self, requests: list[Instance]):
        def _tokenize(e):
            return {
                "question": self.tokenizer(self._wrap_prompt(e["question"]))["input_ids"],
                "question_text": e["question"],
                "until": e["until"],
            }

        ds = [{"question": req.args[0], "until": req.args[1]['until']} for req in requests]
        ds = Dataset.from_list(ds)
        ds = ds.map(_tokenize)
        ds = ds.with_format("torch")

        out = []
        for elem in tqdm(ds, desc="Generating..."):
            prompt = elem["question"].unsqueeze(0).to(self.device)
            stop_tokens = elem["until"]
 
            generated_answer = generate(self.model, prompt, steps=self.steps, gen_length=self.gen_length, block_length=self.block_length,
                                        temperature=0, cfg_scale=self.cfg, remasking=self.remasking, mask_id=self.mask_id,
                                        eos_id=self.eos_id, eot_id=self.eot_id,
                                        logits_eos_inf=self.logits_eos_inf,
                                        confidence_eos_eot_inf=self.confidence_eos_eot_inf)
            
            generated_answer = self.tokenizer.decode(generated_answer[0][prompt.shape[1]:], skip_special_tokens=False)
            for stop_seq in stop_tokens:
                    if stop_seq in generated_answer:
                        generated_answer = generated_answer.split(stop_seq)[0]

            # remove special tokens
            generated_answer_ids = self.tokenizer(generated_answer)["input_ids"]
            generated_answer = self.tokenizer.decode(generated_answer_ids, skip_special_tokens=True)
            out.append(generated_answer)

            if self.accelerator is not None:
                self.accelerator.wait_for_everyone()

        return out


if __name__ == "__main__":
    set_seed(1234)
    cli_evaluate()