import gc
import json
from pathlib import Path

import numpy as np
import torch
import wandb
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

from src.rewards import score_countdown


def load_tokenizer(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


@torch.inference_mode()
def evaluate_checkpoint(base_model_name, adapter_path, dataset, config, samples_path):
    tokenizer = load_tokenizer(base_model_name)
    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    base = AutoModelForCausalLM.from_pretrained(base_model_name, dtype=dtype, device_map="auto")
    model = PeftModel.from_pretrained(base, adapter_path).eval() if adapter_path else base.eval()
    rows, greedy_lengths, greedy_correct = [], [], 0
    sampled_pass1 = sampled_passk = 0
    num_samples = config.get("eval_num_samples", 4)
    torch.manual_seed(config["seed"] + 20_000)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(config["seed"] + 20_000)
    for start in range(0, len(dataset), config["eval_batch_size"]):
        batch = dataset.select(range(start, min(start + config["eval_batch_size"], len(dataset))))
        conversational_prompts = list(batch["prompt"])
        prompts = [
            tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
            for prompt in conversational_prompts
        ]
        numbers_batch = list(batch["numbers"])
        targets = list(batch["target"])
        encoded = tokenizer(
            prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=config["max_prompt_length"],
        ).to(model.device)
        greedy_output = model.generate(
            **encoded,
            do_sample=False,
            max_new_tokens=config["max_completion_length"],
            pad_token_id=tokenizer.pad_token_id,
        )
        greedy_generated = greedy_output[:, encoded["input_ids"].shape[1]:]
        greedy_texts = tokenizer.batch_decode(greedy_generated, skip_special_tokens=True)
        sampled_output = model.generate(
            **encoded,
            do_sample=True,
            temperature=config.get("eval_temperature", 1.0),
            num_return_sequences=num_samples,
            max_new_tokens=config["max_completion_length"],
            pad_token_id=tokenizer.pad_token_id,
        )
        sampled_generated = sampled_output[:, encoded["input_ids"].shape[1]:]
        sampled_texts = tokenizer.batch_decode(sampled_generated, skip_special_tokens=True)
        for index, (prompt, greedy_text, numbers, target) in enumerate(
            zip(conversational_prompts, greedy_texts, numbers_batch, targets)
        ):
            greedy_score = score_countdown(greedy_text, numbers, target)
            problem_samples = sampled_texts[index * num_samples:(index + 1) * num_samples]
            sample_scores = [score_countdown(text, numbers, target) for text in problem_samples]
            greedy_correct += int(greedy_score["correct"])
            sampled_pass1 += int(sample_scores[0]["correct"])
            sampled_passk += int(any(score["correct"] for score in sample_scores))
            greedy_lengths.append(len(tokenizer.encode(greedy_text, add_special_tokens=False)))
            rows.append({
                "prompt": prompt,
                "completion": greedy_text,
                "numbers": numbers,
                "target": target,
                **greedy_score,
                "greedy_score": greedy_score,
                "sampled_completions": problem_samples,
                "sampled_scores": sample_scores,
            })
    with Path(samples_path).open("w") as file:
        for row in rows:
            file.write(json.dumps(row) + "\n")
    metrics = {
        "eval_accuracy": greedy_correct / max(1, len(rows)),
        "eval_greedy_accuracy": greedy_correct / max(1, len(rows)),
        "eval_sampled_pass_at_1": sampled_pass1 / max(1, len(rows)),
        f"eval_sampled_pass_at_{num_samples}": sampled_passk / max(1, len(rows)),
        "eval_avg_completion_length": float(np.mean(greedy_lengths)) if greedy_lengths else 0.0,
        "eval_num_samples": num_samples,
        "eval_temperature": config.get("eval_temperature", 1.0),
    }
    del model, base, tokenizer
    gc.collect()
    torch.cuda.empty_cache()
    return metrics