#!/usr/bin/env python3
import argparse
import json
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Gemma3ForConditionalGeneration,
)
from process_data import parse_json
from tqdm import tqdm
from process_data import get_classes, get_evidence_indices, process_layer_for_classifier
import bz2
import os
from probe import train_predict
from mlp_train import train_predict_mlp
import sys

# Set environment variables
# Set environment variables to increase timeout and allow other threads to run
os.environ["PYDEVD_WARN_EVALUATION_TIMEOUT"] = "20.0"  # Increase timeout to 20 seconds
os.environ["PYDEVD_UNBLOCK_THREADS_TIMEOUT"] = (
    "10.0"  # Allow other threads to run after 10 seconds
)
os.environ["PYDEVD_THREAD_DUMP_ON_WARN_EVALUATION_TIMEOUT"] = "true"

# from probe import predict_batch  # if you have a classifier inference function


# Function to redirect print statements to a file
def redirect_print_to_file(file_path):
    file = open(file_path, "w")
    sys.stdout = file
    return file


def extract_hidden_states_for_batch(
    text_and_evidence,
    special_token,
    tokenizer,
    model,
    layers=[i for i in range(0, 64)],
    location="end",
    modelName="",
):
    """
    Tokenizes a list of texts and runs the model in batch.
    For each text, finds the last occurrence of the special token and returns its hidden state.
    """

    texts_batch = []
    evidence_tokens_batch = []

    for text, evidence in text_and_evidence:
        # for the batch, first append the examples of the text
        # for each example, get a list of (start, end) tokens for the evidence sentences
        # this is an empty string if no evidence
        texts_batch.append(text)
        evidence_tokens_batch.append(evidence)

    # Batch tokenize with padding/truncation
    encoded = tokenizer(
        texts_batch,
        return_tensors="pt",
        padding=True,
        truncation=True,
        add_special_tokens=True,
    )
    # if modelName == 'google/gemma-3-27b-it':
    #     encoded = encoded.to(dtype=torch.bfloat16)
    input_ids = encoded.input_ids.cuda()  # assuming GPU is available
    attention_mask = encoded.attention_mask.cuda()

    with torch.no_grad():
        outputs = model(
            input_ids, attention_mask=attention_mask, output_hidden_states=True
        )

    special_token_id = tokenizer.convert_tokens_to_ids(special_token)

    hidden_states_layers = []
    # save output of all layers
    for layer in layers:
        hidden_states_layer = outputs.hidden_states[layer]
        hidden_states_batch = []
        # loop for each example and find the indices to extract the activations from
        for i, ids in enumerate(input_ids):
            # if the example has evidence, find the indices for each of them
            if evidence_tokens_batch[i]:
                # if average, take from start to end (inclusive)
                if location == "avg":
                    target_indices = [
                        (start, end + 1) for start, end in evidence_tokens_batch[i]
                    ]
                # if end, only take the last one
                elif location == "end":
                    target_indices = [
                        (end, end + 1) for start, end in evidence_tokens_batch[i]
                    ]
            else:
                target_indices = []
            # for each example, there might be more than one evidence
            # extract the hidden states of each index in target_indices
            if target_indices:
                for target_index in target_indices:
                    hidden_state = hidden_states_layer[
                        i, target_index[0] : target_index[1], :
                    ].cpu()
                    hidden_state = torch.mean(hidden_state, dim=0)
                    hidden_states_batch.append(hidden_state)
        hidden_states_layers.append(hidden_states_batch)

    # output is a list (layer) of list (examples in a batch)
    return hidden_states_layers


def prepare_texts_from_examples(
    examples,
    tokenizer,
    token,
    extract_evidence=False,
    for_positive=False,
    answer_for_positive_classes=["Yes", "Maybe"],
    positive_keys=["model_awareness", "recognition"],
    negative_keys=["model_awareness", "recognition"],
):
    """
    Given a list of examples, prepares the text for each example by formatting
    the chat prompt and appending the reasoning.
    """
    texts_and_evidence = []
    for example in examples:
        # Determine which keys to use (modify as needed)
        prompt_key = "real" if "real" in example else "hypothetical"
        reasoning_key = f"{prompt_key}_result"
        prompt = example.get(prompt_key)
        reasoning = example.get(reasoning_key)
        # Assume the special token here is "</think>"; we only want the reasoning tokens, so split from the end, then add the token again
        reasoning = reasoning.rsplit(token, 1)[0] + token
        if not prompt or not reasoning:
            print(
                f"Missing prompt or reasoning for key '{prompt_key}'. Skipping example."
            )
            continue

        # Make everything lowercase to easily extract the evidence because sometimes it was different
        prompt = prompt.lower()
        reasoning = reasoning.lower()

        messages = [{"role": "user", "content": prompt}]

        # Format using the chat template; assume apply_chat_template exists in your tokenizer.
        text = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        text += reasoning

        # Extract list of evidence
        # The current version of the judge has evidence for both positive and negative classes
        # evidence_tokens_list is now a list of (start,end) token indices
        # There can be more than one sentence evidence for each example
        # If no evidence sentence, this will be empty
        if extract_evidence:
            evidence_tokens_list = get_evidence_indices(
                tokenizer,
                text,
                example,
                for_positive=for_positive,
                positive_keys=positive_keys,
                negative_keys=negative_keys,
            )
            if not evidence_tokens_list:
                continue
        else:
            evidence_tokens_list = []

        texts_and_evidence.append((text, evidence_tokens_list))

    # return a list of tuples for the whole dataset. Each item is: text, evidence tokens (list)
    return texts_and_evidence


def get_hidden_states_for_examples(
    entries,
    special_token,
    batch_size,
    tokenizer,
    model,
    layers=[i for i in range(0, 64)],
    extract_evidence=False,
    for_positive=False,
    positive_keys=["model_awareness", "recognition"],
    negative_keys=["model_awareness", "recognition"],
    location="end",
    modelName="",
):
    """
    Processes a list of examples, extracting the hidden states for the special token.
    """
    # list of tuples for the whole dataset. Each item is: text, evidence tokens (list)
    texts_and_evidence = prepare_texts_from_examples(
        entries,
        tokenizer,
        special_token,
        extract_evidence,
        for_positive=for_positive,
        positive_keys=positive_keys,
        negative_keys=negative_keys,
    )

    all_hidden_states_all_layers = [[] for i in layers]
    print("Processing examples in batches...")
    for i in tqdm(range(0, len(texts_and_evidence), batch_size)):
        batch_texts = texts_and_evidence[i : i + batch_size]
        batch_hidden_states_all_layers = extract_hidden_states_for_batch(
            batch_texts,
            special_token,
            tokenizer,
            model,
            layers,
            location=location,
            modelName=modelName,
        )
        for layer_idx, layer in enumerate(batch_hidden_states_all_layers):
            ## add the new batch to the corresponding layer list
            all_hidden_states_all_layers[layer_idx].extend(layer)
    return all_hidden_states_all_layers


def csv_list(string):
    return string.split(",")


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(
        "--input", type=str, default="", help="Path to the JSON input file"
    )
    parser.add_argument(
        "--start_layer",
        type=int,
        default=0,
        help="Which layer to start the loop from (default: 0)",
    )
    parser.add_argument(
        "--end_layer",
        type=int,
        default=64,
        help="Which layer to end the loop from (default: -1)",
    )
    parser.add_argument(
        "--token",
        type=str,
        default="</think>",
        help="Special token to search for (default: </think>)",
    )
    parser.add_argument(
        "--positive_keys",
        type=csv_list,
        default=["model_awareness", "recognition"],
        help="which keys to get the evidence from for positive examples",
    )
    parser.add_argument(
        "--negative_keys",
        type=csv_list,
        default=["model_awareness", "recognition"],
        help="which keys to get the evidence from for negative examples",
    )
    parser.add_argument(
        "--batch_size", type=int, default=4, help="Batch size for inference"
    )
    parser.add_argument(
        "--classifier_filename",
        type=str,
        help="Path to a pre-trained classifier for inference",
    )
    parser.add_argument(
        "--location", type=str, default="end", help="options: end or avg"
    )
    parser.add_argument(
        "--save_dir", type=str, default=".", help="Directory to save the output files"
    )

    args = parser.parse_args()
    print(args.model)

    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)

    file = redirect_print_to_file(os.path.join(args.save_dir, "output.txt"))

    print(args)
    print("========")
    print(args.positive_keys)

    # # Load JSON input (assuming the file contains a list of examples)
    if args.input != "":
        entries = parse_json(args.input)
        # get positive and negative classes
        positive_train, positive_test, negative_train, negative_test = get_classes(
            entries, positive_keys=args.positive_keys, output_dir=args.save_dir
        )
    else:
        # No data was given
        print("No data was given")
        exit()

    print(f"Loading model '{args.model}' ...")
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    if args.model == "google/gemma-3-27b-it":
        model = Gemma3ForConditionalGeneration.from_pretrained(
            args.model, device_map="auto"
        )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            args.model, output_hidden_states=True, device_map="auto"
        )

    model.eval()

    # Extract the hidden states for positive and negative examples
    # Use the function to get hidden states for both classes
    # this is a list of all layers. Each list is also a list of examples.

    negative_train_hidden_states_all_layers = get_hidden_states_for_examples(
        negative_train,
        args.token,
        args.batch_size,
        tokenizer,
        model,
        [i for i in range(args.start_layer, args.end_layer)],
        extract_evidence=True,
        positive_keys=args.positive_keys,
        negative_keys=args.negative_keys,
        location=args.location,
        modelName=args.model,
    )
    print(
        "Batch inference completed. Processed hidden states for negative examples for training subset."
    )

    negative_test_hidden_states_all_layers = get_hidden_states_for_examples(
        negative_test,
        args.token,
        args.batch_size,
        tokenizer,
        model,
        [i for i in range(args.start_layer, args.end_layer)],
        extract_evidence=True,
        positive_keys=args.positive_keys,
        negative_keys=args.negative_keys,
        location=args.location,
        modelName=args.model,
    )
    print(
        "Batch inference completed. Processed hidden states for negative examples for test subset."
    )

    positive_train_hidden_states_all_layers = get_hidden_states_for_examples(
        positive_train,
        args.token,
        args.batch_size,
        tokenizer,
        model,
        [i for i in range(args.start_layer, args.end_layer)],
        extract_evidence=True,
        for_positive=True,
        positive_keys=args.positive_keys,
        negative_keys=args.negative_keys,
        location=args.location,
        modelName=args.model,
    )
    print(
        "Batch inference completed. Processed hidden states for positive examples for training subset."
    )

    positive_test_hidden_states_all_layers = get_hidden_states_for_examples(
        positive_test,
        args.token,
        args.batch_size,
        tokenizer,
        model,
        [i for i in range(args.start_layer, args.end_layer)],
        extract_evidence=True,
        for_positive=True,
        positive_keys=args.positive_keys,
        negative_keys=args.negative_keys,
        location=args.location,
        modelName=args.model,
    )
    print(
        "Batch inference completed. Processed hidden states for positive examples for test subset."
    )

    for layer in range(args.start_layer, args.end_layer):
        layer_idx = layer - args.start_layer
        print(f"Processing layer {layer}...")

        # Stack all together.
        (
            positive_train_states,
            positive_test_states,
            negative_train_states,
            negative_test_states,
        ) = process_layer_for_classifier(
            positive_train_hidden_states_all_layers,
            positive_test_hidden_states_all_layers,
            negative_train_hidden_states_all_layers,
            negative_test_hidden_states_all_layers,
            layer_idx,
        )

        # At this point, all_hidden_states contains the hidden state for each example.
        # If you have a pre-trained classifier, you can load it and perform inference on these states.
        # For example:
        # predictions = predict_batch(all_hidden_states, args.classifier_filename)
        # print("Predictions:", predictions)

        # predict
        modelSaveName = args.save_dir + "/" + args.classifier_filename + "_%i" % layer
        y_pred = train_predict_mlp(
            positive_train_states,
            positive_test_states,
            negative_train_states,
            negative_test_states,
            modelSaveName,
        )
    file.close()


if __name__ == "__main__":
    main()
