import concurrent.futures
import os
import re
import sys
import warnings
from collections import Counter
from pathlib import Path

import numpy as np
import openpyxl
import pandas as pd
from scipy.stats import binom, chi2

from utils_agrb.constants import AGRB_TO_VWA

# Suppress RuntimeWarning for division by zero and PerformanceWarning for DataFrame fragmentation
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning)
from offline_experiments.eval_mapping import EVAL_CRITERIA, TEMPLATE_MAPPINGS, map_eval_to_score_with_context
from utils import timing_utils as timer

DOMAINS = [
    "shopping",
    "classifieds",
    "reddit",
    "chrome",
    "gimp",
    "libreoffice_calc",
    "libreoffice_impress",
    "libreoffice_writer",
    "multi_apps",
    "os",
    "thunderbird",
    "vlc",
    "vs_code",
]
OVEWRITE = True
DELETE_INVALID_FILES = False
IGNORE_HTML = False
file_to_parse_pattern = "**/conversation/**/*.txt"
# file_to_parse_pattern_html = "**/conversation/**/*.html"
file_to_parse_pattern_html = "**/*.html"

# experiments_path = "./offline_experiments/"
experiments_paths = ["./offline_experiments/"]  # "/mnt/d/webarena/experiments_offline"]
output_dir = "./offline_experiments"
envs = ["agrb_vwa"]
gold_scores_dirs = {
    "vwa": "offline_experiments/vwa/gold_scores",
    "osw": "offline_experiments/osw/gold_scores",
    "agrb_vwa": "offline_experiments/agrb_vwa/gold_scores",
}

global agrb_mapping
if "agrb_vwa" in envs:
    agrb_mapping = pd.read_csv(AGRB_TO_VWA)

# eval_criteria is now imported from eval_mapping module as EVAL_CRITERIA
eval_criteria = EVAL_CRITERIA
include_only_strs = [
    # "*gemini-2.5-flash-zero-random/*desc_compare_v2*",
    # "*gemini-2.5-flash-zero-random/*desc_compare_v4*",
    # "*gemini-2.5-flash-zero-random/*desc_compare_v5*",
    # "1p_no_k-tri-basic_cot-actions-rule_efficiency_vision",
    # "gemini-2.5-flash-zero-random",
    # "gemini-2.5-flash-zero-random",
    # "offline_experiments/vwa/gemini-2.5-noverif-2025-06-19/gemini-2.5-flash-zero-random/2p-tri-desc_compare_v5-actions-uncond-expertv4-gemini-2.5-flash-zero-random-k_rule_efficiency_v4/reddit/conversation/114.txt",
    # "offline_experiments/vwa/gemini-2.5-noverif-2025-06-19/gemini-2.5-flash-zero-random/2p-four_num_unc-desc_ss-actions-uncond-expertv4-gemini-2.5-flash-zero-random-k_rule_efficiency_vision_v2/shopping",
    # "2p-four_num_unc-basic_cot-actions-uncond-expertv4-gemini-2.5-flash-zero-random-k_rule_efficiency_vision_v2"
    # "2p-tri-desc_compare_v2-actions-uncond-expertv4-gemini-2.5-flash-zero-random-k_rule_efficiency_vision_v2"
    # "four_room_2"
    # "four_room_score"
    # "four_num_uncond",
]
exclude_strs = [
    "zzOld",
    "zzOther",
    "logs",
    # "claude",
    # "k-cond",
    # "k-uncond",
    "verify_example",
    # "-base",
    "gold_scores",
    # "four",
    # "quad",
    "web_search",
    # "gpt4o",
    # "gemini-2.0-flash-001",
    # "_num_",
    "/k-uncond",
    "/k-cond",
    "_b.html",
]
results_csv_name = "evaluations.csv"

EVAL_SECTION_REGEX = re.compile(
    # r"(?:EVALUATION:|Status:|ANSWER:|# EVALUATION)\s*(.*?)\s*(?:\n|$|FEEDBACK:|</pre>)",
    r"(?:EVALUATION:|Status:|ANSWER:|# EVALUATION)\s*(.*?)(?=\n(?:EVALUATION:|Status:|ANSWER:|# EVALUATION|FEEDBACK:|</pre>|$))",
    # r"(?:EVALUATION:|Status:|ANSWER:|# EVALUATION)\s*(.*?)(?=\n\n|\nFEEDBACK:|\n</pre>|$)",
    re.S | re.IGNORECASE,
)

EVAL_CRITERIA_REGEX = re.compile(
    rf"^\s*:?\s*({'|'.join(map(re.escape, sorted(eval_criteria, key=len, reverse=True)))})(?::)?",
    re.IGNORECASE,
)
GENERATION_BLOCK_REGEX = re.compile(
    r"-+\s*GENERATION\s*\d+\s*-+[\s\S]*?(?=-+\s*GENERATION\s*\d+\s*-+|\Z)",
    re.IGNORECASE,
)


def parse_eval(content: str) -> str:
    """
    Extracts the evaluation score (e.g., SUCCESS, FAILURE) or a numeric value from the EVALUATION section.
    Returns (categorical, numeric) where one may be None/empty.
    """
    # Try categorical first
    match = EVAL_CRITERIA_REGEX.search(content)
    if match:
        return match.group(1).upper()
    # Try numeric (int or float)
    num_match = re.match(r"^\s*:?-?\s*([0-9]+(?:\.[0-9]+)?)", content)
    if num_match:
        return str(num_match.group(1))
    return ""


def mcnemar(b: int, c: int, apply_correction: bool = False) -> float:
    b = int(b)
    c = int(c)
    check_valid = lambda n: isinstance(n, int) or (isinstance(n, float) and n.is_integer())
    if not all(map(check_valid, [b, c])):
        print(f"b: {b}, c: {c}")
        print(check_valid(b), check_valid(c))
        raise ValueError("b and c must be integers!")
    n_min, n_max = sorted([b, c])

    if (n_min + n_max) >= 25:
        chi2_statistic = (abs(n_min - n_max) - int(apply_correction)) ** 2 / (n_min + n_max)
        pvalue = chi2.sf(chi2_statistic, 1)
    else:
        pvalue = 2 * binom.cdf(n_min, n_min + n_max, 0.5) - binom.pmf(n_min, n_min + n_max, 0.5)
    return pvalue


def delete_invalid_files(file: Path | str):
    # Split the extension from the file.
    if not isinstance(file, Path):
        file = Path(file)
    base_path = str(file.with_suffix(""))

    all_files = []
    all_files.append(f"{base_path}.html")
    all_files.append(f"{base_path}.txt")
    parent_dir = Path(file).parent.parent
    all_files.append(f"{parent_dir}/usage/{base_path}.csv")
    for file in all_files:
        if not os.path.exists(file):
            continue
        if os.path.isfile(file):
            os.remove(file)


def get_majority_eval(eval_results: list[str]) -> str:
    if len(eval_results) == 0:
        return ""
    if len(eval_results) == 1:
        return eval_results[-1]

    counter = Counter(eval_results)
    most_common_eval, n = counter.most_common(1)[0]
    return most_common_eval


def _extract_evaluations_from_content(content: str) -> list[str]:
    """Return list containing only the last EVALUATION section from each GENERATION block."""
    eval_sections: list[str] = []
    for block in GENERATION_BLOCK_REGEX.findall(content):
        all_evals = EVAL_SECTION_REGEX.findall(block)
        if all_evals:
            # keep only the last evaluation in the block (ignores placeholders)
            raw_eval = all_evals[-1]
            # Clean up the evaluation content by removing common prefixes
            cleaned_eval = re.sub(
                r"^\s*(?:Score|Value|Rating|Result)\s*:?\s*", "", raw_eval.strip(), flags=re.IGNORECASE
            )
            eval_sections.append(cleaned_eval)
    return eval_sections


def get_eval_from_trace_file(file: Path | str) -> tuple[str, str, list[str] | None, str, Path]:
    if not isinstance(file, Path):
        file = Path(file)

    task_id = get_task_id_from_trace_file(file)
    variation_name = get_variation_name_from_trace_file(file)
    domain = get_domain_from_path(file)
    try:
        with open(file, "r", encoding="utf-8") as f:
            content = f.read().strip()

            # Extract evaluation sections only from inside GENERATION blocks.
            matches = _extract_evaluations_from_content(content)
            if not matches:
                print(f"No evaluation found in {file}", flush=True)
                if DELETE_INVALID_FILES:
                    delete_invalid_files(file)
                return task_id, variation_name, None, domain, file

            # Parse all eval results from all EVALUATION sections
            eval_results = []
            for eval_section_str in matches:
                eval_result = parse_eval(eval_section_str)
                if eval_result:
                    eval_results.append(eval_result)

            if not eval_results:
                print(f"Unable to parse textual score from eval sections for {file}", flush=True)
                print(f"eval_section_str: {eval_section_str}", flush=True)
                if DELETE_INVALID_FILES:
                    delete_invalid_files(file)
                return task_id, variation_name, None, domain, file

            # Take the most frequent answer as the final answer
            # majority_vote = get_majority_eval(eval_results)

            return task_id, variation_name, eval_results, domain, file
    except Exception as e:
        print(f"Error reading file {file}: {e}", flush=True)
        return task_id, variation_name, None, domain, file


def get_task_id_from_trace_file(file: Path | str) -> str:
    if not isinstance(file, Path):
        file = Path(file)

    # 29_k-0.html -> 29 ; 29_k-0.txt -> 29
    filename = file.stem

    task_id = re.sub(r"_k-.*$", "", filename)
    task_id = re.sub(r"conversation_", "", task_id, flags=re.IGNORECASE)
    return task_id


def get_variation_name_from_trace_file(file: Path | str) -> str:
    if not isinstance(file, Path):
        file = Path(file)

    return file.parent.parent.parent.name


def save_csv(
    df: pd.DataFrame, path: Path | str, gold_scores: pd.DataFrame | None = None, key_gold_score: str = "gold_score"
):
    df_to_save = df.copy()
    if gold_scores is not None:
        # Ensure unique_id is of type string for both DataFrames
        gold_df = gold_scores.copy()
        # Create unique_id concatenating domain and task_id
        gold_df["domain_task_id"] = gold_df["domain"].astype(str) + "_" + gold_df["task_id"].astype(str)

        gold_df["domain_task_id"] = gold_df["domain_task_id"].astype(str)
        df_to_save["domain_task_id"] = df_to_save["domain_task_id"].astype(str)

        # Subset gold_scores to only include unique_id and score, and rename score to gold_score
        gold_scores_subset = gold_df[["domain_task_id", key_gold_score, "gold_source"]]

        # Drop the existing gold_score column if it already exists
        if key_gold_score in df_to_save.columns:
            df_to_save = df_to_save.drop(columns=[key_gold_score])

        df_to_save = pd.merge(df_to_save, gold_scores_subset, on="domain_task_id", how="left")

    # Sort by unique_id
    df_to_save = df_to_save.sort_values(by="domain_task_id")
    # Reorder to have metadata cols first
    cols_first = [col for col in df_to_save.columns if not col.endswith("_eval")]
    eval_cols = [col for col in df_to_save.columns if col.endswith("_eval")]
    df_to_save = df_to_save[cols_first + eval_cols]

    # Add per-row confusion label (TP/FP/TN/FN) for each evaluation column
    if gold_scores is not None and "gold_score" in df_to_save.columns:
        df_confusion = pd.DataFrame()
        # Add non-eval columns to df_confusion
        for col in df_to_save.columns:
            if col not in eval_cols:
                df_confusion[col] = df_to_save[col]

        for col in eval_cols:
            eval_upper = df_to_save[col].astype(str).str.upper().str.strip()
            gold = df_to_save["gold_score"]
            confusion = np.where(
                (gold == 1) & (eval_upper == "SUCCESS"),
                "TP",
                np.where(
                    (gold == 0) & (eval_upper == "SUCCESS"),
                    "FP",
                    np.where(
                        (gold == 0) & (eval_upper.isin(["FAILURE", "PARTIAL FAILURE", "PARTIAL SUCCESS"])),
                        "TN",
                        np.where(
                            (gold == 1) & (eval_upper.isin(["FAILURE", "PARTIAL FAILURE", "PARTIAL SUCCESS"])), "FN", ""
                        ),
                    ),
                ),
            )
            df_to_save[f"{col}_confusion"] = confusion
            df_confusion[col] = confusion
        df_confusion.to_csv(Path(path).parent / "confusion.csv", index=False)

    df_to_save.to_csv(path, index=False)

    return df_to_save


def remove_files_to_parse_if_exists(files_to_parse: list[Path], existing_results_df: pd.DataFrame) -> list[Path]:
    variation_names = [get_variation_name_from_trace_file(file) for file in files_to_parse]
    task_ids = [get_task_id_from_trace_file(file) for file in files_to_parse]
    domains = [get_domain_from_path(file) for file in files_to_parse]
    unique_ids = [f"{domain}_{task_id}" for domain, task_id in zip(domains, task_ids)]

    final_files_to_parse = []

    # Create lookup dictionaries for faster access
    existing_values = {
        (str(row["domain_task_id"]), col): row[col]
        for _, row in existing_results_df.iterrows()
        for col in existing_results_df.columns
        if col not in ["domain_task_id", "source", "gold_score", "env", "domain", "task_id"]
    }

    # Create a list of files that need to be parsed
    for file, variation_name, uid in zip(files_to_parse, variation_names, unique_ids):
        if (uid, variation_name) in existing_values:
            val = existing_values[(uid, variation_name)]
            if val and pd.notna(val):
                continue
        final_files_to_parse.append(file)

    return final_files_to_parse


def agrb_vwa_get_domain_from_path(path: Path | str) -> str:
    global agrb_mapping
    task_id = Path(path).stem
    task_id = str(task_id)
    line = agrb_mapping[agrb_mapping["task_id"] == task_id]
    return line["domain"].values[0]


def get_domain_from_path(path: Path | str) -> str:
    path = str(path)
    if "agrb" in str(path):
        return agrb_vwa_get_domain_from_path(path)
    for domain in DOMAINS:
        match = re.match(rf".*/({domain})/.*", path)
        if match:
            return match.group(1)
    raise ValueError(f"Unable to determine domain from {path}")


def _results_to_df(results: dict) -> pd.DataFrame:
    # Convert the nested dictionary to a vertical DataFrame
    rows = []
    for domain_task_id, variations in results.items():
        for variation_name, data in variations.items():
            row = {
                "domain_task_id": domain_task_id,
                "domain": data["domain"],
                "task_id": data["task_id"],
                "config_name": variation_name,
                "eval": data["eval_result"],
            }
            rows.append(row)

    result_df = pd.DataFrame(rows)
    return result_df


def parse_experiment_results(
    base_path,
    existing_results_df: pd.DataFrame | None = None,
):
    num_invalid_files = 0
    print(f"Parsing evaluations in {base_path}", flush=True)

    files_to_parse = list(Path(base_path).glob(file_to_parse_pattern))

    if not IGNORE_HTML:
        files_to_parse_html = list(Path(base_path).glob(file_to_parse_pattern_html))
    else:
        files_to_parse_html = []

    files_to_parse = [file for file in files_to_parse if not any(exclude in str(file) for exclude in exclude_strs)]
    files_to_parse_html = [
        file for file in files_to_parse_html if not any(exclude in str(file) for exclude in exclude_strs)
    ]
    if include_only_strs:
        # Convert wildcard patterns to regex patterns
        regex_patterns = [re.compile(pattern.replace("*", ".*")) for pattern in include_only_strs]
        files_to_parse = [file for file in files_to_parse if any(regex.search(str(file)) for regex in regex_patterns)]
        files_to_parse_html = [
            file for file in files_to_parse_html if any(regex.search(str(file)) for regex in regex_patterns)
        ]

    files_to_parse_search = set(re.sub(r"\.txt$", "", str(file)) for file in files_to_parse)
    # If file is in the files_to_parse already, ignore; else, add it.
    for file in files_to_parse_html:
        file_base_path = re.sub(r"\.html$", "", str(file))
        if not file_base_path:
            continue
        if file_base_path in files_to_parse_search:
            continue
        else:
            files_to_parse.append(file)

    if existing_results_df is not None:
        files_to_parse = remove_files_to_parse_if_exists(files_to_parse, existing_results_df)

    if len(files_to_parse) == 0:
        print(f"No files to parse in {base_path}", flush=True)
        return None

    else:
        print(f"Parsing {len(files_to_parse)} files in {base_path}", flush=True)

    results = {}

    with concurrent.futures.ProcessPoolExecutor(max_workers=16) as executor:
        # Submit all file processing tasks concurrently.
        futures = {executor.submit(get_eval_from_trace_file, file): file for file in files_to_parse}
        for future in concurrent.futures.as_completed(futures):
            task_id, variation_name, eval_result, domain, file = future.result()
            if not eval_result:
                num_invalid_files += 1
                print(f"Not able to parse eval result for {file}")
                continue

            domain_task_id = f"{domain}_{task_id}"
            results.setdefault(domain_task_id, {})

            if variation_name not in results[domain_task_id]:
                results[domain_task_id][variation_name] = {
                    "eval_results": [],
                    "domain": domain,
                    "task_id": task_id,
                }
            if eval_result:
                results[domain_task_id][variation_name]["eval_results"].extend(eval_result)

    if not results:
        return existing_results_df

    # If multiple generations, handle majority vote and individual results
    final_results = {}
    for domain_task_id, data in results.items():
        final_results.setdefault(domain_task_id, {})
        for variation_name, variation_data in data.items():
            final_results[domain_task_id].setdefault(variation_name, {})

            final_results[domain_task_id][variation_name] = {
                "eval_result": get_majority_eval(variation_data["eval_results"]),
                "domain": variation_data["domain"],
                "task_id": variation_data["task_id"],
            }

            if len(variation_data["eval_results"]) > 1:
                # For each sample, also compute the evaluation result.
                n = len(variation_data["eval_results"])
                for i in range(n):
                    variation_name_i = f"{variation_name}-{i}"
                    cat = variation_data["eval_results"][i] if i < len(variation_data["eval_results"]) else ""
                    final_results[domain_task_id][variation_name_i] = {
                        "eval_result": cat,
                        "domain": variation_data["domain"],
                        "task_id": variation_data["task_id"],
                    }
            del variation_data["eval_results"]

    # # Merge additional data
    # for domain_task_id, additional_variations in additional_data.items():
    #     if domain_task_id not in results:
    #         results[domain_task_id] = {}
    #     results[domain_task_id].update(additional_variations)

    # Convert dictionary to DataFrame after processing all subdirs
    result_df = _results_to_df(final_results)

    if existing_results_df is not None:
        result_df = (
            result_df.set_index("domain_task_id")
            .combine_first(existing_results_df.set_index("domain_task_id"))
            .reset_index()
        )

    print(f"Number of invalid files: {num_invalid_files} of {len(files_to_parse)} for {base_path}", flush=True)
    return result_df


def compute_distance_skewness(X, theta):
    X = np.array(X)
    # pairwise distances between the elements of X
    pairwise_distances = np.abs(np.subtract.outer(X, X))

    # numerator and denominator of the distance skewness formula
    numerator = np.sum(pairwise_distances)
    denominator = np.sum(np.abs(np.add.outer(X, X) - 2 * theta))

    # handle the case when Pr(X=theta) = 1
    if denominator == 0:
        return 0
    else:
        return 1 - numerator / denominator


def compute_skewness(X):
    X = np.array(X)
    try:
        return sum([ele**3 for ele in X]) / len(X) / (sum([ele**2 for ele in X]) / len(X)) ** (3 / 2)
    except:
        print(f"Error computing skewness for {X}")
        return np.nan


def compute_confusion_stats(gold_scores: pd.DataFrame, evals: pd.DataFrame, eval_template: str | None = None):
    # Convert task_id to string
    gold_scores.loc[:, "domain_task_id"] = gold_scores["domain_task_id"].astype(str)
    evals.loc[:, "domain_task_id"] = evals["domain_task_id"].astype(str)

    evals = evals.copy()

    # Use context-aware mapping if template is provided
    evals["predicted_score"] = evals["eval"].apply(lambda x: map_eval_to_score_with_context(x, eval_template))

    # Merge evals with the ground-truth scores on task_id
    merged = pd.merge(evals, gold_scores, on="domain_task_id", how="right", suffixes=("_eval", "_true"))

    # Compute confusion matrix values based on the comparison of predicted vs. true
    false_positive = ((merged["predicted_score"] == 1) & (merged["gold_score"] == 0)).sum()
    false_negative = ((merged["predicted_score"] == 0) & (merged["gold_score"] == 1)).sum()
    true_positive = ((merged["predicted_score"] == 1) & (merged["gold_score"] == 1)).sum()
    true_negative = ((merged["predicted_score"] == 0) & (merged["gold_score"] == 0)).sum()
    num_pos = true_positive + false_negative
    num_neg = true_negative + false_positive

    all_diff = merged["predicted_score"] - merged["gold_score"]
    all_diff = all_diff.dropna()
    signed_bias = all_diff.mean()
    unsigned_bias = abs(all_diff).mean()
    skewness = compute_skewness(all_diff)
    distance_skewness = compute_distance_skewness(all_diff, 0)

    tp_ratio = true_positive / (true_positive + false_negative)
    tn_ratio = true_negative / (false_positive + true_negative)
    fp_ratio = false_positive / (false_positive + true_negative)
    fn_ratio = false_negative / (true_positive + false_negative)
    accuracy = (true_positive + true_negative) / (true_positive + true_negative + false_positive + false_negative)
    precision = true_positive / (true_positive + false_positive)
    f1_score = 2 * (precision * tp_ratio) / (precision + tp_ratio)

    f1_pos = 2 * true_positive / (2 * true_positive + false_positive + false_negative)
    f1_neg = 2 * true_negative / (2 * true_negative + false_positive + false_negative)
    f1_balanced = (f1_pos * num_pos + f1_neg * num_neg) / (num_pos + num_neg)
    f1_balanced_inv_weights = (f1_pos * num_neg + f1_neg * num_pos) / (num_pos + num_neg)

    # --- Add Balanced Accuracy and MCC ---
    # Balanced Accuracy
    balanced_accuracy = (tp_ratio + tn_ratio) / 2

    # Matthews Correlation Coefficient (MCC)
    denominator = np.sqrt(
        (true_positive + false_positive)
        * (true_positive + false_negative)
        * (true_negative + false_positive)
        * (true_negative + false_negative)
    )
    if denominator == 0:
        mcc = np.nan
    else:
        mcc = ((true_positive * true_negative) - (false_positive * false_negative)) / denominator
    # --- end addition ---

    # Added counts for SUCCESS, PARTIAL SUCCESS, and FAILURE based on the "eval" column.
    criteria_counts = {}
    for criteria in eval_criteria:
        criteria_count = evals["eval"].str.upper().eq(criteria).sum()
        criteria_counts[criteria] = criteria_count

    # Filter the false negatives (i.e. cases where gold_score==1 but predicted 0)
    # Cast merged["eval"] to string and strip whitespace.
    merged["eval_upper"] = merged["eval"].astype(str).str.upper().str.strip()
    false_negatives = merged[(merged["predicted_score"] == 0) & (merged["gold_score"] == 1)]
    false_neg_partial_success = false_negatives[false_negatives["eval_upper"] == "PARTIAL SUCCESS"].shape[0]
    false_neg_failure = false_negatives[false_negatives["eval_upper"] == "FAILURE"].shape[0]

    total = len(gold_scores)
    na_count = total - false_positive - false_negative - true_positive - true_negative
    total_effectively_completed = total - na_count
    mcnemar_p_value = mcnemar(false_positive, false_negative, apply_correction=True)

    all_data = {  # FIXME: remove rounding
        "tp_ratio": round(tp_ratio * 100, 1),
        "tn_ratio": round(tn_ratio * 100, 1),
        "accuracy": round(accuracy * 100, 1),
        "precision": round(precision * 100, 1),
        "f1_score": round(f1_score * 100, 1),
        "distance_skewness": round(distance_skewness * 100, 2),
        "signed_bias": round(signed_bias * 100, 2),
        "total_effectively_completed": round(total_effectively_completed, 2),
        "NA %": round(100 * na_count / total, 2) if total > 0 else np.nan,
        "skewness": round(skewness * 100, 2),
        "unsigned_bias": round(unsigned_bias * 100, 2),
        "fn_ratio": fn_ratio,
        "fp_ratio": fp_ratio,
        "total": total,
        "f1_pos": f1_pos,
        "f1_neg": f1_neg,
        "f1_balanced": f1_balanced,
        "f1_balanced_inv_weights": f1_balanced_inv_weights,
        "false_positive": false_positive,
        "false_negative": false_negative,
        "true_positive": true_positive,
        "true_negative": true_negative,
        "balanced_accuracy": balanced_accuracy,
        "mcnemar_p_value": mcnemar_p_value,
        "mcc": mcc,
        "false_neg_partial_success": false_neg_partial_success,
        "false_neg_failure": false_neg_failure,
        "NA": na_count,
        "perc failure": (true_negative + false_positive) / total_effectively_completed
        if total_effectively_completed > 0
        else np.nan,
    }

    all_data.update(criteria_counts)

    return all_data


def extract_eval_template_from_variation_name(variation_name: str) -> str | None:
    """
    Extract evaluation template from variation name.

    The pattern is: [number]p-[template]-[other_parts...]
    For example: 2p-tri-desc_cot_v2-actions-...

    Args:
        variation_name: The variation name extracted from directory structure

    Returns:
        The evaluation template name or None if not found
    """
    if not variation_name:
        return None

    # Split by common delimiters
    parts = re.split(r"[-]", variation_name)
    if len(parts) < 2:
        return None

    # Check if first part matches pattern like "1p", "2p", etc.
    first_part = parts[0].lower()
    if re.match(r"\d+p$", first_part):
        # Template should be the second part
        potential_template = parts[1].lower()
    else:
        # Fallback: search through all parts
        parts = re.split(r"[-_]", variation_name)
        for part in parts:
            potential_template = part.lower()
            if potential_template in TEMPLATE_MAPPINGS:
                return potential_template
        return None

    if potential_template in TEMPLATE_MAPPINGS:
        return potential_template

    return None


# ===============================================================
# Main
# ===============================================================
# Assumes folder structure:
# experiments/
#  .../
#   gpt-4o-2024-08-06/
#     experiments_domain_1_date_1/
#       variation_1/
#       variation_2/
#       ...
#     experiments_domain_2_date_2/
#       variation_1/
#       variation_2/
#       ...


def get_gold_score_path_source(model_subdir: str | Path, gold_scores_dir: str | Path) -> tuple[Path, str]:
    model_subdir = str(model_subdir)
    for file in Path(gold_scores_dir).glob("*.csv"):
        if file.stem in model_subdir:
            source = file.stem
            return file, source
    raise ValueError(f"No gold scores found in {gold_scores_dir} for {model_subdir}")


def load_gold_scores(path_to_csv: str | Path) -> pd.DataFrame:
    gold_scores = pd.read_csv(path_to_csv)
    gold_scores.rename(columns={"score": "gold_score"}, inplace=True)
    gold_scores["task_id"] = gold_scores["task_id"].astype(str)

    # Ensure gold_scores has a domain; if not, assign a default.
    if "domain" not in gold_scores.columns:
        raise ValueError("Gold scores must have a domain column.")

    gold_scores = gold_scores.copy()
    gold_scores["domain_task_id"] = gold_scores["domain"].astype(str) + "_" + gold_scores["task_id"].astype(str)

    return gold_scores


def load_test_set_agrb(test_set_path: str) -> pd.DataFrame:
    test_set = pd.read_csv(test_set_path)
    # Get only "test" split
    test_set = test_set[test_set["split"] == "test"]
    test_set["task_id"] = test_set["task_id"].astype(str)
    return test_set


def load_eval_set(eval_set_path: str) -> set:
    """Load eval set file as a set of domain_task_id strings (e.g., 'shopping_0')."""
    eval_set = set()
    with open(eval_set_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            if len(parts) >= 2:
                domain, task_id = parts[0], parts[1]
                eval_set.add(f"{domain}_{task_id}")
    return eval_set


EVAL_SET_VWA_PATH = os.path.join(os.path.dirname(__file__), "eval_set_vwa.txt")
TEST_SET_AGRB_PATH = "utils_agrb/splits.csv"

all_evals = pd.DataFrame()
confusion_per_env = {}
raw_evals_per_env = {}
all_evals_per_env = {}

for env in envs:
    out_dir_env = Path(output_dir) / env

    confusion_matrices = pd.DataFrame()
    all_evals_env = pd.DataFrame()

    model_subdirs = []
    for experiments_path in experiments_paths:
        experiments_path_env = Path(experiments_path) / env
        agent_subdirs = [str(s) for s in Path(experiments_path_env).iterdir() if s.is_dir()]
        model_subdirs.extend(
            [subdir for agent_subdir in agent_subdirs for subdir in Path(agent_subdir).iterdir() if subdir.is_dir()]
        )

    model_subdirs = [str(d) for d in model_subdirs if not any(exclude in str(d) for exclude in exclude_strs)]

    for model_subdir in model_subdirs:
        if any(exclude in model_subdir for exclude in exclude_strs) or not any(env in model_subdir for env in envs):
            continue
        try:
            gold_scores_filepath, source = get_gold_score_path_source(Path(model_subdir).parent, gold_scores_dirs[env])
            gold_scores = load_gold_scores(gold_scores_filepath)
            gold_scores["gold_source"] = source
        except ValueError:
            print(f"No gold scores found in {gold_scores_dirs[env]} for {model_subdir}")
            continue

        if not OVEWRITE and (Path(model_subdir) / results_csv_name).exists():
            existing_results = pd.read_csv(Path(model_subdir) / results_csv_name)
            if existing_results.empty:
                existing_results = None
        else:
            existing_results = None

        # Get the model name from the subdirectory.
        model = Path(model_subdir).name + "__" + Path(model_subdir).parent.name

        model_results_all_domains = parse_experiment_results(
            base_path=model_subdir, existing_results_df=existing_results
        )
        if model_results_all_domains is None:
            continue

        # Add model and env columns to the vertical data
        model_results_all_domains["model"] = model
        model_results_all_domains["env"] = env

        # Merge with gold scores
        model_results_all_domains = pd.merge(
            model_results_all_domains,
            gold_scores[["domain_task_id", "gold_score", "gold_source"]],
            on="domain_task_id",
            how="left",
        )

        # Append to all_evals_env for this environment
        all_evals_env = pd.concat([all_evals_env, model_results_all_domains], ignore_index=True)

        # For backward compatibility, still create consolidated_evals for save_csv
        consolidated_evals = model_results_all_domains.pivot_table(
            index="domain_task_id", columns="config_name", values="eval", aggfunc="first"
        ).reset_index()

        # Append "_eval" to the config column names
        config_cols = [col for col in consolidated_evals.columns if col != "domain_task_id"]
        rename_dict = {col: f"{col}_eval" for col in config_cols}
        consolidated_evals.rename(columns=rename_dict, inplace=True)

        # Add back metadata columns for save_csv
        consolidated_evals = pd.merge(
            consolidated_evals,
            model_results_all_domains[["domain_task_id", "domain", "task_id", "env"]].drop_duplicates(),
            on="domain_task_id",
            how="left",
        )

        # Add experiments_path column
        consolidated_evals["experiments_path"] = model_subdir

        # Saves `evaluations.csv` at model subdirectory.
        save_path = Path(model_subdir) / results_csv_name
        save_csv(consolidated_evals, save_path, gold_scores=gold_scores)
        print(f"Evaluation results consolidated for {model} at {save_path}")

        # Compute confusion stats using the vertical data
        for config_name in model_results_all_domains["config_name"].unique():
            config_data = model_results_all_domains[model_results_all_domains["config_name"] == config_name]

            # Extract evaluation template from config_name
            eval_template = extract_eval_template_from_variation_name(config_name)
            if eval_template:
                print(f"Using template '{eval_template}' for config '{config_name}'")
            else:
                print(f"No template found for config '{config_name}', using default mapping")

            # Prepare the evaluation subset for this config
            evals_subset = config_data[["domain_task_id", "eval", "task_id"]].copy()

            # Overall stats for the entire environment for this config
            overall_confusion_stats = compute_confusion_stats(gold_scores, evals_subset, eval_template)
            overall_col_name = f"{model}--{config_name}--all"
            confusion_matrices[overall_col_name] = pd.Series(overall_confusion_stats)

            # Stats per domain subset
            for domain in gold_scores["domain"].unique():
                domain_evals_subset = config_data[config_data["domain"] == domain][["domain_task_id", "eval"]]
                domain_gold_scores = gold_scores[gold_scores["domain"] == domain]
                if not domain_evals_subset.empty:
                    domain_confusion_stats = compute_confusion_stats(
                        domain_gold_scores, domain_evals_subset, eval_template
                    )
                    domain_col_name = f"{model}--{config_name}--{domain}"
                    confusion_matrices[domain_col_name] = pd.Series(domain_confusion_stats)

            # Stats for eval_set_vwa subset
            if env == "vwa":
                eval_set_vwa = load_eval_set(EVAL_SET_VWA_PATH)
                gold_scores_vwa = gold_scores[gold_scores["domain_task_id"].isin(eval_set_vwa)]
                evals_subset_vwa = evals_subset[evals_subset["domain_task_id"].isin(eval_set_vwa)]
                if not gold_scores_vwa.empty and not evals_subset_vwa.empty:
                    vwa_confusion_stats = compute_confusion_stats(gold_scores_vwa, evals_subset_vwa, eval_template)
                    vwa_col_name = f"{model}--{config_name}--eval_set_vwa"
                    confusion_matrices[vwa_col_name] = pd.Series(vwa_confusion_stats)

            if env == "agrb_vwa":
                test_set_agrb = load_test_set_agrb(TEST_SET_AGRB_PATH)
                gold_scores["task_id"] = gold_scores["task_id"].apply(lambda x: x.replace(".resized.", "."))
                evals_subset["task_id"] = evals_subset["task_id"].apply(lambda x: x.replace(".resized.", "."))

                gold_scores_agrb = gold_scores[gold_scores["task_id"].isin(test_set_agrb["task_id"])]
                evals_subset_agrb = evals_subset[evals_subset["task_id"].isin(test_set_agrb["task_id"])]
                if not gold_scores_agrb.empty and not evals_subset_agrb.empty:
                    agrb_confusion_stats = compute_confusion_stats(gold_scores_agrb, evals_subset_agrb, eval_template)
                    agrb_col_name = f"{model}--{config_name}--test_set_agrb"
                    confusion_matrices[agrb_col_name] = pd.Series(agrb_confusion_stats)

    # Add a column for the model (LHS of '__') before saving all_evals_env.csv
    if not all_evals_env.empty and "model" in all_evals_env.columns:
        all_evals_env["model_name"] = all_evals_env["model"].apply(
            lambda x: x.split("__")[0] if isinstance(x, str) and "__" in x else x
        )
    all_evals_env.to_csv(out_dir_env / "all_evals_env.csv", index=False)
    print(f"All evals for {env} saved to", out_dir_env / "all_evals_env.csv")

    # Store the vertical data for this environment
    all_evals_per_env[env] = all_evals_env

    # === Convert the flat column names into a MultiIndex (domain, model, config) ===
    new_columns = []
    for col in confusion_matrices.columns:
        if "--" in col:
            parts = col.split("--")
            if len(parts) == 3:
                # Note that parts are in order: model, config, domain.
                # We want the MultiIndex in the order: (domain, model, config)
                new_columns.append((parts[2], parts[0], parts[1]))
            else:
                new_columns.append(("", col, ""))
    confusion_matrices.columns = pd.MultiIndex.from_tuples(new_columns, names=["domain", "model", "config"])
    # Optional: sort columns by the multi-index levels.
    confusion_matrices = confusion_matrices.sort_index(axis=1, level=["domain", "model", "config"])

    # Optional: sort the columns so that the first level (model) is in order.
    if env == "agrb_vwa":
        # Aggregate by env_domain, evaluator, config
        # get evaluator column gemini-2.5-flash-zero-random__agrb_claude-3.7 -> gemini-2.5-flash-zero-random

        # Filter for test cases only
        test_columns = [col for col in confusion_matrices.columns if "test" in str(col).lower()]
        confusion_matrices_test = confusion_matrices[test_columns]

        # Group by evaluator and config, then average the values
        grouped_data = {}
        for col in confusion_matrices_test.columns:
            domain, model, config = col
            # Extract evaluator from model (first part before '__')
            evaluator = model.split("__")[0] if isinstance(model, str) and "__" in model else model
            # Create new column key
            new_col = ("agrb_vwa_all", evaluator, config)

            if new_col not in grouped_data:
                grouped_data[new_col] = []
            grouped_data[new_col].append(confusion_matrices_test[col])

        # Calculate averages for each group
        averaged_data = {}
        for new_col, data_list in grouped_data.items():
            if len(data_list) == 1:
                averaged_data[new_col] = data_list[0]
            else:
                # Average across the columns, handling NaN values
                stacked_data = pd.concat(data_list, axis=1)
                averaged_data[new_col] = stacked_data.mean(axis=1, skipna=True)

        # Create DataFrame with averaged data and add to existing confusion_matrices
        averaged_df = pd.DataFrame(averaged_data)
        averaged_df.columns = pd.MultiIndex.from_tuples(list(averaged_data.keys()), names=["domain", "model", "config"])

        # Concatenate the averaged data with the existing confusion_matrices
        confusion_matrices = pd.concat([confusion_matrices, averaged_df], axis=1)
        # Sort by the new multi-index levels
        confusion_matrices = confusion_matrices.sort_index(axis=1, level=["domain", "model", "config"])

    confusion_matrices.T.to_csv(out_dir_env / "consolidated_stats.csv", index=True)
    print(f"Offline stats for {env} saved to", out_dir_env / "consolidated_stats.csv")
    confusion_per_env[env] = confusion_matrices


# Aggregate confusion matrices across all environments
# This creates a MultiIndex where the top level indicates the environment.
if confusion_per_env:
    # First, concatenate the confusion matrices from the different environments.
    consolidated_confusion = pd.concat(confusion_per_env, axis=1)
    # consolidated_output_path = Path(experiments_path) / "consolidated_stats_multilevel.csv"
    # consolidated_confusion.to_csv(consolidated_output_path, index=True)
    # print(f"Consolidated confusion matrices saved at {consolidated_output_path}")

    # --- Merge only the env and domain levels ---
    # We assume the MultiIndex is structured as follows:
    # Level 0: env
    # Level 1: domain
    # Level 2: model
    # Level 3: config
    #
    # Merge levels 0 and 1 (env and domain) into a single value.
    merged_env_domain = [
        f"{env}_{domain}"
        for env, domain in zip(
            consolidated_confusion.columns.get_level_values(0), consolidated_confusion.columns.get_level_values(1)
        )
    ]

    # Rebuild the MultiIndex with the merged env_domain level and preserve any remaining levels.
    if consolidated_confusion.columns.nlevels > 2:
        new_tuples = list(
            zip(
                merged_env_domain,
                consolidated_confusion.columns.get_level_values(2),
                consolidated_confusion.columns.get_level_values(3),
            )
        )
        new_index = pd.MultiIndex.from_tuples(new_tuples, names=["env_domain", "model", "config"])
    else:
        new_index = pd.Index(merged_env_domain, name="env_domain")

    consolidated_confusion.columns = new_index

    # Transpose the DataFrame so that the merged env_domain (and remaining levels)
    # become the row labels. Note that when saving a DataFrame with a MultiIndex header,
    # Pandas writes multiple header rows.
    transposed_confusion = consolidated_confusion.T

    transposed_output_path = Path(output_dir) / "consolidated_stats.csv"
    transposed_confusion.to_csv(transposed_output_path, index=True)
    print(f"Consolidated stats saved at {transposed_output_path}")

    # Save to XLSX
    transposed_confusion_flat = transposed_confusion.copy()
    transposed_confusion_flat.columns = [
        "_".join([str(level) for level in col if str(level) != ""]) if isinstance(col, tuple) else str(col)
        for col in transposed_confusion_flat.columns
    ]
    consolidated_xlsx_path = Path(output_dir) / "consolidated_stats.xlsx"
    print(f"Saving to {consolidated_xlsx_path}")
    # Reset index so index becomes a column, and avoid merged cells in Excel
    transposed_confusion_flat_reset = transposed_confusion_flat.reset_index()
    # Instead of overwriting the file, update only the values in the sheet, preserving formatting/filters
    wb = openpyxl.load_workbook(consolidated_xlsx_path)
    ws = wb["Sheet1"]

    # Clear only the cell values (not formatting/filters)
    for row in ws.iter_rows(min_row=2, max_row=ws.max_row, min_col=1, max_col=ws.max_column):
        for cell in row:
            cell.value = None
    for cell in ws[1]:
        cell.value = None

    # Write DataFrame values (including header)
    for col_idx, col_name in enumerate(transposed_confusion_flat_reset.columns, 1):
        ws.cell(row=1, column=col_idx, value=col_name)
    for row_idx, row in enumerate(transposed_confusion_flat_reset.itertuples(index=False), 2):
        for col_idx, value in enumerate(row, 1):
            ws.cell(row=row_idx, column=col_idx, value=value)

    wb.save(consolidated_xlsx_path)
    print(f"Consolidated stats saved at {consolidated_xlsx_path}")

else:
    print("No confusion matrices to consolidate.")

# Aggregate all vertical evaluation data from all environments
if all_evals_per_env:
    all_evals_consolidated = pd.concat(all_evals_per_env.values(), ignore_index=True)
    all_evals_output_path = Path(output_dir) / "all_evals_consolidated.csv"
    all_evals_consolidated.to_csv(all_evals_output_path, index=False)
    print(f"All consolidated vertical evals saved at {all_evals_output_path}")
else:
    print("No vertical evals to consolidate.")
    all_evals_consolidated = pd.concat(all_evals_per_env.values(), ignore_index=True)
    all_evals_output_path = Path(output_dir) / "all_evals_consolidated.csv"
    all_evals_consolidated.to_csv(all_evals_output_path, index=False)
    print(f"All consolidated vertical evals saved at {all_evals_output_path}")
    print(f"All consolidated vertical evals saved at {all_evals_output_path}")
