"""Define Logger class for logging information to stdout and disk."""

import json

import os

from os.path import join

from pytorch_lightning.loggers import WandbLogger

from pytorch_lightning.loggers.test_tube import TestTubeLogger

from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

def get_ckpt_dir(save_path, exp_name):

    return os.path.join(save_path, exp_name, "ckpts")

def get_ckpt_callback(save_path, exp_name):

    ckpt_dir = os.path.join(save_path, exp_name, "ckpts")

    return ModelCheckpoint(dirpath=ckpt_dir,

                           save_top_k=1,

                           verbose=True,

                           monitor='val_loss',

                           mode='min')

def get_early_stop_callback(patience=10):

    return EarlyStopping(monitor='val_loss',

                         patience=patience,

                         verbose=True,

                         mode='min')

def get_logger(logger_type, save_path, exp_name, project_name=None):

    if logger_type == 'wandb':

        if project_name is None:

            raise ValueError("Must supply project name when using wandb logger.")

        return WandbLogger(name=exp_name,

                           project=project_name)

    elif logger_type == 'test_tube': 

        exp_dir = os.path.join(save_path, exp_name)

        return TestTubeLogger(save_dir=exp_dir,

                              name='lightning_logs',

                              version="0")

    else:

        raise ValueError(f'{logger_type} is not a supported logger.')