import subprocess
import argparse

# python run_experiments.py --dataset "Birds-Nest" --model yolov10n yolov10s yolov10m yolov10l

def run_experiment(base_command, run_mode, use_pretrained):
    """
    Constructs and runs a single experiment command.
    """
    command = base_command + ["--run", run_mode]
    if use_pretrained:
        command.append("--pretrained")

    print("="*80)
    print(f"Starting run: {run_mode}")
    print(f"Command: {' '.join(command)}")
    print("="*80)

    try:
        # subprocess.run is a blocking call, ensuring sequential execution.
        subprocess.run(command, check=True)
        print(f"\nSUCCESS: Run '{run_mode}' completed.\n")
    except subprocess.CalledProcessError as e:
        print(f"\nERROR: Run '{run_mode}' failed with exit code {e.returncode}.\n")
        # Decide if you want to stop all subsequent runs on failure
        # raise e # Uncomment to stop the entire sequence on error

def main():
    """
    Parses arguments and launches a sequence of training experiments for each specified model.
    """
    parser = argparse.ArgumentParser(description="Run a sequence of YOLO training experiments.")

    # Define arguments that will be common to all training runs
    parser.add_argument('--dataset', type=str, required=True, choices=["Birds-Nest", "Common-VALID", "Electric-Substation", "InsPLAD-det"], help='Dataset name to be used.')
    parser.add_argument("--model", nargs='+', required=True, choices=["yolov8n", "yolov8s", "yolov8m", "yolov8l", "yolov10n", "yolov10s", "yolov10m", "yolov10l"], help="One or more models to use for the experiments.")
    parser.add_argument("--epochs", type=int, default=1000, help="Number of epochs.")
    parser.add_argument("--batch", type=int, default=16, help="Batch size.")
    parser.add_argument("--plots", action="store_true", default=True, help="Generate plots for all runs.")

    args = parser.parse_args()

    # Define the sequence of experiments to run for each model
    # Each tuple is (run_mode, use_pretrained_flag)
    experiment_sequence = [
        ("From_Scratch",    False),
        ("Finetuning",      True),
        ("freeze_[P1-P3]",  True),
        ("freeze_Backbone", True),
        ("freeze_[P1-23]",  True)
    ]

    # Iterate over each specified model variant
    for model_name in args.model:
        print(f"\n{'='*25} Starting Experiments for Model: {model_name.upper()} {'='*25}\n")

        # Base command list, specific to the current model
        base_command = [
            "python", "main.py",
            "--dataset", args.dataset,
            "--epochs", str(args.epochs),
            "--batch", str(args.batch),
            "--model", model_name
        ]
        if args.plots:
            base_command.append("--plots")

        # Execute each experiment in sequence for the current model
        for run_mode, use_pretrained in experiment_sequence:
            run_experiment(base_command, run_mode, use_pretrained)

    print("="*80)
    print("All experiments for all specified models have been completed.")
    print("="*80)

if __name__ == "__main__":
    main()