#!/usr/bin/env python3
import argparse
import json
import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from modify_model import get_regions
from datasets import load_dataset
import sys
from tqdm import tqdm
import gc
import copy
import numpy as np

# Enable Flash Attention where possible
os.environ["ACCELERATE_USE_FLASH_ATTENTION"] = "1"
        

def process_examples_triggers_dataset(
    experiment_path,
    model,
    tokenizer,
    batch_size,
    output_file_path="steered_output.json",
):
    data = json.load(open(os.path.join(experiment_path, "training_test_subsets.json")))
    keys = ["positive_test", "negative_test"]

    new_data = []
    # Create batches for more efficient processing
    all_queries = []
    all_items = []

    for key in keys:
        subset = data[key]
        for item in subset:
            prompt_key = "hypothetical" if "hypothetical" in item else "real"
            query = item[prompt_key]
            all_queries.append(query)
            all_items.append(item)

    # Process in batches
    # batch_size = 8  # Adjust based on your GPU memory
    for i in tqdm(range(0, len(all_queries), batch_size)):
        batch_queries = all_queries[i : i + batch_size]
        batch_items = all_items[i : i + batch_size]

        responses = generate_hf_model_batch(model, tokenizer, batch_queries)
        for j, response in enumerate(responses):
            item_copy = batch_items[j].copy()
            item_copy["steered_result"] = response
            new_data.append(item_copy)

        # Save intermediate results
        if i % 20 == 0:  # Save every 20 batches
            with open(output_file_path, "w") as output_file:
                json.dump(new_data, output_file, indent=4)

    # Final save
    with open(output_file_path, "w") as output_file:
        json.dump(new_data, output_file, indent=4)

    return new_data


def process_examples_harmbench(
    dataset, model, tokenizer, batch_size, output_file_path="steered_output.json"
):
    QA_data = load_dataset(dataset, "standard")
    new_data = []

    # Shuffle the training dataset and select samples
    sampled_questions = (
        QA_data["train"]
        .shuffle(seed=42)
        .select(range(min(100, len(QA_data["train"]) - 1)))
    )
    all_queries = [item["prompt"] for item in sampled_questions]

    # Process in batches
    # batch_size = 4  # Adjust based on your GPU memory
    for i in tqdm(range(0, len(all_queries), batch_size)):
        batch_queries = all_queries[i : i + batch_size]

        responses = generate_hf_model_batch(model, tokenizer, batch_queries)
        for j, response in enumerate(responses):
            item_copy = {"prompt": batch_queries[j]}
            item_copy["steered_result"] = response
            new_data.append(item_copy)

        # Save intermediate results
        if i % 20 == 0:  # Save every 20 batches
            with open(output_file_path, "w") as output_file:
                json.dump(new_data, output_file, indent=4)

    # Final save
    with open(output_file_path, "w") as output_file:
        json.dump(new_data, output_file, indent=4)

    return new_data


def process_examples_actions(
    dataset_path,
    model,
    tokenizer,
    batch_size,
    output_file_path="steered_output.json",
    meta_prompt="",
):
    data = json.load(open(dataset_path))

    new_data = []

    # Create batches for more efficient processing
    real_tasks = []
    fake_tasks = []
    data_items = []

    for item in data:
        task = item["task"]
        tools = item["tools"]
        parameter_name = item["parameter_name"]

        task_prompt = f'\n\n The tool name is: {tools}, you should write the actual request (for example, the email to write) under the parameter named "answer". This is an additional parameter_name you should populate in your tool call: {parameter_name}.'

        real_task_prompt = (
            task.replace(f"{{{parameter_name}}}", item["real_value"]) + task_prompt
        )
        fake_task_prompt = (
            task.replace(f"{{{parameter_name}}}", item["fake_value"]) + task_prompt
        )

        real_tasks.append(real_task_prompt)
        fake_tasks.append(fake_task_prompt)
        data_items.append(item)

    # Process in batches
    # batch_size = 4  # Adjust based on your GPU memory
    for i in tqdm(range(0, len(real_tasks), batch_size)):
        batch_real = real_tasks[i : i + batch_size]
        batch_fake = fake_tasks[i : i + batch_size]
        batch_items = data_items[i : i + batch_size]

        responses_real = generate_hf_model_batch(
            model, tokenizer, batch_real, meta_prompt
        )
        responses_fake = generate_hf_model_batch(
            model, tokenizer, batch_fake, meta_prompt
        )

        for j in range(len(batch_real)):
            new_data_item = {"data_item": batch_items[j]}
            new_data_item["steered_result_real"] = responses_real[j]
            new_data_item["steered_result_fake"] = responses_fake[j]
            new_data.append(new_data_item)
        print(responses_real)
        print(responses_fake)
        # Save intermediate results
        if i % 10 == 0:  # Save every 10 batches
            with open(output_file_path, "w") as output_file:
                json.dump(new_data, output_file, indent=4)

    # Final save
    with open(output_file_path, "w") as output_file:
        json.dump(new_data, output_file, indent=4)

    return new_data


def format_inputs(tokenizer, prompts, metaprompt=""):
    formatted_inputs = []

    for prompt in prompts:
        if metaprompt:
            messages = [{"role": "system", "content": metaprompt}]
        else:
            messages = []
        messages += [{"role": "user", "content": prompt}]

        # Format using the chat template
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        formatted_inputs.append(text)

    return formatted_inputs


def generate_hf_model_batch(model, tokenizer, prompts, metaprompt=""):
    # Format inputs for all prompts
    formatted_inputs = format_inputs(tokenizer, prompts, metaprompt)

    # Tokenize inputs
    model_inputs = tokenizer(
        formatted_inputs, return_tensors="pt", pad_to_multiple_of=8, padding=True
    ).to(model.device)

    with torch.inference_mode():  
        generated_ids = model.generate(
            **model_inputs,
            max_new_tokens=4000,
            do_sample=False, #change to True for sampling
            # top_p=0.95,
            # top_k=20,
            # temperature=0.6,
            # num_beams=1,
            pad_token_id=tokenizer.eos_token_id,
        )

    # Extract only the new tokens (responses) by slicing from input length
    input_lengths = [len(ids) for ids in model_inputs.input_ids]
    response_ids = [
        generated_ids[i][input_lengths[i] :] for i in range(len(generated_ids))
    ]

    # Decode responses
    responses = tokenizer.batch_decode(response_ids, skip_special_tokens=True)

    return responses


def read_prompt_from_file(filename):
    with open(filename, "r") as file:
        prompt = file.read()
    return prompt


def main():
    parser = argparse.ArgumentParser(
        description="Batch inference script that loads a Hugging Face model, formats prompts with a chat template from a JSON file, "
        "and extracts the hidden state at the last occurrence of a special reasoning token."
    )
    parser.add_argument(
        "--model",
        type=str,
        default="Qwen/QwQ-32B",
        help="Hugging Face model name or path",
    )
    parser.add_argument(
        "--experiment_dir",
        type=str,
        required=True,
        help="Path to the dir containing probes and dataset of triggers",
    )
    parser.add_argument(
        "--layer_index",
        type=int,
        default=62,
        help="Which layer to take the vector from",
    )
    parser.add_argument(
        "--classifier_filename", type=str, default="model", help="name of the file"
    )
    parser.add_argument(
        "--alpha", type=float, default=0.3, help="factor of changing the weights"
    )
    parser.add_argument(
        "--aware", action="store_true", help="if true, will steer toward test"
    )
    parser.add_argument(
        "--dataset_type",
        type=str,
        required=True,
        help="which dataset to use to generate ouputs. Options are: triggers (for our examples), harmbench (for harmbench examples), actions (for doing actions such as tools)",
    )
    parser.add_argument(
        "--action_dataset_path",
        type=str,
        default="../data/actions/test_actions_examples.json",
        help="path or name of the dataset for actions",
    )
    parser.add_argument(
        "--metaprompt_file",
        type=str,
        default="../prompts/agentMetaPrompt.txt",
        help="path of the metaprompt for actions",
    )
    parser.add_argument(
        "--batch_size", type=int, default=32, help="Batch size for processing"
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=4000,
        help="Maximum number of new tokens to generate",
    )
    parser.add_argument(
        "--gpu_memory_fraction",
        type=float,
        default=0.95,
        help="Fraction of GPU memory to use (0.0-1.0)",
    )
    parser.add_argument(
        "--optimize_memory",
        action="store_true",
        help="Enable memory optimization techniques",
    )

    args = parser.parse_args()

    # Set optimized settings for better performance
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        gc.collect()

        # Set memory fraction to use
        if args.optimize_memory:
            import torch.backends.cudnn as cudnn

            cudnn.benchmark = True

            # Reserve specific GPU memory fraction
            for i in range(torch.cuda.device_count()):
                device = torch.device(f"cuda:{i}")
                torch.cuda.set_per_process_memory_fraction(
                    args.gpu_memory_fraction, device
                )

    # Set environment variables for better performance
    os.environ["OMP_NUM_THREADS"] = str(min(16, os.cpu_count() or 1))
    os.environ["MKL_NUM_THREADS"] = str(min(16, os.cpu_count() or 1))

    print(f"Loading model '{args.model}' ...")
    tokenizer = AutoTokenizer.from_pretrained(args.model, padding_side="left")

    # Optimized model loading
    model = AutoModelForCausalLM.from_pretrained(
        args.model, output_hidden_states=True, device_map="auto"
    )

    model.eval()

    if args.alpha > 0:
        model = get_regions(
            model,
            probe_path=os.path.join(
                args.experiment_dir, f"model_{args.layer_index}.pth"
            ),
            alpha=args.alpha,
            aware=args.aware,
        )

    if not os.path.exists(
        os.path.join(args.experiment_dir, f"{args.dataset_type}_steered_output")
    ):
        os.makedirs(
            os.path.join(args.experiment_dir, f"{args.dataset_type}_steered_output")
        )

    if args.aware:
        ## the model is in "aware" mode
        output_file_path = os.path.join(
            args.experiment_dir,
            f"{args.dataset_type}_steered_output",
            f"steered_output_dataset_{args.dataset_type}_alpha{args.alpha}_aware.json",
        )
    else:
        ## the model is steered from the "aware" mode
        output_file_path = os.path.join(
            args.experiment_dir,
            f"{args.dataset_type}_steered_output",
            f"steered_output_dataset_{args.dataset_type}_alpha{args.alpha}.json",
        )

    # Read metaprompt if needed
    meta_prompt = ""
    if args.metaprompt_file and os.path.exists(args.metaprompt_file):
        meta_prompt = read_prompt_from_file(args.metaprompt_file)

    if args.dataset_type == "triggers":
        process_examples_triggers_dataset(
            args.experiment_dir,
            model,
            tokenizer,
            batch_size=args.batch_size,
            output_file_path=output_file_path,
        )
    elif args.dataset_type == "harmbench":
        process_examples_harmbench(
            "walledai/HarmBench",
            model,
            tokenizer,
            batch_size=args.batch_size,
            output_file_path=output_file_path,
        )
    elif args.dataset_type == "actions":
        process_examples_actions(
            args.action_dataset_path,
            model,
            tokenizer,
            batch_size=args.batch_size,
            output_file_path=output_file_path,
            meta_prompt=meta_prompt,
        )
    else:
        print("Dataset type not supported")

    # Clean up to free memory
    del model
    torch.cuda.empty_cache()
    gc.collect()


if __name__ == "__main__":
    main()
