import argparse
import subprocess

def run_transfer_learning(dataset, epochs, batch_size, imgsz, patience, cache, pretrained, cos_lr, profile, plots, resume, augment, model, run):
    # Construct the command with user-provided arguments
    command = [
        "python", "./experiment/transfer_learning_train&test.py",
        "--dataset", str(dataset),
        "--epochs", str(epochs),
        "--batch", str(batch_size),
        "--imgsz", str(imgsz),
        "--patience", str(patience),
        "--cache", cache,
        "--model", model,
        "--run", run
    ]

    # Append boolean flags conditionally
    if pretrained:
        command.append("--pretrained")
    if cos_lr:
        command.append("--cos_lr")
    if profile:
        command.append("--profile")
    if plots:
        command.append("--plots")
    if resume:
        command.append("--resume")
    if augment:
        command.append("--augment")
    
    # Use subprocess to run the script with the arguments
    subprocess.run(command, check=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run transfer learning with YOLO model.")
    
    # Define all arguments with default values from your script
    parser.add_argument('--dataset', type=str, choices=["Birds-Nest", "Common-VALID", "Electric-Substation", "InsPLAD-det"], help='Dataset name to be used')
    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("--imgsz", type=int, default=640, help="Image size")
    parser.add_argument("--patience", type=int, default=30, help="Patience for early stopping")
    parser.add_argument("--cache", type=str, default='ram', help="Cache option")
    parser.add_argument("--pretrained", action="store_true", help="Use pretrained weights")
    parser.add_argument("--cos_lr", action="store_true", help="Use cosine learning rate")
    parser.add_argument("--profile", action="store_true", help="Profile the training")
    parser.add_argument("--plots", action="store_true", help="Generate plots")
    parser.add_argument("--resume", action="store_true", help="Resume run")
    parser.add_argument("--augment", action="store_true", help="Apply augmentation techniques during training")
    parser.add_argument("--model", type=str, choices=["yolov8n", "yolov8s", "yolov8m", "yolov8l", "yolov10n", "yolov10s", "yolov10m", "yolov10l"], help="Model to use")
    parser.add_argument("--run", type=str, choices=["From_Scratch", "Finetuning", "freeze_[P1-P3]", "freeze_Backbone", "freeze_[P1-23]"], help="Run mode")

    args = parser.parse_args()

    # Call the function to run transfer learning
    run_transfer_learning(
        dataset=args.dataset,
        epochs=args.epochs, 
        batch_size=args.batch, 
        imgsz=args.imgsz, 
        patience=args.patience, 
        cache=args.cache, 
        pretrained=args.pretrained, 
        cos_lr=args.cos_lr, 
        profile=args.profile, 
        plots=args.plots,
        resume=args.resume,
        augment=args.augment,
        model=args.model, 
        run=args.run
    )