src/egg_damage/utils.py

from __future__ import annotations

import json
import logging
import math
import platform
import time
from pathlib import Path
from typing import Any, Iterable

import numpy as np


LOGGER_NAME = "egg_damage"


def get_logger(name: str | None = None) -> logging.Logger:
    logger = logging.getLogger(name or LOGGER_NAME)
    if not logger.handlers:
        handler = logging.StreamHandler()
        handler.setFormatter(
            logging.Formatter("%(asctime)s | %(levelname)s | %(name)s | %(message)s")
        )
        logger.addHandler(handler)
    logger.setLevel(logging.INFO)
    return logger


def save_json(data: Any, path: str | Path) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(make_json_safe(data), f, indent=2)


def load_json(path: str | Path) -> Any:
    with Path(path).open("r", encoding="utf-8") as f:
        return json.load(f)


def make_json_safe(value: Any) -> Any:
    if isinstance(value, dict):
        return {str(k): make_json_safe(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [make_json_safe(v) for v in value]
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, (np.integer,)):
        return int(value)
    if isinstance(value, (np.floating,)):
        value = float(value)
        return value if math.isfinite(value) else None
    if isinstance(value, float):
        return value if math.isfinite(value) else None
    if isinstance(value, np.ndarray):
        return value.tolist()
    return value


def now_stamp() -> str:
    return time.strftime("%Y%m%d-%H%M%S")


def model_file_size_mb(path: str | Path | None) -> float | None:
    if not path:
        return None
    p = Path(path)
    if not p.exists():
        return None
    return p.stat().st_size / (1024 * 1024)


def timer() -> float:
    return time.perf_counter()


def elapsed_ms(start: float, count: int = 1) -> float:
    if count <= 0:
        return 0.0
    return (time.perf_counter() - start) * 1000.0 / count


def format_distribution(rows: Iterable[dict[str, Any]]) -> str:
    lines = ["split | label | count", "--- | --- | ---"]
    for row in rows:
        lines.append(f"{row['split']} | {row['label']} | {row['count']}")
    return "\n".join(lines)


def environment_summary() -> dict[str, Any]:
    summary = {"python": platform.python_version(), "platform": platform.platform()}
    try:
        import torch

        summary["torch"] = torch.__version__
        summary["cuda_available"] = bool(torch.cuda.is_available())
        summary["cuda_device"] = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
    except Exception:
        summary["torch"] = "not importable"
        summary["cuda_available"] = False
        summary["cuda_device"] = None
    return summary