import torch
from torch.distributions import Kumaraswamy


def exp_in_prev_range(x, factor):
    mini, maxi = x.min(0)[0], x.max(0)[0]
    expx = (factor * x).exp()
    expx_01 = (expx - expx.min(0)[0]) / (expx.max(0)[0] - expx.min(0)[0])
    return expx_01 * (maxi - mini) + mini


@torch.no_grad()
def get_batch(*args, hyperparameters, get_batch, **kwargs):
    """
    This `get_batch` can be used to wrap another `get_batch` and apply a Kumaraswamy transform to the input.
    The x's have to be in [0,1] for this to work!
    """
    returns = get_batch(*args, hyperparameters=hyperparameters, **kwargs)

    input_warping_type = hyperparameters.get("input_warping_type", "kumar")
    # controls what part of the batch ('x', 'y' or 'xy') to apply the warping to
    input_warping_groups = hyperparameters.get("input_warping_groups", "x")
    # whether to norm inputs between 0 and 1 before warping, as warping is only possible in that range.
    input_warping_norm = hyperparameters.get("input_warping_norm", False)
    use_icdf = hyperparameters.get("input_warping_use_icdf", False)

    def norm_to_0_1(x):
        eps = 0.00001
        maxima = torch.max(x, 0)[0]
        minima = torch.min(x, 0)[0]
        normed_x = (x - minima) / (maxima - minima + eps)

        def denorm(normed_x):
            return normed_x * (maxima - minima + eps) + minima

        return normed_x, denorm

    def warp_input(x):
        if input_warping_norm:
            x, denorm = norm_to_0_1(x)

        if input_warping_type == "kumar":
            if "input_warping_c_std" in hyperparameters:
                assert (
                    "input_warping_c0_std" not in hyperparameters
                    and "input_warping_c1_std" not in hyperparameters
                )
                hyperparameters["input_warping_c0_std"] = hyperparameters[
                    "input_warping_c_std"
                ]
                hyperparameters["input_warping_c1_std"] = hyperparameters[
                    "input_warping_c_std"
                ]
            inside = 0
            while not inside:
                c1 = (
                    torch.randn(*x.shape[1:], device=x.device)
                    * hyperparameters.get("input_warping_c1_std", 0.75)
                ).exp()
                c0 = (
                    torch.randn(*x.shape[1:], device=x.device)
                    * hyperparameters.get("input_warping_c0_std", 0.75)
                ).exp()
                if not hyperparameters.get("input_warping_in_range", False):
                    inside = True
                elif (
                    (c1 < 10).all()
                    and (c1 > 0).all()
                    and (c0 < 10).all()
                    and (c0 > 0).all()
                ):
                    inside = True
                else:
                    inside -= 1
                    if inside < -100:
                        print("It seems that the input warping is not working.")
            if c1_v := hyperparameters.get("fix_input_warping_c1", False):
                c1[:] = c1_v
            if c0_v := hyperparameters.get("fix_input_warping_c0", False):
                c0[:] = c0_v
            if hyperparameters.get("verbose", False):
                print(f"c1: {c1}, c0: {c0}")
            k = Kumaraswamy(concentration1=c1, concentration0=c0)
            x_transformed = k.icdf(x) if use_icdf else k.cdf(x)
        elif input_warping_type == "exp":
            transform_likelihood = hyperparameters.get(
                "input_warping_transform_likelihood", 0.2
            )
            to_be_transformed = torch.rand_like(x[0, 0]) < transform_likelihood
            transform_factors = torch.rand_like(x[0, 0]) * hyperparameters.get(
                "input_warping_transform_factor", 1.0
            )
            log_direction = torch.rand_like(x[0, 0]) < 0.5
            exp_x = exp_in_prev_range(x, transform_factors)
            minus_exp_x = 1.0 - exp_in_prev_range(1.0 - x, transform_factors)
            exp_x = torch.where(log_direction, exp_x, minus_exp_x)
            x_transformed = torch.where(to_be_transformed[None, None, :], exp_x, x)
        elif input_warping_type is None or input_warping_type == "none":
            x_transformed = x
        else:
            raise ValueError(f"Unknown input_warping_type: {input_warping_type}")

        if input_warping_norm:
            x_transformed = denorm(x_transformed)
        return x_transformed

    if "x" in input_warping_groups:
        returns.x = warp_input(returns.x)
    if "y" in input_warping_groups:
        returns.y = warp_input(returns.y)
        returns.target_y = warp_input(returns.target_y)

    return returns
