import json
import numpy as np
import matplotlib.pyplot as plt
from src.marlAcrl import marlAcrl


def load_config(config_path):
    """Load configuration from a JSON file."""
    with open(config_path, 'r') as f:
        return json.load(f)


def main(config_path):
    # Load configuration
    config = load_config(config_path)

    # Initialize the MARL Manager
    manager = marlAcrl(
        connections=config["connections"],
        buildings=config["buildings"],
        models=config["models"],
        nu_val=config["nu_val"],
        lambda_val=config["lambda_val"],
        derenv=config["derenv"]
    )

    # Set constraints
    manager.set_costraint(thresh=config["constraint_thresh"])

    # Run experiments
    print(f"Running main experiment for {config['epochs']} epochs...")
    manager.run_episode(
        epochs=config["epochs"],
        consensus=True,
        lamdaStep=0.01,
        consensus_steps=config["consensus_steps"]
    )

    # Compare algorithms
    print(f"Comparing algorithms for {config['comparison_epochs']} epochs...")
    manager.compare_algs(
        epochs=config["comparison_epochs"],
        consensus_steps=config["consensus_steps"]
    )
    manager.compare_algs2(
        epochs=config["comparison_epochs"],
        consensus_steps=config["consensus_steps"],
        from_start=True,
        cum=True,
        num_runs=config["num_runs"]
    )
    manager.compare_algs3(
        epochs=config["epochs"],
        consensus_steps=config["consensus_steps"],
        from_start=True,
        cum=True,
        constraints=config["comparison_constraints"],
        run=True
    )
    print("Experiments completed successfully!")


if __name__ == "__main__":
    import argparse

    # Argument parser for command-line execution
    parser = argparse.ArgumentParser(description="Run MARL experiments.")
    parser.add_argument(
        "--config", type=str, required=True,
        help="Path to the configuration file (e.g., config.json)."
    )
    args = parser.parse_args()

    # Run the main script with the provided config
    main(args.config)
