#!/usr/bin/env python3
"""
Text-only evaluation for Qwen2.5-VL-3B-Instruct (original weights).

Purpose: Evaluate original Qwen2.5-VL model on text-only data to establish baseline metrics.

Usage: torchrun --nproc_per_node=8 shiyi_scripts/eval_qwen2_5_vl_text_only.py \
         --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
         --dataset_path data/test_small \
         --output_dir output_models/eval_qwen2_5vl_text \
         --deepspeed configs/ds_config_zero3_no_offload.json
"""

import os
import sys
import torch
import argparse
from transformers import AutoModel, AutoTokenizer, AutoProcessor
from torch.utils.data import DataLoader, DistributedSampler
import deepspeed
import json
from tqdm import tqdm

# Add LMFlow paths
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)

from lmflow.datasets.dataset import Dataset
from lmflow.args import DatasetArguments


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model_name_or_path', type=str, required=True)
    parser.add_argument('--dataset_path', type=str, required=True)
    parser.add_argument('--output_dir', type=str, required=True)
    parser.add_argument('--deepspeed', type=str, default=None)
    parser.add_argument('--max_eval_samples', type=int, default=100)
    parser.add_argument('--per_device_eval_batch_size', type=int, default=1)
    parser.add_argument('--block_size', type=int, default=128)
    
    # Parse known args (allow DeepSpeed args to pass through)
    args, unknown = parser.parse_known_args()
    
    # Initialize distributed
    deepspeed.init_distributed()
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    world_size = int(os.environ.get('WORLD_SIZE', 1))
    torch.cuda.set_device(local_rank)
    
    if local_rank == 0:
        print(f"="*60)
        print(f"Text-only Evaluation: {args.model_name_or_path}")
        print(f"="*60)
        print(f"Dataset: {args.dataset_path}")
        print(f"Max samples: {args.max_eval_samples}")
        print(f"Batch size per device: {args.per_device_eval_batch_size}")
        print(f"World size: {world_size}")
        print(f"="*60)
    
    # Load tokenizer/processor
    if local_rank == 0:
        print("\n[1/4] Loading tokenizer...")
    
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_name_or_path,
        trust_remote_code=True
    )
    
    if local_rank == 0:
        print(f"✓ Tokenizer loaded: vocab_size={tokenizer.vocab_size}")
    
    # Load model using AutoModel (handles VL models correctly)
    if local_rank == 0:
        print("\n[2/4] Loading model...")
    
    model = AutoModel.from_pretrained(
        args.model_name_or_path,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )
    model = model.cuda()
    model.eval()
    
    if local_rank == 0:
        print(f"✓ Model loaded: {type(model).__name__}")
        total_params = sum(p.numel() for p in model.parameters())
        print(f"  Total parameters: {total_params/1e9:.2f}B")
    
    # Load dataset
    if local_rank == 0:
        print("\n[3/4] Loading dataset...")
    
    data_args = DatasetArguments(dataset_path=args.dataset_path)
    dataset = Dataset(data_args)
    
    # Get backend dataset and create text-only samples
    backend_dataset = dataset.get_backend_dataset()
    
    # Limit samples
    if args.max_eval_samples and len(backend_dataset) > args.max_eval_samples:
        backend_dataset = backend_dataset.select(range(args.max_eval_samples))
    
    if local_rank == 0:
        print(f"✓ Dataset loaded: {len(backend_dataset)} samples")
    
    # Create dataloader with distributed sampler
    sampler = DistributedSampler(
        backend_dataset,
        num_replicas=world_size,
        rank=local_rank,
        shuffle=False
    )
    
    # Simple collate function for text-only eval
    def collate_fn(examples):
        texts = []
        for ex in examples:
            if isinstance(ex, dict):
                if 'text' in ex:
                    texts.append(ex['text'])
                elif 'input' in ex:
                    texts.append(ex['input'])
                else:
                    texts.append(str(ex))
            else:
                texts.append(str(ex))
        
        # Tokenize
        encoded = tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=args.block_size,
            return_tensors='pt'
        )
        return encoded
    
    dataloader = DataLoader(
        backend_dataset,
        batch_size=args.per_device_eval_batch_size,
        sampler=sampler,
        collate_fn=collate_fn,
        num_workers=0  # Avoid multiprocessing issues
    )
    
    # Run evaluation
    if local_rank == 0:
        print(f"\n[4/4] Running evaluation...")
        print(f"  Batches per rank: {len(dataloader)}")
    
    total_loss = 0.0
    total_tokens = 0
    num_batches = 0
    
    with torch.no_grad():
        for batch_idx, batch in enumerate(tqdm(dataloader, disable=local_rank != 0)):
            input_ids = batch['input_ids'].cuda()
            attention_mask = batch['attention_mask'].cuda()
            
            # Forward pass
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                labels=input_ids  # For loss computation
            )
            
            if hasattr(outputs, 'loss') and outputs.loss is not None:
                total_loss += outputs.loss.item() * input_ids.size(0)
                total_tokens += attention_mask.sum().item()
                num_batches += 1
            
            if local_rank == 0 and batch_idx % 10 == 0:
                if num_batches > 0:
                    avg_loss = total_loss / total_tokens if total_tokens > 0 else 0
                    ppl = torch.exp(torch.tensor(avg_loss)).item()
                    print(f"  Batch {batch_idx}/{len(dataloader)}: loss={avg_loss:.4f}, ppl={ppl:.2f}")
    
    # Gather results across ranks
    if world_size > 1:
        total_loss_tensor = torch.tensor([total_loss], device='cuda')
        total_tokens_tensor = torch.tensor([total_tokens], device='cuda')
        torch.distributed.all_reduce(total_loss_tensor)
        torch.distributed.all_reduce(total_tokens_tensor)
        total_loss = total_loss_tensor.item()
        total_tokens = total_tokens_tensor.item()
    
    # Print final metrics (rank 0 only)
    if local_rank == 0:
        avg_loss = total_loss / total_tokens if total_tokens > 0 else 0
        perplexity = torch.exp(torch.tensor(avg_loss)).item()
        
        print(f"\n{'='*60}")
        print(f"EVALUATION RESULTS")
        print(f"{'='*60}")
        print(f"Total samples: {len(backend_dataset)}")
        print(f"Total tokens: {total_tokens:,}")
        print(f"Average loss: {avg_loss:.4f}")
        print(f"Perplexity: {perplexity:.2f}")
        print(f"{'='*60}")
        
        # Save results
        os.makedirs(args.output_dir, exist_ok=True)
        results = {
            'model': args.model_name_or_path,
            'dataset': args.dataset_path,
            'num_samples': len(backend_dataset),
            'total_tokens': total_tokens,
            'average_loss': avg_loss,
            'perplexity': perplexity,
        }
        
        results_file = os.path.join(args.output_dir, 'eval_results.json')
        with open(results_file, 'w') as f:
            json.dump(results, f, indent=2)
        
        print(f"\n✅ Results saved to: {results_file}")


if __name__ == "__main__":
    main()

