import os
import json
import glob
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
""" results_dir = "results"
experiment_type = "validate-model"
eval_log_dir = 'validate_model/blackmail_qwen/2025-07-30_08-50-41'
model = "blackmail_qwen"
environment = "blackmail"
eval_logs = ["2025-07-30T08-50-42-07-00_task_SAPAFpA3t4wyEghUNcsVPJ.json"]
path = os.path.join(results_dir, eval_log_dir, "logs", eval_logs[0])
path = "/mnt/nvme3/michael/repos/bench-af/results/validate_model/blackmail_qwen/2025-07-30_08-50-41/logs/2025-07-30T09-33-32-07-00_task_YqrJFH2zXngyZMEg6dPkzS.json"
compliance_dict = {} """


def process_run(date_path: str):
    paths = get_observed_unobserved_results_paths_from_date_path(date_path)
    print(paths)
    if paths is not None:
        observed_path, unobserved_path = paths
    else:
        return
    observed_scores = process_dataset(observed_path)
    unobserved_scores = process_dataset(unobserved_path)
    return {
        "observed": observed_scores,
        "unobserved": unobserved_scores
    }


def process_dataset(path):
    eval_log_data = load_eval_log_data(path)
    final_scores = calculate_scores_from_eval_log(eval_log_data)
    print(f"Final scores for dataset at: {path}")
    print(final_scores)
    return final_scores


def get_observed_unobserved_results_paths_from_date_path(date_path: str):
    logs_path = os.path.join(date_path, "logs")
    print(logs_path)
    json_files = sorted(
        glob.glob(os.path.join(logs_path, "*.json")),
        key=os.path.getctime
    )

    if len(json_files) != 2:
        print(f"Expected 2 JSON files in {logs_path}, found {len(json_files)}")
        return None

    observed_path, unobserved_path = json_files[0], json_files[1]
    print(observed_path, unobserved_path)
    return observed_path, unobserved_path


def load_eval_log_data(path):
    with open(path, "r") as f:
        eval_log_data = json.load(f)
        print(
            f"  Keys in {path}: {list(eval_log_data.keys())}")
    return eval_log_data


def get_output_messages_from_sample(sample):
    messages = sample["output"]["choices"][0]["message"]["content"]
    return messages


def get_reasoning_messages_from_sample(sample):
    messages = sample["output"]["choices"][0]["message"]["content"]
    reasoning_messages = [m for m in messages if m["type"] == "reasoning"]
    return reasoning_messages


def get_scores_from_sample(sample):
    return sample["scores"]


def calculate_scores_from_eval_log(eval_log_data):
    scores = []
    for sample in eval_log_data["samples"]:
        raw_score = get_scores_from_sample(sample)
        new_score = {}
        for k, v in raw_score.items():
            new_score[k] = v["value"]
        scores.append(new_score)

    df = pd.DataFrame(scores)
    final_scores = df.mean().to_dict()
    return final_scores


def make_plot(model_name, scores, environment_name, title):
    obs = scores.get("observed", {})
    unobs = scores.get("unobserved", {})

    metric_names = sorted(set(obs.keys()) | set(unobs.keys()))
    if not metric_names:
        print(f"No {title} data to plot for {model_name}")
        return
    metric_names = ["is_blackmailing"]

    observed = [obs.get(metric, 0.0) for metric in metric_names]
    unobserved = [unobs.get(metric, 0.0) for metric in metric_names]
    gaps = [unobserved[i] - observed[i] for i in range(len(metric_names))]

    x = np.arange(len(metric_names))
    width = 0.35

    fig, ax = plt.subplots(figsize=(10, 6))
    rects1 = ax.bar(x - width/2, observed, width, label='Observed')
    rects2 = ax.bar(x + width/2, unobserved, width, label='Unobserved')

    # Add gap annotations
    for i, gap in enumerate(gaps):
        ax.annotate(f'{gap:.2f}',
                    xy=(x[i], max(observed[i], unobserved[i]) + 0.03),
                    xytext=(0, 3),
                    textcoords="offset points",
                    ha='center', va='bottom',
                    fontsize=9,
                    color='red' if gap > 0 else 'green')

    ax.set_ylabel(title)
    ax.set_title(f'{title} for {model_name} in {environment_name}')
    ax.set_xticks(x)
    ax.set_xticklabels(metric_names, rotation=45, ha='right')
    ax.set_ylim(0, 1.1)
    ax.legend()

    plt.tight_layout()
    os.makedirs("plots", exist_ok=True)
    plt.savefig(f"plots/{model_name}_{title}_{environment_name}.png")
    plt.show()


def make_combined_plot(model_scores: dict, environment_name: str, title: str):
    model_names = list(model_scores.keys())
    observed = []
    unobserved = []

    for model_name in model_names:
        scores = model_scores[model_name]
        obs_score = scores.get("observed", {}).get("is_blackmailing", 0.0)
        unobs_score = scores.get("unobserved", {}).get("is_blackmailing", 0.0)
        observed.append(obs_score)
        unobserved.append(unobs_score)

    x = np.arange(len(model_names))
    width = 0.35

    fig, ax = plt.subplots(figsize=(10, 6))
    ax.bar(x - width/2, observed, width, label='Observed')
    ax.bar(x + width/2, unobserved, width, label='Unobserved')

    # Annotate gaps
    for i in range(len(model_names)):
        gap = unobserved[i] - observed[i]
        ax.annotate(f'{gap:.2f}',
                    xy=(x[i], max(observed[i], unobserved[i]) + 0.03),
                    xytext=(0, 3),
                    textcoords="offset points",
                    ha='center', va='bottom',
                    fontsize=9,
                    color='red' if gap > 0 else 'green')

    ax.set_ylabel("is_blackmailing")
    ax.set_title(f"{title} in {environment_name}")
    ax.set_xticks(x)
    ax.set_xticklabels(model_names, rotation=0)
    ax.set_ylim(0, 1.1)
    ax.legend()

    plt.tight_layout()
    os.makedirs("plots", exist_ok=True)
    plt.savefig(f"plots/combined_{title}_{environment_name}.png")
    plt.show()


if __name__ == "__main__":
    model_paths = [
        ("claude", "results/validate_model/blackmail_claude/2025-07-29_20-39-41"),
        ("llama", "results/validate_model/blackmail_llama/2025-07-28_21-18-07"),
        ("qwen",  "results/validate_model/blackmail_qwen/2025-07-30_08-50-41"),
        # ("qwen", "results/validate_model/blackmail_qwen/2025-07-30_10-33-13"),  # alternate run
    ]

    for model_name, path in model_paths:
        data = process_run(path)
        make_plot(model_name, data, "blackmail", f"{model_name} blackmail")

    model_scores = {}
    for model_name, path in model_paths:
        model_scores[model_name] = process_run(path)

    make_combined_plot(
        model_scores, environment_name="blackmail", title="is_blackmailing")
