"""
Calculates acceptance rate (LCP) and ranking quality (NDCG) for a language
model with a token order prediction head.

The evaluation compares the sequence proposed by the special head against the
sequence greedily verified by the standard NTP head.
"""
import fla
import argparse
import random
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from tqdm import tqdm

def calculate_lcp_from_sequences(predicted_sequences, verified_sequences):
    """
    Calculates LCP and second-token-match rate by comparing two token sequences.
    This version implements the original notebook's logic in a vectorized way.
    """
    batch_size, eval_length = predicted_sequences.size()

    if eval_length == 0:
        return 0, 0
    if eval_length == 1:
        # LCP is always 1 for a single token, no second token to check.
        return batch_size, 0

    # Check for second token match (at index 1)
    second_token_matches_tensor = (predicted_sequences[:, 1] == verified_sequences[:, 1])
    second_token_matches = torch.sum(second_token_matches_tensor).item()

    # Match the tail of the sequence (from the second token onwards)
    tail_matches = (predicted_sequences[:, 1:] == verified_sequences[:, 1:])

    # Use cumprod to find the length of the prefix of matches in the tail
    # This finds how many tokens in a row match, starting from the second token
    prefix_matches = torch.cumprod(tail_matches.int(), dim=1)
    tail_lcp_lengths = torch.sum(prefix_matches, dim=1)

    # Per the notebook's logic, total LCP is 1 (for the first token) + LCP of the tail
    total_lcp = torch.sum(1 + tail_lcp_lengths).item()

    return total_lcp, second_token_matches

def calculate_ndcg_from_sequences(predicted_sequences, verified_sequences):
    """
    Calculates NDCG based on position-wise matches between two sequences.
    The relevance score (gain) is 1 for a match and 0 for a mismatch.
    """
    k = predicted_sequences.size(1)
    if k == 0:
        return 0.0

    # Gains are 1 for a match, 0 for a mismatch
    gains = (predicted_sequences == verified_sequences).float()

    # Calculate DCG (Discounted Cumulative Gain)
    discounts = torch.log2(torch.arange(k, device=gains.device).float() + 2.0)
    dcg = torch.sum(gains / discounts, dim=1)

    # Calculate IDCG (Ideal DCG) by sorting the gains to rank all matches first
    ideal_gains, _ = torch.sort(gains, dim=1, descending=True)
    idcg = torch.sum(ideal_gains / discounts, dim=1)

    # Calculate NDCG, handling the case where IDCG is 0 (no matches)
    ndcg = dcg / (idcg + 1e-8)
    ndcg[idcg == 0] = 0.0

    return torch.sum(ndcg).item()

def main(args):
    """Main function to run the metric calculations."""
    print("Starting model evaluation with corrected sequence-based metrics...")
    print(f"Arguments: {vars(args)}")

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    if device == "cpu": print("⚠️ Warning: No CUDA device found. Running on CPU will be very slow.")

    print(f"Loading model: {args.model_name}...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name, cache_dir=args.cache_dir
    ).to(torch.bfloat16).to(device)
    model.eval()

    print(f"Loading dataset: {args.dataset_name} (Subset: {args.dataset_subset}, Split: {args.dataset_split})...")
    dataset = load_dataset(
        args.dataset_name,
        name=args.dataset_subset,
        split=args.dataset_split,
        cache_dir=args.cache_dir,
        streaming=True
    )

    total_samples_processed = 0
    total_lcp_length = 0
    total_second_token_matches = 0
    total_ndcg_scores = {k: 0.0 for k in args.ndcg_k}

    max_k_needed = max(args.ndcg_k + [args.lcp_eval_length])
    print(f"Will generate and verify top {max_k_needed} candidates per sample.")

    dataset_iterator = iter(dataset)
    pbar_total = args.num_samples if args.num_samples != -1 else None

    with tqdm(total=pbar_total, desc="Processing samples") as pbar:
        while True:
            if args.num_samples != -1 and total_samples_processed >= args.num_samples: break

            batch_input_ids = []
            num_to_fetch = min(args.batch_size, (args.num_samples - total_samples_processed) if args.num_samples != -1 else args.batch_size)

            while len(batch_input_ids) < num_to_fetch:
                try:
                    sample = next(dataset_iterator)
                    text = sample[args.text_column]
                    if not text or not isinstance(text, str): continue
                    tokens = tokenizer.encode(text)
                    if len(tokens) >= args.context_length:
                        start = random.randint(0, len(tokens) - args.context_length)
                        cropped_tokens = tokens[start : start + args.context_length]
                        batch_input_ids.append(torch.tensor(cropped_tokens))
                except StopIteration:
                    break

            if not batch_input_ids: break

            input_ids = torch.stack(batch_input_ids).to(device)

            with torch.no_grad():
                # Step 1: Get the proposed sequence from the special head
                input_dict = {'input_ids': input_ids, 'attention_mask': torch.ones_like(input_ids)}
                _, top_logits = model(**input_dict, output_top_logits=True).logits
                top_sequences_full = torch.argsort(top_logits[:, -1, :], dim=-1, descending=True)[:, :max_k_needed]

                # Step 2: Run verification pass to get the ground-truth greedy sequence
                verify_input_ids = torch.cat([input_ids, top_sequences_full], dim=1)
                verify_dict = {'input_ids': verify_input_ids, 'attention_mask': torch.ones_like(verify_input_ids)}
                verify_ntp_logits = model(**verify_dict).logits
                verify_output_sequence = torch.argmax(verify_ntp_logits, dim=-1)

                # Extract the verified tokens corresponding to the proposed tokens
                verified_tokens_full = verify_output_sequence[:, args.context_length-1:-1]

                # --- Metric Calculations ---
                # A) Calculate LCP
                predicted_lcp = top_sequences_full[:, :args.lcp_eval_length]
                verified_lcp = verified_tokens_full[:, :args.lcp_eval_length]
                lcp_sum, second_token_sum = calculate_lcp_from_sequences(predicted_lcp, verified_lcp)

                # B) Calculate NDCG for each specified k
                for k in args.ndcg_k:
                    predicted_ndcg = top_sequences_full[:, :k]
                    verified_ndcg = verified_tokens_full[:, :k]
                    ndcg_sum_k = calculate_ndcg_from_sequences(predicted_ndcg, verified_ndcg)
                    total_ndcg_scores[k] += ndcg_sum_k

            # --- Update statistics ---
            current_batch_size = input_ids.size(0)
            total_lcp_length += lcp_sum
            total_second_token_matches += second_token_sum
            total_samples_processed += current_batch_size
            pbar.update(current_batch_size)

            if total_samples_processed > 0:
                avg_lcp = total_lcp_length / total_samples_processed
                ndcg_descs = [f"NDCG@{k}: {(total_ndcg_scores[k] / total_samples_processed):.3f}" for k in sorted(args.ndcg_k)]
                pbar.set_description(f"Avg LCP: {avg_lcp:.2f} | " + " | ".join(ndcg_descs))

    print("\n--- Calculation Complete ---")
    if total_samples_processed == 0:
        print("No valid samples were processed.")
        return

    print(f"Total samples processed: {total_samples_processed}")
    print("\n--- LCP Metrics ---")
    print(f"Average Acceptance Rate (LCP): {(total_lcp_length / total_samples_processed):.4f}")
    print(f"Second Token Accuracy Rate: {(total_second_token_matches / total_samples_processed):.4%}")

    print("\n--- NDCG Metrics (Sequence Match) ---")
    for k in sorted(args.ndcg_k):
        final_avg_ndcg = total_ndcg_scores[k] / total_samples_processed
        print(f"Average NDCG@{k}: {final_avg_ndcg:.4f}")
    print("-------------------------------------")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Calculate LCP and NDCG for a custom HF language model.")

    parser.add_argument("--model_name", type=str, required=True, help="Hugging Face model name.")
    parser.add_argument("--dataset_name", type=str, required=True, help="Hugging Face dataset name.")
    parser.add_argument("--dataset_subset", type=str, default=None, help="The subset/configuration of the dataset to use (e.g., 'un-common_crawl').")
    parser.add_argument("--dataset_split", type=str, default="train", help="Dataset split to use.")
    parser.add_argument("--text_column", type=str, default="text", help="The name of the text column in the dataset.")
    parser.add_argument("--num_samples", type=int, default=1024, help="Total samples to process. Set to -1 for the entire dataset.")
    parser.add_argument("--batch_size", type=int, default=8, help="Number of samples to process in a single batch.")
    parser.add_argument("--context_length", type=int, default=512, help="The length of the token sequence to use as context.")
    parser.add_argument("--lcp_eval_length", type=int, default=10, help="The number of candidate tokens to verify for the LCP metric.")
    parser.add_argument("--ndcg_k", type=int, nargs='+', default=[10], help="A list of k-values for which to calculate NDCG@k (e.g., 5 10 20).")
    parser.add_argument("--cache_dir", type=str, default=None, help="Directory to cache models and datasets.")

    args = parser.parse_args()
    main(args)