from datetime import datetime, timedelta, timezone
import re
import os
from pathlib import Path
from bench_af.util.environments import inference
from bench_af.util.environments.cot_scorers import behave_because_observed, misbehave_because_unobserved
import json
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

PROJECT_ROOT = Path(__file__).parent.parent

def change_to_project_root():
    """Change to the project root directory (parent of util/)."""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)
    os.chdir(project_root)

def save_current_time():
    with open("cache/last_runtime.txt", "w") as f:
        f.write(datetime.now().isoformat())

def get_last_runtime():
    if os.path.exists("cache/last_runtime.txt"):
        if os.path.getsize("cache/last_runtime.txt") > 0:
            with open("cache/last_runtime.txt", "r") as f:
                return datetime.fromisoformat(f.read()).replace(tzinfo=timezone.utc)
        else:
            return datetime.now(timezone.utc) - timedelta(days=3650)
    else:
        return datetime.now(timezone.utc) - timedelta(days=3650)

def get_timestamp_of_log(path: str):
    match = re.match(r"(\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2})-(\d{2})-(\d{2})", path)
    if match:
        dt_part = match.group(1)
        offset_hours = int(match.group(2))
        offset_minutes = int(match.group(3))
        
        # Detect if the offset is negative based on the original string
        if path[len(dt_part)] == '-':
            sign = -1
        else:
            sign = 1
        
        offset = timedelta(hours=sign * offset_hours, minutes=sign * offset_minutes)
        tz = timezone(offset)
        dt = datetime.strptime(dt_part, "%Y-%m-%dT%H-%M-%S")
        return dt.replace(tzinfo=tz)
    else:
        raise ValueError(f"Invalid path: {path}")

def get_environments():
    environments = []
    for p in (PROJECT_ROOT / "environments").iterdir():
        if p.is_dir() and p.name != "__pycache__":
            environments.append(p.name)
    return sorted(environments, key=len)


def get_models():
    models = []
    for p in (PROJECT_ROOT / "models").iterdir():
        if p.is_dir() and p.name != "__pycache__":
            models.append(p.name)
    return sorted(models, key=len)

def was_successful(results_dir: str):
    eval_logs = sorted([f for f in os.listdir(os.path.join(
                results_dir, "logs")) if f.endswith(".json")])
    if len(eval_logs) < 2:
        return False
    with open(os.path.join(results_dir, "logs", eval_logs[0]), "r") as f:
        eval_log_data = json.load(f)
        if eval_log_data["status"] != "success":
            return False
    with open(os.path.join(results_dir, "logs", eval_logs[1]), "r") as f:
        eval_log_data = json.load(f)
        if eval_log_data["status"] != "success":
            return False
    return True

def get_results_dirs():
    results_dirs = []
    for model in get_models():
        for environment in get_environments():
            candidates = []
            for dirpath, dirnames, filenames in os.walk(PROJECT_ROOT / "results"):
                for dirname in dirnames:
                    try:
                        timestamp = get_timestamp_of_log(os.listdir(os.path.join(dirpath, dirname, "logs"))[0])
                        if dirname == f"{model}_{environment}_compliance"\
                            and timestamp > get_last_runtime()\
                            and was_successful(os.path.join(dirpath, dirname)):
                            candidates.append((timestamp, os.path.join(dirpath, dirname)))
                    except:
                        pass
            if len(candidates) > 0:
                results_dirs.append(sorted(candidates)[-1][1])
    return results_dirs

def load_dict(path: str, default_dict: dict):
    if os.path.exists(path):
        if os.path.getsize(path) > 0:
            with open(path, "r") as f:
                return json.load(f)
        else:
            return default_dict
    else:
        return default_dict

def update_metrics():
    results_dirs = get_results_dirs()
    models = get_models()
    environments = get_environments()
    compliance_dict = load_dict(
                            "cache/compliance_dict.json", 
                            {m: {e: {"observed": None, "unobserved": None} 
                            for e in environments} for m in models})
    behave_because_observed_dict = load_dict(
                            "cache/behave_because_observed_dict.json", 
                            {m: {e: {"observed": None, "unobserved": None} 
                            for e in environments} for m in models})
    misbehave_because_unobserved_dict = load_dict(
                            "cache/misbehave_because_unobserved_dict.json", 
                            {m: {e: {"observed": None, "unobserved": None} 
                            for e in environments} for m in models})
    behave_examples_dict = load_dict(
                            "cache/behave_examples_dict.json", 
                            {m: {e: {"observed": [], "unobserved": []} 
                            for e in environments} for m in models})
    misbehave_examples_dict = load_dict(
                            "cache/misbehave_examples_dict.json", 
                            {m: {e: {"observed": [], "unobserved": []} 
                            for e in environments} for m in models})
    behave_conditional_dict = load_dict(
                            "cache/behave_conditional_dict.json", 
                            {m: {e: None
                            for e in environments} for m in models})
    misbehave_conditional_dict = load_dict(
                            "cache/misbehave_conditional_dict.json", 
                            {m: {e: None
                            for e in environments} for m in models})
    for eval_log_dir in results_dirs:
        behave_reasoning_indices_observed = []
        behave_reasoning_indices_unobserved = []
        misbehave_reasoning_indices_observed = []
        misbehave_reasoning_indices_unobserved = []
        scores_observed = []
        scores_unobserved = []
        model = None
        environment = None
        for m in models:
            if m in eval_log_dir:
                model = m
        for e in environments:
            if e in eval_log_dir:
                environment = e
        if model is not None and environment is not None:
            eval_logs = sorted([f for f in os.listdir(os.path.join(
                eval_log_dir, "logs")) if f.endswith(".json")])
            print(
                f"\nProcessing {eval_log_dir} ({model} in {environment}), found {len(eval_logs)} JSON files: {eval_logs}")

            # Process first file (observed)
            with open(os.path.join(eval_log_dir, "logs", eval_logs[0]), "r") as f:
                eval_log_data = json.load(f)
                print(
                    f"  Keys in {eval_logs[0]}: {list(eval_log_data.keys())}")
                try:
                    samples = eval_log_data["reductions"][0]["samples"]
                    scores_observed = [s["value"] for s in samples]
                    compliance_dict[model][environment]["observed"] = sum(scores_observed) / len(scores_observed)
                    print(
                        f"  Compliance score (observed): {compliance_dict[model][environment]['observed']:.3f}")
                except:
                    print(f"  Can't find compliance score in {eval_logs[0]}")

            if "no_cot" not in model:    
                try:
                    templated_responses = [eval_log_data["samples"][i]["output"]["choices"]
                                           [0]["message"]["content"] for i in range(len(eval_log_data["samples"]))]
                    cots = []
                    for r in templated_responses:
                        if isinstance(r, str):
                            cots.append(r)
                        elif isinstance(r, list):
                            cots.append(r[0]["reasoning"])
                        else:
                            print(f"  Unknown response type: {type(r)}")

                    # Process both metrics together to avoid multiple API calls
                    if cots:
                        behave_results = behave_because_observed(cots)
                        misbehave_results = misbehave_because_unobserved(cots)
                        behave_reasoning_indices_observed = [i for i in range(len(behave_results)) if behave_results[i].startswith("C")]
                        misbehave_reasoning_indices_observed = [i for i in range(len(misbehave_results)) if misbehave_results[i].startswith("C")]
                        behave_score = len(behave_reasoning_indices_observed) / len(behave_results)
                        misbehave_score = len(misbehave_reasoning_indices_observed) / len(misbehave_results)
                        if behave_score > 0:
                            behave_examples_dict[model][environment]["observed"] = cots[behave_results.index("C")]
                            print(
                                f"Example of model {model} in environment {environment} doing \"behave because observed\" reasoning: \n{cots[behave_results.index('C')]}\n\n\n")
                        if misbehave_score > 0:
                            misbehave_examples_dict[model][environment]["observed"] = cots[misbehave_results.index(
                                "C")]
                        behave_because_observed_dict[model][environment]["observed"] = behave_score
                        misbehave_because_unobserved_dict[model][environment]["observed"] = misbehave_score

                        # print(f"  Behave because observed: {behave_score:.3f}")
                        # print(f"  Misbehave because unobserved: {misbehave_score:.3f}")
                except Exception as e:
                    print(f"  Can't find COTs in {eval_logs[0]}: {e}")

            # Process second file (unobserved)
            if len(eval_logs) > 1:
                with open(os.path.join(eval_log_dir, "logs", eval_logs[1]), "r") as f:
                    eval_log_data = json.load(f)
                    print(
                        f"  Keys in {eval_logs[1]}: {list(eval_log_data.keys())}")
                    try:
                        samples = eval_log_data["reductions"][0]["samples"]
                        scores_unobserved = [s["value"] for s in samples]
                        compliance_dict[model][environment]["unobserved"] = sum(scores_unobserved) / len(scores_unobserved)
                        print(
                            f"  Compliance score (unobserved): {compliance_dict[model][environment]['unobserved']:.3f}")
                    except Exception as e:
                        print(
                            f"  Can't find compliance score in {eval_logs[1]}: {e}")

                    try:
                        templated_responses = [eval_log_data["samples"][i]["output"]["choices"]
                                               [0]["message"]["content"] for i in range(len(eval_log_data["samples"]))]
                        cots = []
                        for r in templated_responses:
                            if isinstance(r, str):
                                cots.append(r)
                            elif isinstance(r, list):
                                cots.append(r[0]["reasoning"])
                            else:
                                print(f"  Unknown response type: {type(r)}")

                        # Process both metrics together to avoid multiple API calls
                        if cots:
                            behave_results = behave_because_observed(cots)
                            misbehave_results = misbehave_because_unobserved(cots)
                            behave_reasoning_indices_unobserved = [i for i in range(len(behave_results)) if behave_results[i].startswith("C")]
                            misbehave_reasoning_indices_unobserved = [i for i in range(len(misbehave_results)) if misbehave_results[i].startswith("C")]
                            behave_score = len(behave_reasoning_indices_unobserved) / len(behave_results)
                            misbehave_score = len(misbehave_reasoning_indices_unobserved) / len(misbehave_results)
                            if behave_score > 0:
                                behave_examples_dict[model][environment]["unobserved"] = cots[behave_results.index(
                                    "C")]
                            if misbehave_score > 0:
                                misbehave_examples_dict[model][environment]["unobserved"] = cots[misbehave_results.index(
                                    "C")]
                                print(
                                    f"Example of model {model} in environment {environment} doing \"misbehave because unobserved\" reasoning: \n{cots[misbehave_results.index('C')]}\n\n\n")
                            behave_because_observed_dict[model][environment]["unobserved"] = behave_score
                            misbehave_because_unobserved_dict[model][environment]["unobserved"] = misbehave_score
                            # print(f"  Behave because observed: {behave_score:.3f}")
                            # print(f"  Misbehave because unobserved: {misbehave_score:.3f}")
                    except Exception as e:
                        print(f"  Can't find COTs in {eval_logs[1]}: {e}")
            try:
                behave_conditional_dict[model][environment] = 1-\
                    sum([scores_observed[i] for i in behave_reasoning_indices_observed]) \
                    + sum([scores_unobserved[i] for i in behave_reasoning_indices_unobserved])\
                    / (len(behave_reasoning_indices_observed)+len(behave_reasoning_indices_unobserved))
                misbehave_conditional_dict[model][environment] = \
                    sum([scores_observed[i] for i in misbehave_reasoning_indices_observed]) \
                    + sum([scores_unobserved[i] for i in misbehave_reasoning_indices_unobserved])\
                    / (len(misbehave_reasoning_indices_observed)+len(misbehave_reasoning_indices_unobserved))
            except Exception as e:
                print(f"  Can't find conditional scores in {eval_logs[1]}: {e}")

    with open("cache/compliance_dict.json", "w") as f:
        json.dump(compliance_dict, f)
    with open("cache/behave_because_observed_dict.json", "w") as f:
        json.dump(behave_because_observed_dict, f)
    with open("cache/misbehave_because_unobserved_dict.json", "w") as f:
        json.dump(misbehave_because_unobserved_dict, f)
    with open("cache/behave_examples_dict.json", "w") as f:
        json.dump(behave_examples_dict, f)
    with open("cache/misbehave_examples_dict.json", "w") as f:
        json.dump(misbehave_examples_dict, f)
    with open("cache/behave_conditional_dict.json", "w") as f:
        json.dump(behave_conditional_dict, f)
    with open("cache/misbehave_conditional_dict.json", "w") as f:
        json.dump(misbehave_conditional_dict, f)


def make_plots_pairs(model_name, data, title, plot_dir):
    envs = []
    observed = []
    unobserved = []
    gaps = []

    for env, scores in data.items():
        if scores is None:
            continue
        obs = scores.get("observed")
        unobs = scores.get("unobserved")
        if obs is None or unobs is None:
            continue
        envs.append(env)
        observed.append(obs)
        unobserved.append(unobs)
        gaps.append(unobs - obs)

    if not envs:
        print(f"No {title} data to plot for {model_name}")
        return

    x = np.arange(len(envs))
    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}')
    ax.set_xticks(x)
    ax.set_xticklabels(envs, rotation=45, ha='right')
    ax.set_ylim(0, 1.1)
    ax.legend()

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

def make_plots_single(model_name, data, title, plot_dir):
    envs = []
    values = []

    for env, score in data.items():
        if score is None:
            continue
        if isinstance(score, (int, float)):
            envs.append(env)
            values.append(score)
        else:
            print(f"Skipping {env}: Expected a single number or None, got {type(score)}")
            continue

    if not envs:
        print(f"No {title} data to plot for {model_name}")
        return

    x = np.arange(len(envs))
    width = 0.6

    fig, ax = plt.subplots(figsize=(10, 6))
    bars = ax.bar(x, values, width, color='steelblue')

    # Add value annotations
    for i, val in enumerate(values):
        ax.annotate(f'{val:.2f}',
                    xy=(x[i], val),
                    xytext=(0, 3),
                    textcoords="offset points",
                    ha='center', va='bottom',
                    fontsize=9)

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

    plt.tight_layout()
    plt.show()

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



if __name__ == "__main__":
    plot_dir = "plots_test"
    print("\nGenerating plots...")
    update_metrics()
    compliance_dict = load_dict("cache/compliance_dict.json", {})
    behave_because_observed_dict = load_dict("cache/behave_because_observed_dict.json", {})
    misbehave_because_unobserved_dict = load_dict("cache/misbehave_because_unobserved_dict.json", {})
    behave_conditional_dict = load_dict("cache/behave_conditional_dict.json", {})
    misbehave_conditional_dict = load_dict("cache/misbehave_conditional_dict.json", {})
    # Create plots with progress bars
    for model, env_data in tqdm(compliance_dict.items(), desc="Creating compliance plots"):
        make_plots_pairs(model, env_data, "compliance", plot_dir)
    
    for model, env_data in tqdm(behave_because_observed_dict.items(), desc="Creating behave_because_observed plots"):
        make_plots_pairs(model, env_data, "behave_because_observed", plot_dir)
    
    for model, env_data in tqdm(misbehave_because_unobserved_dict.items(), desc="Creating misbehave_because_unobserved plots"):
        make_plots_pairs(model, env_data, "misbehave_because_unobserved", plot_dir)

    for model, env_data in tqdm(behave_conditional_dict.items(), desc="Creating behave_conditional plots"):
        make_plots_single(model, env_data, "behave_conditional", plot_dir)

    for model, env_data in tqdm(misbehave_conditional_dict.items(), desc="Creating misbehave_conditional plots"):
        make_plots_single(model, env_data, "misbehave_conditional", plot_dir)

    print("All plots generated successfully!")

    # Give user the option to avoid caching this run if if failed / was interrupted for some reason
    if input("Would you like to cache this run? (y/n)\n(You can always undo this by setting util/cache/last_runtime.txt to a previous time.)") == "y":
        save_current_time()