# ================== evaluate_perplexity_final.py ==================
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoProcessor
import numpy as np
from tqdm import tqdm
import argparse
import os

# ================== Command Line Arguments ==================
parser = argparse.ArgumentParser(description="Calculate Perplexity for the final merged Gemma-4 model")
parser.add_argument("--model_path", type=str, required=True,
                    help="Path to the final merged model directory")
parser.add_argument("--dataset_path", type=str, default="data/gemma-4-26B-A4B-it/generated_outputs.jsonl",
                    help="Path to the evaluation dataset (jsonl)")
parser.add_argument("--max_seq_length", type=int, default=8192,
                    help="Maximum sequence length for evaluation")
parser.add_argument("--stride", type=int, default=512,
                    help="Stride for sliding window evaluation")
parser.add_argument("--num_samples", type=int, default=100,
                    help="Number of samples to use for evaluation (0 or None = use all)")
parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16", "fp32"],
                    help="Model data type (bf16 recommended)")
args = parser.parse_args()

# ================== Configuration ==================
if args.dtype == "bf16":
    torch_dtype = torch.bfloat16
elif args.dtype == "fp16":
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

print(f"Model Path     : {args.model_path}")
print(f"Eval Samples   : {args.num_samples if args.num_samples else 'All'}")
print(f"Max Length     : {args.max_seq_length}")
print(f"Stride         : {args.stride}")

# ================== Load Model and processor ==================
print("\nLoading model and processor...")
processor = AutoProcessor.from_pretrained(args.model_path, trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained(
    args.model_path,
    dtype=torch_dtype,
    device_map="auto",
    trust_remote_code=True,
    # attn_implementation="flash_attention_2",   # Uncomment if you have Flash Attention 2 installed
)

model.eval()
print("Model loaded successfully!")

# ================== Data Preparation ==================
def prepare_text(example):
    """Prepare text with the same format used during training."""
    text = example["solution"].strip()
    if not text.startswith("<bos>"):
        text = "<bos>" + text
    if not text.endswith("<turn|>"):
        text = text.rstrip() + "<turn|>"
    return text

print("Loading evaluation dataset...")
dataset = load_dataset("json", data_files=args.dataset_path, split="train")
texts = [prepare_text(ex) for ex in dataset]

if args.num_samples and args.num_samples > 0:
    texts = texts[:args.num_samples]

print(f"Using {len(texts)} samples for Perplexity evaluation.")

# ================== Perplexity Calculation Function ==================
def calculate_perplexity(model, processor, texts, max_length=8192, stride=512):
    """Calculate perplexity using sliding window method."""
    total_loss = 0.0
    total_tokens = 0
    
    with torch.no_grad():
        for text in tqdm(texts, desc="Calculating Perplexity"):
            encodings = processor.tokenizer(text, return_tensors="pt", truncation=False)
            input_ids = encodings.input_ids.to(model.device)
            seq_len = input_ids.size(1)
            
            prev_end_loc = 0
            for begin_loc in range(0, seq_len, stride):
                end_loc = min(begin_loc + max_length, seq_len)
                trg_len = end_loc - prev_end_loc
                
                if trg_len <= 1:
                    break
                    
                input_ids_chunk = input_ids[:, begin_loc:end_loc]
                
                outputs = model(input_ids_chunk)
                logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
                
                # Shift for next token prediction
                shift_logits = logits[:, :-1, :].contiguous()
                shift_labels = input_ids_chunk[:, 1:].contiguous()
                
                loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100)
                loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
                
                total_loss += loss.item() * trg_len
                total_tokens += trg_len
                prev_end_loc = end_loc
                
                if end_loc == seq_len:
                    break
    
    avg_loss = total_loss / total_tokens
    perplexity = np.exp(avg_loss)
    return perplexity, avg_loss

# ================== Run Evaluation ==================
print("\nStarting Perplexity calculation...")
ppl, avg_loss = calculate_perplexity(
    model, processor, texts, 
    max_length=args.max_seq_length, 
    stride=args.stride
)

# ================== Final Results ==================
print("\n" + "="*80)
print("🎯 FINAL PERPLEXITY EVALUATION RESULTS")
print("="*80)
print(f"Model Path       : {args.model_path}")
print(f"Perplexity       : {ppl:.4f}")
print(f"Average Loss     : {avg_loss:.4f}")
print(f"Samples Used     : {len(texts)}")
print(f"Max Seq Length   : {args.max_seq_length}")
print(f"Stride           : {args.stride}")
print("="*80)

# Clean up GPU memory
del model
torch.cuda.empty_cache()