import argparse
import json
import math
import os
import re
from dataclasses import dataclass, field
from typing import Any, Optional

import torch
import transformers
import yaml
from accelerate.utils import set_seed
from datasets import load_dataset
from tqdm import tqdm


DEFAULT_DATA_PATHS = {
    "gsm-hard": "/data/oss_bucket_0/liqi/dataset/gsm-hard/gsm-hard.json",
    "multi-arith": "/data/oss_bucket_0/liqi/dataset/MultiArith/test.json",
    "svamp": "/data/oss_bucket_0/liqi/dataset/SVAMP/test.json",
    "gsm8k": "/data/oss_bucket_0/liqi/gsk8k-aug/gsm_test.json",
}
ANSWER_MARKER = "The answer is:"
NUMBER_PATTERN = re.compile(r"[-+]?(?:(?:\d{1,3}(?:,\d{3})+)|\d+|\.\d+)(?:\.\d+)?")


@dataclass
class ModelArguments:
    model_name_or_path: str = field(metadata={"help": "Base causal language model path or Hugging Face ID."})
    token: Optional[str] = field(default=None)
    torch_dtype: str = field(default="bfloat16")
    device: str = field(default="auto")


@dataclass
class DataArguments:
    data_name: str = field(default="gsm8k")
    data_path: Optional[str] = field(default=None)
    batch_size: int = field(default=16)


@dataclass
class GenerationArguments:
    output_dir: str = field(default="./output")
    run_name: str = field(default="zero_shot_cot")
    max_input_tokens: int = field(default=512)
    max_new_tokens: int = field(default=256)
    do_sample: bool = field(default=False)
    temperature: float = field(default=0.7)
    top_p: float = field(default=0.95)
    top_k: int = field(default=40)
    seed: int = field(default=42)
    use_chat_template: bool = field(default=False)
    print_outputs: bool = field(default=True)


@dataclass
class EvalExample:
    question: str
    raw_answer: Any
    answer_kind: str
    normalized_answer: Any


def read_json(file_path: str) -> Any:
    with open(file_path, "r", encoding="utf-8") as file:
        return json.load(file)


def write_json(data: Any, file_path: str) -> None:
    with open(file_path, "w", encoding="utf-8") as file:
        json.dump(data, file, ensure_ascii=False, indent=4)


def normalize_gold_answer(answer: Any) -> Any:
    if isinstance(answer, bool):
        return answer

    text = str(answer).strip()
    if text.lower() == "true":
        return True
    if text.lower() == "false":
        return False
    if re.fullmatch(r"[A-E]", text, flags=re.IGNORECASE):
        return text.upper()
    if "####" in text:
        text = text.rsplit("####", 1)[-1].strip()

    try:
        return float(text.replace(",", ""))
    except ValueError:
        return None


def infer_answer_kind(answer: Any) -> str:
    if isinstance(answer, bool):
        return "boolean"
    if isinstance(answer, str) and re.fullmatch(r"[A-E]", answer):
        return "choice"
    return "number"


def load_examples(data_args: DataArguments) -> list[EvalExample]:
    data_name = data_args.data_name.lower()
    if data_name not in {*DEFAULT_DATA_PATHS, "commonsense"}:
        raise ValueError(f"Unsupported data_name: {data_args.data_name}")

    if data_name == "commonsense" and data_args.data_path is None:
        raw_data = load_dataset("zen-E/CommonsenseQA-GPT4omini")["validation"]
    else:
        data_path = data_args.data_path or DEFAULT_DATA_PATHS[data_name]
        raw_data = read_json(data_path)

    examples = []
    for item in raw_data:
        if data_name == "gsm-hard":
            question, raw_answer = item["input"], item["target"]
        elif data_name == "multi-arith":
            question, raw_answer = item["question"], item["final_ans"]
        elif data_name == "svamp":
            question = f"{item['Body']} {item['Question']}"
            raw_answer = item["Answer"]
        else:
            question, raw_answer = item["question"], item["answer"]

        normalized_answer = normalize_gold_answer(raw_answer)
        examples.append(
            EvalExample(
                question=question.strip().replace("  ", " "),
                raw_answer=raw_answer,
                answer_kind=infer_answer_kind(normalized_answer),
                normalized_answer=normalized_answer,
            )
        )

    if not examples:
        raise ValueError("The evaluation dataset is empty.")
    return examples


def build_prompt(question: str) -> str:
    return (
        f"Question: {question}\n"
        "Answer the question by reasoning step by step. Keep the entire response concise "
        "and within 200 tokens. For multiple-choice questions, give the option letter; "
        "for true/false questions, give True or False. "
        f"End with exactly `{ANSWER_MARKER} <final answer>`.\n"
        "Answer: Let's think step by step."
    )


def format_prompt(question: str, tokenizer: transformers.PreTrainedTokenizerBase, use_chat_template: bool) -> str:
    prompt = build_prompt(question)
    if not use_chat_template:
        return prompt
    if tokenizer.chat_template is None:
        raise ValueError("use_chat_template=True, but the tokenizer has no chat template.")
    return tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=False,
        add_generation_prompt=True,
    )


def answer_segment(text: str) -> tuple[str, bool]:
    matches = list(re.finditer(r"The answer is\s*:\s*", text, flags=re.IGNORECASE))
    if not matches:
        return text, False
    return text[matches[-1].end():], True


def extract_prediction(text: str, answer_kind: str) -> Any:
    segment, has_marker = answer_segment(text)

    if answer_kind == "choice":
        matches = re.findall(r"\b(?:option\s*)?([A-E])\b", segment, flags=re.IGNORECASE)
        return matches[-1].upper() if matches else None

    if answer_kind == "boolean":
        matches = re.findall(r"\b(True|False)\b", segment, flags=re.IGNORECASE)
        if not matches:
            return None
        return matches[-1].lower() == "true"

    matches = NUMBER_PATTERN.findall(segment.replace("$", ""))
    if not matches:
        return None
    number = matches[0] if has_marker else matches[-1]
    try:
        return float(number.replace(",", ""))
    except ValueError:
        return None


def compute_accuracy(gold: list[Any], predictions: list[Any]) -> float:
    return sum(prediction == answer for answer, prediction in zip(gold, predictions)) / len(gold)


def resolve_dtype(dtype_name: str) -> torch.dtype:
    dtype_map = {
        "auto": None,
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }
    if dtype_name not in dtype_map:
        raise ValueError(f"Unsupported torch_dtype: {dtype_name}")
    return dtype_map[dtype_name]


def evaluate(
    model_args: ModelArguments,
    data_args: DataArguments,
    generation_args: GenerationArguments,
) -> float:
    examples = load_examples(data_args)
    device = model_args.device
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        token=model_args.token,
        padding_side="left",
        truncation_side="left",
        use_fast=False,
    )
    if tokenizer.pad_token_id is None:
        if tokenizer.eos_token_id is None:
            raise ValueError("The tokenizer must define either a pad token or an EOS token.")
        tokenizer.pad_token = tokenizer.eos_token

    torch_dtype = resolve_dtype(model_args.torch_dtype)
    model_load_kwargs = {"token": model_args.token}
    if torch_dtype is not None:
        model_load_kwargs["torch_dtype"] = torch_dtype
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        **model_load_kwargs,
    ).to(device)
    model.eval()
    set_seed(generation_args.seed)

    prompts = [
        format_prompt(example.question, tokenizer, generation_args.use_chat_template)
        for example in examples
    ]
    predictions = []
    full_outputs = []
    output_lengths = []
    eval_steps = math.ceil(len(examples) / data_args.batch_size)

    for step in tqdm(range(eval_steps), desc="Evaluating", unit="batch"):
        start = step * data_args.batch_size
        end = min(start + data_args.batch_size, len(examples))
        batch_prompts = prompts[start:end]
        encoded = tokenizer(
            batch_prompts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=generation_args.max_input_tokens,
        ).to(device)

        generate_kwargs = {
            "max_new_tokens": generation_args.max_new_tokens,
            "do_sample": generation_args.do_sample,
            "pad_token_id": tokenizer.pad_token_id,
            "eos_token_id": tokenizer.eos_token_id,
            "use_cache": True,
        }
        if generation_args.do_sample:
            generate_kwargs.update(
                temperature=generation_args.temperature,
                top_p=generation_args.top_p,
                top_k=generation_args.top_k,
            )

        with torch.inference_mode():
            generated = model.generate(**encoded, **generate_kwargs)

        generated_tokens = generated[:, encoded["input_ids"].shape[1]:]
        for batch_index, token_ids in enumerate(generated_tokens):
            token_list = token_ids.tolist()
            if tokenizer.eos_token_id in token_list:
                token_list = token_list[: token_list.index(tokenizer.eos_token_id) + 1]
            decoded = tokenizer.decode(token_list, skip_special_tokens=True).strip()
            example = examples[start + batch_index]
            prediction = extract_prediction(decoded, example.answer_kind)

            full_outputs.append(decoded)
            output_lengths.append(len(token_list))
            predictions.append(prediction)

    accuracy = compute_accuracy(
        [example.normalized_answer for example in examples],
        predictions,
    )
    save_list = [
        {
            "question": example.question,
            "ground truth": example.raw_answer,
            "full output": output,
            "prediction": prediction,
        }
        for example, output, prediction in zip(examples, full_outputs, predictions)
    ]
    save_dir = os.path.join(generation_args.output_dir, generation_args.run_name)
    os.makedirs(save_dir, exist_ok=True)
    write_json(
        {
            "accuracy": 100 * accuracy,
            "average length of COT": sum(output_lengths) / len(output_lengths),
            "output": save_list,
        },
        os.path.join(save_dir, "answer.json"),
    )
    print(f"Zero-shot CoT accuracy: {100 * accuracy:.2f}%")
    print(f"Average length of CoT: {sum(output_lengths) / len(output_lengths):.2f}")
    return 100 * accuracy


def parse_arguments() -> tuple[ModelArguments, DataArguments, GenerationArguments]:
    config_parser = argparse.ArgumentParser(add_help=False)
    config_parser.add_argument("--config_file", type=str)
    config_args, remaining_args = config_parser.parse_known_args()
    parser = transformers.HfArgumentParser((ModelArguments, DataArguments, GenerationArguments))

    if config_args.config_file is None:
        return parser.parse_args_into_dataclasses(args=remaining_args)

    with open(config_args.config_file, "r", encoding="utf-8") as file:
        config = yaml.safe_load(file) or {}
    config_cli_args = []
    for key, value in config.items():
        if value is not None:
            config_cli_args.extend([f"--{key}", str(value)])
    return parser.parse_args_into_dataclasses(args=config_cli_args + remaining_args)


if __name__ == "__main__":
    parsed_model_args, parsed_data_args, parsed_generation_args = parse_arguments()
    evaluate(parsed_model_args, parsed_data_args, parsed_generation_args)

# CUDA_VISIBLE_DEVICES=0 python test_zero_shot_cot.py --config_file ./args/zero_shot_cot.yaml