import os
import hydra
import wandb
import omegaconf
import warnings
from pathlib import Path
import sklearn.exceptions
warnings.filterwarnings('ignore', category=sklearn.exceptions.UndefinedMetricWarning)
from lib.da.tools.trainer import cross_domain_trainer

def make_wandb_tags(config):
    tags = []
    for key in ['experiment_name', 'seed']:
        if key in config:
            val = config[key]
            if key == 'seed':
                val = 'seed_' + str(val)
            tags.append(val)
    return tags

def setup_wandb(config):
    print('Setting up logging to Weights & Biases.')
    config_dict = omegaconf.OmegaConf.to_container(config, resolve=True, throw_on_missing=True)

    wandb_proj_name = '_'.join([config.experiment_name,
                                #config.run_description,
                                config.dre_method, config.dataset])
    # hydra changes working directories
    log_dir = str(Path.joinpath(Path(os.getcwd()), config.run_params.log_dir))
    # make 'wandb' path, otherwise WSL might block writing to dir
    wandb_path = Path.joinpath(Path(log_dir), 'wandb')
    wandb_path.mkdir(exist_ok=True, parents=True)
    wandb.login()
    # tracks everything that TensorBoard tracks
    # writes to same dir as TensorBoard
    wandb.tensorboard.patch(root_logdir=log_dir)
    wandb_run = wandb.init(project=wandb_proj_name, name=log_dir,
                           dir=log_dir, save_code=False, config=config_dict)
    wandb_run.tags = wandb_run.tags + tuple(make_wandb_tags(config_dict))
    print(f'Writing Weights & Biases logs to: {str(wandb_path)}')
    return wandb_run


@hydra.main(config_path='lib/da/configs', config_name='config')
def main(config):
    if config.use_wandb:
        wandb_run = setup_wandb(config)
    trainer = cross_domain_trainer(config)
    trainer.train()
    # necessary for Hydra multiruns
    wandb.finish()
    wandb.tensorboard.unpatch()


if __name__ == '__main__':
    main()
