import os
import random
import re
import json
import numpy as np
from pathlib import Path
from natsort import natsorted
import jax


def objects_to_txt():
    curdir = "/".join(__file__.split("/")[:-1])

    with open(os.path.join(curdir, "pickupables_prior.json")) as f:
        pickupable_prior = json.load(f)
    with open(os.path.join(curdir, "receptacles_prior.json")) as f:
        receptacle_prior = json.load(f)

    def split_camel_preserve_acronyms(name):
        # Insert space between lowercase → uppercase
        # OR between acronym → normal word
        s = re.sub(r"(?<=[a-z])(?=[A-Z])", " ", name)
        s = re.sub(r"(?<=[A-Z])(?=[A-Z][a-z])", " ", s)
        return s.lower()

    objects = set(pickupable_prior.keys()) | set(receptacle_prior.keys())
    objects = {split_camel_preserve_acronyms(obj) for obj in objects}

    # Save to txt file
    with open(os.path.join(curdir, "objects.txt"), "w") as f:
        for obj in sorted(objects):
            f.write(f"{obj}\n")


def seed_everything(seed: int = 42) -> jax.random.PRNGKey:
    random.seed(seed)
    np.random.seed(seed)
    # Make hashing deterministic
    os.environ["PYTHONHASHSEED"] = str(seed)
    return jax.random.PRNGKey(seed)

def is_sublevel(name: str) -> str:
    # Check if receptacle is a sublevel one by looking for '___' pattern
    return bool(re.search(r"___", name))


def format_receptacle(name: str) -> str:
    # Extract all digit sequences and join them
    numbers = "".join(re.findall(r"\d+", name))

    # Remove everything after last non-digit separator, then strip trailing '|' or whitespace
    label = re.split(r"\d+", name)[0].rstrip("| ").strip()

    return f"{label} # {numbers}" if numbers else label


def evaluate_runs_accuracy(res_root: str) -> dict:
    folder_path = Path(res_root)
    run_folders = natsorted(folder_path.glob("run_*"))
    scene_accuracy = 0.0
    num_runs = len(run_folders)
    for i, run_folder in enumerate(run_folders):
        accuracy_file = run_folder / "accuracy.txt"
        if accuracy_file.exists():
            with open(accuracy_file, "r") as f:
                run_accuracy = float(f.read().strip())
            scene_accuracy += run_accuracy / num_runs
            print(f"Run {i} ({run_folder.name}): Accuracy = {run_accuracy:.4f}")
        else:
            print(f"Run {i} ({run_folder.name}): Accuracy file not found.")

    print("####################################")
    print(f"Overall accuracy: {scene_accuracy}")
    print("####################################")


if __name__ == "__main__":
    run_id = 453
    split = "test"
    seed = 0
    # objects_to_txt()
    evaluate_runs_accuracy(
        f"generated_data/semistaticsim/{split}/{run_id}/{seed}/privileged"
    )  # Replace with actual path
