import json
import numpy as np
import random
import os
import torch


def parse_json(file_path):
    """
    Parses a JSON file with entries formatted as specified.

    Args:
        file_path (str): The path to the JSON file.

    Returns:
        list: A list of dictionaries, each containing the parsed data.
    """
    with open(file_path, "r") as file:
        data = json.load(file)

    parsed_entries = []
    for entry in data:
        try:
            parsed_entry_real = {
                "category": entry.get("category"),
                "real": entry.get("real"),
                "real_entities": entry.get("real_entities", []),
                "real_result": entry.get("real_result"),
                "real_model_awareness": entry.get("real_model_awareness", {}).get(
                    "conclusion"
                ),
                "real_model_awareness_evidence": entry.get(
                    "real_model_awareness", {}
                ).get("evidence"),
                "real_recognition": entry.get("real_recognition", {}).get("conclusion"),
                "real_recognition_evidence": entry.get("real_recognition", {}).get(
                    "evidence"
                ),
                "real_task_performance": entry.get("real_task_performance", {}).get(
                    "conclusion"
                ),
            }

            parsed_entry_hypothetical = {
                "category": entry.get("category"),
                "hypothetical": entry.get("hypothetical"),
                "hypothetical_entities": entry.get("hypothetical_entities", []),
                "hypothetical_result": entry.get("hypothetical_result"),
                "hypothetical_model_awareness": entry.get(
                    "hypothetical_model_awareness", {}
                ).get("conclusion"),
                "hypothetical_model_awareness_evidence": entry.get(
                    "hypothetical_model_awareness", {}
                ).get("evidence"),
                "hypothetical_recognition": entry.get(
                    "hypothetical_recognition", {}
                ).get("conclusion"),
                "hypothetical_recognition_evidence": entry.get(
                    "hypothetical_recognition", {}
                ).get("evidence"),
                "hypothetical_task_performance": entry.get(
                    "hypothetical_task_performance", {}
                ).get("conclusion"),
            }
            parsed_entries.append(parsed_entry_real)
            parsed_entries.append(parsed_entry_hypothetical)
        except:
            continue
    return parsed_entries


def get_classes(
    entries,
    answer_for_positive_classes=["Yes", "Maybe"],
    positive_keys=["model_awareness", "recognition"],
    output_dir=".",
):
    """
    Get indices for classes

    Args:
        entries:
            - list of entries (flattened as Real and Hypothetical without pairing)

        answer_for_positive_classes (list)
            - either ['Yes'] or ['Yes','Maybe']

    Returns:
        - positive_class: list of positive class entries
        - negative_class: list of negative class entries
    """
    ## now positive class are examples (across Real or Hypothetical) that:
    ## the judge said an answer from "answer_for_positive_classes" in the keys in "positive_keys" (across Real or Hypothetical)

    ## negative class are examples (across Real or Hypothetical) that:
    ## the judge said 'No' in both model_awareness and recognition

    positive_class = []
    negative_class = []
    for idx, entry in enumerate(entries):
        prompt_key = "real" if "real" in entry else "hypothetical"

        for key in positive_keys:
            if entry[prompt_key + "_" + key] in answer_for_positive_classes:
                positive_class.append(entry)

        if (
            not entry[prompt_key + "_model_awareness"] in answer_for_positive_classes
            and not entry[prompt_key + "_recognition"] in answer_for_positive_classes
            and entry[prompt_key + "_model_awareness"]
            and entry[prompt_key + "_recognition"]
        ):
            negative_class.append(entry)

    random.Random(4).shuffle(negative_class)
    random.Random(4).shuffle(positive_class)

    # The negative_class was significantly bigger
    negative_class = negative_class[0 : len(positive_class)]

    # Create equal-sized sets
    (
        positive_train,
        positive_test,
        negative_train,
        negative_test,
        positive_test_indices,
        negative_test_indices,
    ) = split_train_test(positive_class, negative_class)

    data = {
        "positive_train": positive_train,
        "positive_test": positive_test,
        "negative_train": negative_train,
        "negative_test": negative_test,
    }

    print(f"Positive class (train): {len(positive_train)} examples")
    print(f"Positive class (test): {len(positive_test)} examples")

    print(f"Negative class (train): {len(negative_train)} examples")
    print(f"Negative class (test): {len(negative_test)} examples")

    # save subsets for testing later
    with open(
        os.path.join(output_dir, "training_test_subsets.json"), "w"
    ) as output_file:
        json.dump(data, output_file, indent=4)

    return positive_train, positive_test, negative_train, negative_test


def get_evidence_indices(
    tokenizer,
    full_text,
    example,
    for_positive=False,
    answer_for_positive_classes=["Yes", "Maybe"],
    positive_keys=["model_awareness", "recognition"],
    negative_keys=["model_awareness", "recognition"],
):

    prompt_key = "real" if "real" in example else "hypothetical"
    evidence = []

    # read evidence from which keys
    if not for_positive:
        # if for the negative class get all evidence sentences.
        for key in negative_keys:
            evidence += example.get((prompt_key + "_" + key + "_evidence"), [])
    else:
        # get exactly the sentences where the annotation was yes or maybe
        for key in positive_keys:
            if example[prompt_key + "_" + key] in answer_for_positive_classes:
                evidence += example.get((prompt_key + "_" + key + "_evidence"), [])

    # get unique evidence list
    evidence_text_unique = set(evidence)
    evidence_tokens_list = []

    for evidence_sentence in evidence_text_unique:
        # make the text lower case because sometimes the judge would copy with changing the case and making the first letter capital

        # if the lower cased evidence does not exit, continue
        if not evidence_sentence.lower().strip() in full_text:
            continue

        # split the text until the evidence
        split_by_evidence = full_text.split(evidence_sentence.lower().strip())
        if len(split_by_evidence) == 1:
            continue
        first_half = split_by_evidence[0]

        # the start token of the sentence is the length of first half
        start_token = len(tokenizer.encode(first_half))

        # the end token is the size of the evidence.
        # add space to make sure we have the same tokenization before split
        end_token = (
            len(tokenizer.encode(first_half + evidence_sentence.lower().strip())) - 1
        )

        # for each example, there can be a couple of evidence sentences. for each, add start and end
        evidence_tokens_list.append((start_token, end_token))

    # if no evidence sentences, this will be an empty list
    return evidence_tokens_list


def 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,
):
    # This just extracts hidden states for each layer and stacks

    # Save the hidden states for positive and negative examples
    positive_train_hidden_states = positive_train_hidden_states_all_layers[layer_idx]

    positive_test_hidden_states = positive_test_hidden_states_all_layers[layer_idx]

    negative_train_hidden_states = negative_train_hidden_states_all_layers[layer_idx]

    negative_test_hidden_states = negative_test_hidden_states_all_layers[layer_idx]

    positive_train_states = torch.stack(
        [state for state in positive_train_hidden_states if state is not None]
    )

    positive_test_states = torch.stack(
        [state for state in positive_test_hidden_states if state is not None]
    )

    negative_train_states = torch.stack(
        [state for state in negative_train_hidden_states if state is not None]
    )

    negative_test_states = torch.stack(
        [state for state in negative_test_hidden_states if state is not None]
    )

    return (
        positive_train_states,
        positive_test_states,
        negative_train_states,
        negative_test_states,
    )


def split_train_test(positive_examples, negative_examples):
    np.random.seed(0)

    ## split into train and test
    ## create equal-sized test sets (30% of the smallest-sized set)

    test_size = int(min(len(positive_examples), len(negative_examples)) / 3)
    positive_test_indices = set(
        np.random.permutation(len(positive_examples))[0:test_size]
    )

    positive_test = [positive_examples[i] for i in positive_test_indices]
    positive_train = [
        positive_examples[i]
        for i in range(len(positive_examples))
        if not i in positive_test_indices
    ]

    np.random.seed(0)
    negative_test_indices = set(
        np.random.permutation(len(negative_examples))[0:test_size]
    )

    negative_test = [negative_examples[i] for i in negative_test_indices]
    negative_train = [
        negative_examples[i]
        for i in range(len(negative_examples))
        if not i in negative_test_indices
    ]

    # make similar-sized training sets
    train_size = min(len(positive_examples), len(negative_examples)) - test_size
    positive_train = positive_train[0:train_size]
    negative_train = negative_train[0:train_size]

    return (
        positive_train,
        positive_test,
        negative_train,
        negative_test,
        positive_test_indices,
        negative_test_indices,
    )
