from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.callbacks.early_stopping import EarlyStopping
from pytorch_lightning.loggers import WandbLogger

import hydra
import pytorch_lightning as pl

from dadeit.model import LitDADeiT


@hydra.main(config_path=".", config_name="common")
def run(config):
    pl.seed_everything(config.seed)

    logger = None
    if config.logger:
        logger = WandbLogger(
            project=config.project,
            name=config.model_name,
            config=config,
        )

    callbacks = []
    callbacks.append(ModelCheckpoint(
        dirpath=config.checkpoint_dir,
        filename=f"{config.model_name}-"+"{epoch:02d}-"+"{val_loss:.4f}",
        monitor="val_loss",
        mode="min"
    ))

    trainer = pl.Trainer(
        precision=16,
        gradient_clip_val=0.5,
        callbacks=callbacks,
        deterministic=True,
        check_val_every_n_epoch=1,
        gpus=config.gpus,
        logger=logger,
        max_epochs=config.max_epochs,
        weights_summary="top",
        accelerator="ddp",
    )

    model = LitDADeiT(config)

    if config.lr_finder:
        lr_finder = trainer.tuner.lr_find(
        model, min_lr=1e-8, max_lr=1e-1, num_training=100)
        model.hparams.lr = lr_finder.suggestion()
        print(model.hparams.lr)
    else:
        trainer.fit(model)
        trainer.test()


if __name__ == '__main__':
    run()