import os
from argparse import ArgumentParser

import pytorch_lightning as pl
import torch

import models
import datamodules

from transformers_lightning import utils, callbacks, datamodules
from transformers_lightning.defaults import DefaultConfig


# Print high precision tensor values
torch.set_printoptions(precision=16)

def main(hyperparameters):

    # instantiate PL model
    model = TransformerModel(hyperparameters)

    # default tensorboard logger
    test_tube_logger = pl.loggers.TestTubeLogger(
        os.path.join(hyperparameters.output_dir, hyperparameters.tensorboard_dir), name=hyperparameters.name
    )

    # Save pre-trained models to
    save_transformers_callback = callbacks.TransformersModelCheckpointCallback(hyperparameters)

    # instantiate PL trainer
    trainer = pl.Trainer.from_argparse_args(
        hyperparameters,
        default_root_dir=hyperparameters.output_dir,
        profiler=True,
        logger=test_tube_logger,
        callbacks=[save_transformers_callback],
        log_gpu_memory='all',
        weights_summary='full'
    )

    # Datasets
    datamodule = YourDataModule(hyperparameters, model, trainer)

    # Train!
    if datamodule.do_train():
        trainer.fit(model, datamodule=datamodule)

    # Test!
    if datamodule.do_test():
        trainer.test(model, datamodule=datamodule)


if __name__ == '__main__':

    parser = ArgumentParser()

    # Experiment name, used both for checkpointing, pre_trained_names, logging and tensorboard
    parser.add_argument('--name', type=str, required=True, help='Name of the experiment, well be used to correctly retrieve checkpoints and logs')

    # I/O folders
    DefaultConfig.add_argparse_args(parser)

    # add model specific cli arguments
    TransformerModel.add_argparse_args(parser)
    YourDataModule.add_argparse_args(parser)

    # add callback / logger specific cli arguments
    callbacks.TransformersModelCheckpointCallback.add_argparse_args(parser)

    # add all the available trainer options to argparse
    # ie: now --gpus --num_nodes ... --fast_dev_run all work in the cli
    pl.Trainer.add_argparse_args(parser)

    # get NameSpace of parameters
    hyperparameters = parser.parse_args()
    main(hyperparameters)