from typing import Dict, Any, Tuple, Callable
from datetime import timedelta
from spacy.util import registry
from spacy.errors import Errors
from wasabi import msg


@registry.loggers("spacy-ray.ConsoleLogger.v1")
def ray_console_logger():
    def setup_printer(
        nlp: "Language",
    ) -> Tuple[Callable[[Dict[str, Any]], None], Callable]:
        score_cols = list(nlp.config["training"]["score_weights"])
        score_widths = [max(len(col), 6) for col in score_cols]
        loss_cols = [f"Loss {pipe}" for pipe in nlp.pipe_names]
        loss_widths = [max(len(col), 8) for col in loss_cols]
        table_header = ["T", "E", "#", "W"] + loss_cols + score_cols + ["Score"]
        table_header = [col.upper() for col in table_header]
        table_widths = [8, 3, 6, 6] + loss_widths + score_widths + [6]
        table_aligns = ["r" for _ in table_widths]
        msg.row(table_header, widths=table_widths)
        msg.row(["-" * width for width in table_widths])

        def log_step(info: Dict[str, Any]):
            try:
                losses = [
                    "{0:.2f}".format(float(info["losses"][pipe_name]))
                    for pipe_name in nlp.pipe_names
                ]
            except KeyError as e:
                raise KeyError(
                    Errors.E983.format(
                        dict="scores (losses)",
                        key=str(e),
                        keys=list(info["losses"].keys()),
                    )
                ) from None

            try:
                scores = [
                    "{0:.2f}".format(float(info["other_scores"].get(col, 0.0)) * 100)
                    for col in score_cols
                ]
            except KeyError as e:
                raise KeyError(
                    Errors.E983.format(
                        dict="scores (other)",
                        key=str(e),
                        keys=list(info["other_scores"].keys()),
                    )
                ) from None
            time = timedelta(seconds=info["seconds"])
            data = (
                [str(time), info["epoch"], info["step"], info["words"]]
                + losses
                + scores
                + ["{0:.2f}".format(float(info["score"]))]
            )
            msg.row(data, widths=table_widths, aligns=table_aligns)

        def finalize():
            pass

        return log_step, finalize

    return setup_printer
