ezpz.log⚓︎

def get_active_enrich_handlers(logger: logging.Logger) -> list:
    """Return ``(index, handler)`` pairs for active ``RichHandler`` instances."""
    from ezpz.log.handler import RichHandler as EnrichHandler

    return [
        (idx, h)
        for idx, h in enumerate(logger.handlers)
        if isinstance(h, EnrichHandler)
    ]

def get_console_from_logger(logger: logging.Logger) -> Console:
    """Return the ``Console`` attached to *logger* or synthesise a new one."""
    from ezpz.log.handler import RichHandler as EnrichHandler

    for handler in logger.handlers:
        if isinstance(handler, (RichHandler, EnrichHandler)):
            return handler.console  # type: ignore
    from ezpz.log.console import get_console

    return get_console()

def get_enrich_logging_config_as_yaml(
    name: str = "enrich", level: str = "INFO"
) -> str:
    """Render the Enrich logging YAML snippet with the requested name/level."""
    return rf"""
    ---
    # version: 1
    handlers:
      {name}:
        (): ezpz.log.handler.RichHandler
        show_time: true
        show_level: true
        enable_link_path: false
        level: {level.upper()}
    root:
      handlers: [{name}]
    disable_existing_loggers: false
    ...
    """

def get_file_logger(
    name: Optional[str] = None,
    level: str = "INFO",
    rank_zero_only: bool = True,
    fname: Optional[str] = None,
    # rich_stdout: bool = True,
) -> logging.Logger:
    """Create a file-backed logger, optionally emitting only on rank zero."""
    # logging.basicConfig(stream=DummyTqdmFile(sys.stderr))
    import logging

    from ezpz.dist import get_rank

    fname = "output" if fname is None else fname
    log = logging.getLogger(name)
    if rank_zero_only:
        fh = logging.FileHandler(f"{fname}.log")
        if get_rank() == 0:
            log.setLevel(level)
            fh.setLevel(level)
        else:
            log.setLevel("CRITICAL")
            fh.setLevel("CRITICAL")
    else:
        fh = logging.FileHandler(f"{fname}-{get_rank()}.log")
        log.setLevel(level)
        fh.setLevel(level)
    # create formatter and add it to the handlers
    formatter = logging.Formatter(
        "[%(asctime)s][%(name)s][%(levelname)s] - %(message)s"
    )
    fh.setFormatter(formatter)
    log.addHandler(fh)
    return log

def get_logger(
    name: Optional[str] = None,
    level: Optional[str] = None,
    rank_zero_only: bool = True,
    rank: Optional[int | str] = None,
    colored_logs: Optional[bool] = True,
) -> logging.Logger:
    """Return a logger initialised with the project's logging configuration."""
    if rank is None and rank_zero_only:
        from ezpz.dist import get_rank

        rank = get_rank()
    assert rank is not None
    # if is_interactive():
    #     return get_rich_logger(name=name, level=level)
    ezpz_log_level = (
        os.environ.get("EZPZ_LOG_LEVEL", os.environ.get("LOG_LEVEL", "INFO"))
        if level is None
        else level
    )
    if not colored_logs:
        os.environ["NO_COLOR"] = "1"

    logging_config = get_logging_config()

    LOG_FROM_ALL_RANKS = os.environ.get(
        "LOG_FROM_ALL_RANKS", os.environ.get("EZPZ_LOG_FROM_ALL_RANKS")
    )
    if LOG_FROM_ALL_RANKS is not None and to_bool(LOG_FROM_ALL_RANKS):
        LOG_FROM_ALL_RANKS = True
        print(
            f"[{rank:>2}] Logging from all ranks at level {ezpz_log_level} enabled via EZPZ_LOG_FROM_ALL_RANKS"
        )
        logging_config["handlers"]["term"] |= {"rank": rank}

    logging.config.dictConfig(logging_config)
    logger = logging.getLogger(name if name is not None else __name__)
    if LOG_FROM_ALL_RANKS:
        logger.setLevel(ezpz_log_level)
        return logger

    elif rank_zero_only:
        if int(rank) == 0:
            logger.setLevel(ezpz_log_level)
        else:
            logger.setLevel("CRITICAL")
    else:
        logger.setLevel(ezpz_log_level)

    return logger

def get_logger_new(
    name: str,
    level: str = "INFO",
):
    """Return a logger configured solely via the Enrich YAML template."""
    import yaml

    config = yaml.safe_load(
        get_enrich_logging_config_as_yaml(name=name, level=level),
    )
    logging.config.dictConfig(config)
    log = logging.getLogger(name=name)
    log.setLevel(level)
    return log

def get_rich_logger(
    name: Optional[str] = None, level: Optional[str] = None
) -> logging.Logger:
    """Return a logger backed by a single :class:`RichHandler`."""
    from ezpz.dist import get_world_size
    from ezpz.log.handler import RichHandler

    level = "INFO" if level is None else level
    # log: logging.Logger = get_logger(name=name, level=level)
    log = logging.getLogger(name)
    log.handlers = []
    console = get_console(
        markup=True,
        redirect=(get_world_size() > 1),
    )
    handler = RichHandler(
        level,
        rich_tracebacks=False,
        console=console,
        show_path=False,
        enable_link_path=False,
    )
    log.handlers = [handler]
    log.setLevel(level)
    return log

def make_layout(ratio: int = 4, visible: bool = True) -> Layout:
    """Define the layout."""
    layout = Layout(name="root", visible=visible)
    layout.split_row(
        Layout(name="main", ratio=ratio, visible=visible),
        Layout(name="footer", visible=visible),
    )
    return layout

def print_config(
    config: DictConfig | dict | Any,
    resolve: bool = True,
    print_order: Sequence[str] | None = None,
) -> None:
    """Prints content of DictConfig using Rich library and its tree structure.

    Args:
        config (DictConfig): Configuration composed by Hydra.
        print_order (Sequence[str], optional): Determines in what order config
            components are printed.
        resolve (bool, optional): Whether to resolve reference fields of
            DictConfig.
    """
    import pandas as pd

    tree = rich.tree.Tree("CONFIG")  # , style=style, guide_style=style)
    quee = []
    if print_order:
        quee.extend([f for f in print_order if f not in quee])
    for f in config:
        if f not in quee:
            quee.append(f)
    dconfig = {}
    for f in quee:
        branch = tree.add(f)  # , style=style, guide_style=style)
        config_group = config[f]
        if isinstance(config_group, DictConfig):
            branch_content = OmegaConf.to_yaml(config_group, resolve=resolve)
            cfg = OmegaConf.to_container(config_group, resolve=resolve)
        else:
            branch_content = str(config_group)
            cfg = str(config_group)
        dconfig[f] = cfg
        branch.add(rich.syntax.Syntax(branch_content, "yaml"))
    outfile = Path(os.getcwd()).joinpath("config_tree.log")
    from rich.console import Console

    with outfile.open("wt") as f:
        console = Console(file=f)
        console.print(tree)
    with open("config.json", "w") as f:
        f.write(json.dumps(dconfig))
    cfgfile = Path("config.yaml")
    OmegaConf.save(config, cfgfile, resolve=True)
    cfgdict = OmegaConf.to_object(config)
    logdir = Path(os.getcwd()).resolve().as_posix()
    if not config.get("debug_mode", False):
        dbfpath = Path(os.getcwd()).joinpath("logdirs.csv")
    else:
        dbfpath = Path(os.getcwd()).joinpath("logdirs-debug.csv")
    if dbfpath.is_file():
        mode = "a"
        header = False
    else:
        mode = "w"
        header = True
    df = pd.DataFrame({logdir: cfgdict})
    df.T.to_csv(dbfpath.resolve().as_posix(), mode=mode, header=header)
    os.environ["LOGDIR"] = logdir

def print_styles():
    """Print the configured logging styles (optionally exporting to HTML)."""
    import argparse

    parser = argparse.ArgumentParser()
    from rich.text import Text

    from ezpz.log.console import Console

    parser.add_argument(
        "--html", action="store_true", help="Export as HTML table"
    )
    args = parser.parse_args()
    html: bool = args.html
    from rich.table import Table

    console = Console(record=True, width=120) if html else Console()
    table = Table("Name", "Styling")
    for style_name, style in STYLES.items():
        table.add_row(Text(style_name, style=style), str(style))

    console.print(table)
    if html:
        outfile = "enrich_styles.html"
        print(f"Saving to `{outfile}`")
        with open(outfile, "w") as f:
            f.write(console.export_html(inline_styles=True))

def print_styles_alt(
    html: bool = False,
    txt: bool = False,
):
    """Variant of :func:`print_styles` with HTML and plain-text exports."""
    from pathlib import Path

    from rich.table import Table
    from rich.text import Text

    from ezpz.log.console import get_console
    from ezpz.log.style import DEFAULT_STYLES

    console = get_console(record=html, width=150)
    table = Table("Name", "Styling")
    styles = DEFAULT_STYLES
    styles |= STYLES
    for style_name, style in styles.items():
        table.add_row(Text(style_name, style=style), str(style))
    console.print(table)
    if html:
        outfile = "ezpz_styles.html"
        print(f"Saving to `{outfile}`")
        with open(outfile, "w") as f:
            f.write(console.export_html(inline_styles=True))
    if txt:
        file1 = "ezpz_styles.txt"
        text = console.export_text()
        # with open(file1, "w") as file:
        with Path(file1).open("w") as file:
            file.write(text)

def printarr(*arrs, float_width=6):
    """
    Print a pretty table giving name, shape, dtype, type, and content
    information for input tensors or scalars.

    Call like: printarr(my_arr, some_other_arr, maybe_a_scalar). Accepts a
    variable number of arguments.

    Inputs can be:
        - Numpy tensor arrays
        - Pytorch tensor arrays
        - Jax tensor arrays
        - Python ints / floats
        - None

    It may also work with other array-like types, but they have not been tested

    Use the `float_width` option specify the precision to which floating point
    types are printed.

    Author: Nicholas Sharp (nmwsharp.com)
    Canonical source:
        https://gist.github.com/nmwsharp/54d04af87872a4988809f128e1a1d233
    License: This snippet may be used under an MIT license, and it is also
    released into the public domain. Please retain this docstring as a
    reference.
    """
    import inspect

    frame_ = inspect.currentframe()
    assert frame_ is not None
    frame = frame_.f_back
    # if frame_ is not None:
    #     frame = frame_.f_back
    # else:
    #     frame = inspect.getouterframes()
    default_name = "[temporary]"

    # helpers to gather data about each array

    def name_from_outer_scope(a):
        if a is None:
            return "[None]"
        name = default_name
        if frame_ is not None:
            for k, v in frame_.f_locals.items():
                if v is a:
                    name = k
                    break
        return name

    def dtype_str(a):
        if a is None:
            return "None"
        if isinstance(a, int):
            return "int"
        if isinstance(a, float):
            return "float"
        return str(a.dtype)

    def shape_str(a):
        if a is None:
            return "N/A"
        if isinstance(a, int):
            return "scalar"
        if isinstance(a, float):
            return "scalar"
        return str(list(a.shape))

    def type_str(a):
        # TODO this is is weird... what's the better way?
        return str(type(a))[8:-2]

    def device_str(a):
        if hasattr(a, "device"):
            device_str = str(a.device)
            if len(device_str) < 10:
                # heuristic: jax returns some goofy long string we don't want,
                # ignore it
                return device_str
        return ""

    def format_float(x):
        return f"{x:{float_width}g}"

    def minmaxmean_str(a):
        if a is None:
            return ("N/A", "N/A", "N/A")
        if isinstance(a, int) or isinstance(a, float):
            return (format_float(a), format_float(a), format_float(a))

        # compute min/max/mean. if anything goes wrong, just print 'N/A'
        min_str = "N/A"
        try:
            min_str = format_float(a.min())
        except Exception:
            pass
        max_str = "N/A"
        try:
            max_str = format_float(a.max())
        except Exception:
            pass
        mean_str = "N/A"
        try:
            mean_str = format_float(a.mean())
        except Exception:
            pass

        return (min_str, max_str, mean_str)

    try:
        props = [
            "name",
            "dtype",
            "shape",
            "type",
            "device",
            "min",
            "max",
            "mean",
        ]

        # precompute all of the properties for each input
        str_props = []
        for a in arrs:
            minmaxmean = minmaxmean_str(a)
            str_props.append(
                {
                    "name": name_from_outer_scope(a),
                    "dtype": dtype_str(a),
                    "shape": shape_str(a),
                    "type": type_str(a),
                    "device": device_str(a),
                    "min": minmaxmean[0],
                    "max": minmaxmean[1],
                    "mean": minmaxmean[2],
                }
            )

        # for each property, compute its length
        maxlen = {}
        for p in props:
            maxlen[p] = 0
        for sp in str_props:
            for p in props:
                maxlen[p] = max(maxlen[p], len(sp[p]))

        # if any property got all empty strings,
        # don't bother printing it, remove if from the list
        props = [p for p in props if maxlen[p] > 0]

        # print a header
        header_str = ""
        for p in props:
            prefix = "" if p == "name" else " | "
            fmt_key = ">" if p == "name" else "<"
            header_str += f"{prefix}{p:{fmt_key}{maxlen[p]}}"
        print(header_str)
        print("-" * len(header_str))
        # now print the acual arrays
        for strp in str_props:
            for p in props:
                prefix = "" if p == "name" else " | "
                fmt_key = ">" if p == "name" else "<"
                print(f"{prefix}{strp[p]:{fmt_key}{maxlen[p]}}", end="")
            print("")

    finally:
        del frame