import argparse
from src.acrlAgent import ACRLAgent, ACRLMultiAgent

def main():
    # Set up command line argument parsing
    parser = argparse.ArgumentParser(description='Train the ACRLAgent model.')
    parser.add_argument('--T0', type=int, default=80, help='Maximum steps per episode.')
    parser.add_argument('--total_timesteps', type=int, default=2000000, help='Total number of timesteps for training.')
    parser.add_argument('--filename', type=str, default='ppo_acrl', help='Filename to save the trained model.')
    parser.add_argument('--verbose', type=int, default=0, help='Verbosity level for PPO algorithm log.')
    parser.add_argument('--building', type=int, nargs='+', required=True, help='List of building types for each agent. Use a single value for single-agent training.')
    parser.add_argument('--dual', action='store_true', help='Whether or not to use the augmented state space.')
    parser.add_argument('--nu', type=float, default=1, help='Penalization for local constraint.')
    parser.add_argument('--lambd', type=float, default=1, help='Penalization for global constraint.')
    # Parse arguments from the command line
    args = parser.parse_args()

    # Determine the number of agents from the length of the building list
    num_agents = len(args.building)

    # Adjust filename based on the building argument
    if num_agents > 1:
        building_str = "_".join(map(str, args.building))
        filename = f"{args.filename}_multi_b{building_str}_n{num_agents}_T{args.total_timesteps}"
    else:
        filename = f"{args.filename}_b{args.building[0]}_T{args.total_timesteps}"

    if args.dual:
        filename += "_dd"  # Dual dynamics
    else:
        nu = str(args.nu).replace('.', '_')
        lambd = str(args.lambd).replace('.', '_')
        filename += f"_nu{nu}_lambd{lambd}"

    # Create an instance of ACRLAgent or ACRLMultiAgent based on the number of agents
    if num_agents > 1:
        acrl = ACRLMultiAgent(
            verb=args.verbose,
            building_types=args.building,
            nu_val=args.nu,
            lambda_val=args.lambd,
            n_agents=num_agents
        )
    else:
        acrl = ACRLAgent(
            verb=args.verbose,
            currentBuilding=args.building[0],
            derenv=args.dual,
            nu_val=args.nu,
            lambda_val=args.lambd
        )

    # Train the model with the specified parameters
    acrl.train(T0=args.T0, total_timesteps=args.total_timesteps, model_name=filename)

    if args.verbose == 1:
        print(f"Training complete. Model saved as {filename}") 

if __name__ == "__main__":
    main()
