import os
import argparse
import yaml
from dotenv import load_dotenv
from colmind import MasterProcess

OPENAI_KEY="ADD_OPENAI_KEY"
TOGETHER_API_KEY = "ADD_TOGETHER_KEY"
MISTRAL_API_KEY = "ADD_MISTRAL_KEY"
os.environ["TOGETHER_API_KEY"] = TOGETHER_API_KEY
os.environ["MISTRAL_API_KEY"] = MISTRAL_API_KEY
os.environ["OPENAI_API_KEY"] = OPENAI_KEY


def main(config):
    load_dotenv()
    master = MasterProcess(config)
    master.start(cooperation=False, learning=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run the ColMind collaborative experiment.")
    parser.add_argument("--config_path", type=str, help="Path to the YAML config file.")
    parser.add_argument("--resume", action="store_true", help="Resume the experiment.")

    args = parser.parse_args()
    with open(args.config_path, 'r') as file:
        config = yaml.safe_load(file)
    main(config | {"resume": args.resume})
