from anomalib.data import MVTec
from anomalib.engine import Engine
from anomalib.models import EfficientAd

def main():
    # 1. Initialize the MVTec datamodule (defaults to downloading/using the 'bottle' category or similar)
    datamodule = MVTec(
        category="bottle",
        train_batch_size=16,
        eval_batch_size=16,
    )

    # 2. Choose an anomaly detection model 
    model = EfficientAd()

    # 3. Initialize the PyTorch Lightning based Engine 
    # (Setting max_epochs low here just for a quick demonstration)
    engine = Engine(
        max_epochs=5,
        default_root_dir="results"
    )

    print("Starting training...")
    # 4. Train the model
    engine.fit(datamodule=datamodule, model=model)

    print("Running prediction/inference on test images...")
    # 5. Run prediction using the best checkpoint generated during training
    predictions = engine.predict(
        datamodule=datamodule,
        model=model,
        ckpt_path=engine.trainer.checkpoint_callback.best_model_path,
    )

    print(f"Predictions complete. Results generated for {len(predictions)} batches.")

if __name__ == "__main__":
    main()