from functools import wraps
from typing import Callable

import numpy as np


def adjust_dims(func) -> Callable:

    @wraps(func)
    def wrapper(*args, **kwargs):
        """
        ALlow outer product for omega*tau and remove redundant dimension afterwards
        """
        x1, x2, *rest = args

        x1 = np.asanyarray(x1)
        x2 = np.asanyarray(x2)

        if len(x1.shape) > 1 or len(x2.shape) > 1:
            raise ValueError('Only numbers or 1d-arrays are supported')

        result = func(x1[..., None], x2[None, ...], *rest, **kwargs)
        if result.ndim == 1:
            return result[0]
        else:
            return result.squeeze()

    return wrapper
    #
    # return func
