# Copyright 2025 Thinking Machines Lab
# Licensed under the Apache License, Version 2.0
#
# Modifications:
# - Adapted for HPC-AI cloud fine-tuning workflow
# Copyright © 2025 HPC-AI.COM
"""
Small utilities requiring only basic python libraries.
"""

import importlib
import logging
import time
from contextlib import contextmanager
from typing import Any, Sequence, TypeVar, cast

import numpy as np

logger = logging.getLogger(__name__)

T = TypeVar("T")


@contextmanager
def timed(key: str, metrics: dict[str, Any]):
    logger.info(f"Starting {key}")
    tstart = time.time()
    yield
    logger.info(f"{key} took {time.time() - tstart:.2f} seconds")
    metrics[f"time/{key}"] = time.time() - tstart


safezip = cast(type[zip], lambda *args, **kwargs: zip(*args, **kwargs, strict=True))


def dict_mean(list_of_dicts: list[dict[str, float | int]]) -> dict[str, float]:
    key2values = {}
    for d in list_of_dicts:
        for k, v in d.items():
            key2values.setdefault(k, []).append(v)
    return {k: float(np.mean(values)) for k, values in key2values.items()}


def all_same(xs: list[Any]) -> bool:
    return all(x == xs[0] for x in xs)


def lookup_func(path_to_func: str, default_module: str | None = None):
    """
    path.to.module:func_name or func_name (assumes default_module)
    """
    colon_count = path_to_func.count(":")
    if colon_count == 0 and default_module is not None:
        module_name = default_module
        func_name = path_to_func
    elif colon_count == 1:
        module_name, func_name = path_to_func.rsplit(":", 1)
    else:
        raise ValueError(f"Invalid path: {path_to_func}")
    module = importlib.import_module(module_name)
    return getattr(module, func_name)


def split_list(lst: Sequence[T], num_splits: int) -> list[list[T]]:
    """
    Split a sequence into a list of lists, where the sizes are as equal as possible,
    and the long and short lists are as uniformly distributed as possible.

    Args:
        lst: The sequence to split
        num_splits: Number of sublists to create

    Returns:
        A list of sublists with sizes differing by at most 1

    Raises:
        ValueError: If num_splits > len(lst) or num_splits <= 0

    Examples:
        >>> split_list([1, 2, 3, 4, 5], 2)
        [[1, 2, 3], [4, 5]]
        >>> split_list([1, 2, 3, 4, 5], 3)
        [[1, 2], [3, 4], [5]]
    """
    if num_splits <= 0:
        raise ValueError(f"num_splits must be positive, got {num_splits}")
    if num_splits > len(lst):
        raise ValueError(f"Cannot split list of length {len(lst)} into {num_splits} parts")

    edges = np.linspace(0, len(lst), num_splits + 1).astype(int)
    return [list(lst[edges[i] : edges[i + 1]]) for i in range(num_splits)]


def concat_lists(list_of_lists: list[list[Any]]) -> list[Any]:
    return [item for sublist in list_of_lists for item in sublist]


def not_none(x: T | None) -> T:
    assert x is not None, f"{x=} must not be None"
    return x
