""" Layer/Module Helpers

Hacked together by / Copyright 2020 Ross Wightman
"""
from itertools import repeat
import collections.abc

import torch

_HAS_COMPILER_IS_EXPORTING = hasattr(torch, 'compiler') and hasattr(torch.compiler, 'is_exporting')  # PyTorch >= 2.7


# From PyTorch internals
def _ntuple(n):
    """Return a function that converts input to an n-tuple.

    Scalar values are repeated n times, while iterables are converted to tuples.
    Strings are treated as scalars to avoid character-level splitting.

    Args:
        n: Target tuple length.

    Returns:
        Function that converts input to n-tuple.
    """
    def parse(x):
        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
            return tuple(x)
        return tuple(repeat(x, n))
    return parse


to_1tuple = _ntuple(1)
to_2tuple = _ntuple(2)
to_3tuple = _ntuple(3)
to_4tuple = _ntuple(4)
to_ntuple = _ntuple


def get_device_dtype(module, device=None, dtype=None):
    """Resolve factory device and dtype kwargs from a parent module.

    Explicit values take precedence. Missing values are inferred from the
    module's first parameter, if it has one.
    """
    parameter = next(module.parameters(), None)
    if parameter is not None:
        if device is None:
            device = parameter.device
        if dtype is None:
            dtype = parameter.dtype
    return {'device': device, 'dtype': dtype}


def make_divisible(v, divisor=8, min_value=None, round_limit=.9):
    """Adjust value to be divisible by a divisor, typically for channel counts.

    Rounds to the nearest multiple of divisor while ensuring the result doesn't
    fall below min_value or decrease by more than (1 - round_limit).

    Args:
        v: Value to adjust.
        divisor: Target divisor.
        min_value: Minimum acceptable value.
        round_limit: Prevent decrease beyond this fraction of original value.

    Returns:
        Adjusted value divisible by divisor.
    """
    min_value = min_value or divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < round_limit * v:
        new_v += divisor
    return new_v


def extend_tuple(x, n):
    """Pad a tuple to length n by repeating the last value.

    If input is shorter than n, extends by repeating the last element.
    If input is longer than n, truncates to n.

    Args:
        x: Input value, tuple, or list.
        n: Target length.

    Returns:
        Tuple of length n.
    """
    if not isinstance(x, (tuple, list)):
        x = (x,)
    else:
        x = tuple(x)
    pad_n = n - len(x)
    if pad_n <= 0:
        return x[:n]
    return x + (x[-1],) * pad_n


@torch.jit.unused
def _compiler_is_exporting() -> bool:
    return _HAS_COMPILER_IS_EXPORTING and torch.compiler.is_exporting()


def is_exporting() -> bool:
    """Return True while the model is being exported with torch.export (incl. the dynamo ONNX exporter).

    Use to avoid Python branches on input shape or data that would specialize the exported graph to the
    example input. Always False when scripting or tracing, and with PyTorch < 2.7 (no is_exporting()).

    NOTE: also True within torch.compile for PyTorch 2.7 - 2.11, so code paths
    selected by this must produce the same result as the non-export path.
    """
    if torch.jit.is_scripting():
        return False
    return _compiler_is_exporting()
