""" logging_utils.py
    Utility functions for logging experiments to CometML and TensorBoard
    Developed as part of DeepThinking project
    July 2021
"""

import os

import comet_ml
from torch.utils.tensorboard import SummaryWriter

# Ignore statements for pylint:
#     Too many branches (R0912), Too many statements (R0915), No member (E1101),
#     Not callable (E1102), Invalid name (C0103), No exception (W0702),
#     Too many local variables (R0914), Missing docstring (C0116, C0115).
# pylint: disable=R0912, R0915, E1101, E1102, C0103, W0702, R0914, C0116, C0115


def log_to_comet(comet_exp, train_acc, val_acc, test_acc, epoch, out_str=None):
    metrics = {}
    for key in train_acc.keys():
        metrics.update(
            {
                f"val_acc_{key}": val_acc[key],
                f"train_acc_{key}": train_acc[key],
                f"test_acc_{key}": test_acc[key],
            }
        )
    comet_exp.log_metrics(metrics, epoch=epoch)
    if out_str:
        comet_exp.log_parameter("ckpt_path", out_str)


def setup_comet(args):
    if args.use_comet:
        experiment = comet_ml.Experiment(
            project_name="deepthinking",
            auto_param_logging=False,
            auto_metric_logging=False,
            disabled=False,
            parse_args=False,
        )

        experiment.add_tag(args.train_mode)
        experiment.log_parameters(vars(args))
        return experiment
    return None


def setup_tb(train_log, output):
    writer = SummaryWriter(log_dir=f"{output}/runs/{train_log[:-4]}")
    return writer


def write_to_tb(stats, stat_names, epoch, writer):
    for name, stat in zip(stat_names, stats):
        stat_name = os.path.join("val", name)
        writer.add_scalar(stat_name, stat, epoch)
