from lightning import Trainer, seed_everything
from lightning.pytorch.callbacks import ModelCheckpoint, TQDMProgressBar

from acfg.modelconfig import ModelConfig
from ml.app.data import ImageDataModule
from ml.app.lm import ClassificationModule
from ml.app.models.classification import DiseaseClassificationModel


ckpt_callback = ModelCheckpoint(
    filename="classification" + "_{epoch:02d}_{VA:.2f}",
    save_top_k=1,
    mode="min",
    monitor=ModelConfig.VAL_LOSS,
)

tqdm_callback = TQDMProgressBar(refresh_rate=10)


model = DiseaseClassificationModel(ModelConfig.PRETRAINED_MODEL_NAME)

datamodule = ImageDataModule(
    train_path=ModelConfig.TRAIN_DATA_PATH,
    val_path=ModelConfig.VAL_DATA_PATH,
    test_path=ModelConfig.TEST_DATA_PATH,
    batch_size=ModelConfig.BATCH_SIZE,
    img_size=ModelConfig.IMG_SIZE,
)

l_module = ClassificationModule(
    model=model,
    num_classes=ModelConfig.NUM_OUTPUT_CLASSES,
)

seed_everything(42)
trainer = Trainer(
    max_epochs=25,
    callbacks=[ckpt_callback, tqdm_callback],
    num_sanity_val_steps=2,
)


if __name__ == "__main__":
    trainer.fit(
        model=l_module,
        datamodule=datamodule,
    )