import json
import os

class AxiomLoader:
    """
    Loads axioms from JSON files in a specified directory.
    """

    def __init__(self, axiom_directory):
        """
        Initializes the AxiomLoader with the directory containing axiom files.

        Args:
            axiom_directory (str): The path to the directory containing axiom JSON files.
        """
        self.axiom_directory = axiom_directory
        self.axioms = {}

    def load_axioms(self):
        """
        Loads all axiom files from the specified directory and merges them into a single dictionary.
        """
        for filename in os.listdir(self.axiom_directory):
            if filename.endswith(".json"):
                filepath = os.path.join(self.axiom_directory, filename)
                try:
                    with open(filepath, "r") as f:
                        axiom_data = json.load(f)
                        axiom_set_name = axiom_data.get("axiom_set_name", "Unnamed Axiom Set")
                        self.axioms[axiom_set_name] = axiom_data
                except FileNotFoundError:
                    print(f"Error: Axiom file not found: {filepath}")
                except json.JSONDecodeError:
                    print(f"Error: Invalid JSON format in file: {filepath}")
                except Exception as e:
                    print(f"Error loading axiom file {filename}: {e}")

    def get_axioms(self):
        """
        Returns all loaded axioms.

        Returns:
            dict: A dictionary where keys are axiom set names and values are the corresponding axiom data.
        """
        return self.axioms

    def get_axiom_set(self, axiom_set_name):
        """
        Retrieves a specific set of axioms by its name.

        Args:
            axiom_set_name (str): The name of the axiom set to retrieve.

        Returns:
            dict: The axiom data for the specified set, or None if not found.
        """
        return self.axioms.get(axiom_set_name)

# Example usage (to demonstrate, not part of the class itself):
if __name__ == "__main__":
    # Assuming your axiom directory is named 'data/axioms' relative to the script
    axiom_dir = os.path.join(os.path.dirname(__file__), "../data/axioms")
    axiom_loader = AxiomLoader(axiom_dir)
    axiom_loader.load_axioms()

    all_axioms = axiom_loader.get_axioms()
    print("All Loaded Axioms:", all_axioms)

    ralph_axioms = axiom_loader.get_axiom_set("Ralph Matrix True Method Axioms")
    if ralph_axioms:
        print("\nRalph Matrix Axioms:", ralph_axioms)
    else:
        print("\nRalph Matrix Axioms not found.")