import numpy as np


def sample_stable(alpha, beta=0, sigma=1, mu=0, size=1):
    """
    Generates iid samples with with shape equal to `size` according to a stable
    distribution with:
    - `alpha` : stability index
    - `beta`  : skewness
    - `sigma` : scale
    - `mu`    : shift
    according to the parametrization of "Stable non-Gaussian random processes"
    Samorodnitsky[07].
    Algorithm from "On the Chambers-Mallows-Stuck method for simulating skewed stable
    random variables" (and correction) Weron[95-96].
    """
    assert 0 < alpha <= 2
    assert -1 <= beta <= 1
    assert sigma > 0
    U = np.random.uniform(-np.pi / 2, np.pi / 2, size=size)
    W = -np.log(np.random.uniform(0, 1, size=size))
    pi2 = np.pi / 2
    Z = -beta * np.tan(pi2 * alpha)
    if alpha == 1:
        E = pi2
        X = (
            (pi2 + beta * U) * np.tan(U)
            - beta * np.log((pi2 * W * np.cos(U)) / (pi2 + beta * U))
        ) / E
        Y = sigma * X + mu + 2 / np.pi * beta * sigma * np.log(sigma)
    else:
        E = 1 / alpha * np.arctan(-Z)
        X = (
            (1 + Z ** 2) ** (1 / (2 * alpha))
            * np.sin(alpha * (U + E))
            / (np.cos(U) ** (1 / alpha))
            * (np.cos(U - alpha * (U + E)) / W) ** ((1 - alpha) / alpha)
        )
        Y = sigma * X + mu
    return Y


def norm(x):
    """
    Computes l_2 norm of an array `x` of shape [n].
    """
    return np.sqrt(np.sum(x ** 2))


def sample_mv_stable_discrete(alpha, weights, points, n):
    r"""
    Generates `n` iid samples from the multivariate stable distribution with discrete
    spectral measure \Gamma where:
    - `alpha`  : stability index
    - `points` : array of shape [m,k]; points[i] needs to be on the unit sphere in R^k;
                 defines the mass of \Gamma
    - `weights`: array of shape [m] of non-negative floats; defines the weights of
                 points of \Gamma
    For more details see "An overview of multivariate stable distributions" Nolan[94],
    section 4.
    """
    assert weights.shape[0] == points.shape[0], str(weights.shape) + str(points.shape)
    assert len(weights.shape) == 1
    assert len(points.shape) == 2
    for point in points:
        assert np.abs(norm(point) - 1) < 1e-8
    m = weights.shape[0]
    k = points.shape[1]
    if alpha == 1:
        X = []
        for _ in range(n):
            xs = (
                weights
                * (
                    sample_stable(alpha=alpha, beta=1, size=[m])
                    + 2 / np.pi * np.log(weights)
                )
            )[..., None] * points
            X.append(np.sum(xs, axis=0))
        X = np.array(X)
    else:
        X = []
        for _ in range(n):
            xs = (
                weights ** (1 / alpha) * sample_stable(alpha=alpha, beta=1, size=[m])
            )[..., None] * points
            X.append(np.sum(xs, axis=0))
        X = np.array(X)
    assert X.shape == (n, k)
    return X


def sample_net(alpha, sigma_w, sigma_b, phi, x, D, L):
    """
    Generates a sample from a feedforward nn with weights and biases where:
    - `alpha`   : stability index of w and b
    - `sigma_w` : scale parameter of (unscaled) w
    - `sigma_b` : scale parameter of b
    - `phi`     : scalar activation
    - `x`       : input array of shape [k,I] (k: number of inputs; I: dimensionality of
                  each input)
    - `D`       : width of each layer
    - `L`       : total number of layers
    """
    for i in range(L):
        if i > 0:
            x = phi(x)
        d_in = x.shape[1]
        w = sample_stable(
            alpha=alpha, sigma=sigma_w / (d_in ** (1 / alpha)), size=[d_in, D]
        )
        b = sample_stable(alpha=alpha, sigma=sigma_b, size=[D])
        x = x @ w + b
    return x


def gamma_bias(alpha, sigma_b, k):
    """
    Computes the weight and point for the bias part of the spectral measure, same
    paramaters as sample_net().
    """
    ones = np.ones([k])
    weight = np.array([sigma_b ** alpha * norm(ones) ** alpha])
    point = ones / norm(ones)
    return weight, point


def gamma_weight(alpha, sigma_w, x):
    """
    Computes the weights and points for the weight part of the spectral measure, `x` is
    an array of shape [m,k].
    """
    weights = sigma_w ** alpha / len(x) * np.array([norm(x_j) ** alpha for x_j in x])
    points = np.array([x_j / max(norm(x_j), 1e-12) for x_j in x])
    return weights, points


def sample_stable_net(alpha, sigma_w, sigma_b, phi, x, M, M_last, L):
    r"""
    Generates `M_last` iid samples from the multivariate stable distribution
    corresponding to the infinitely wide network generated by sample_net(), where
    `alpha`, `sigma_w`, `sigma_b`, `phi`, `L` are as in sample_net() and:
    - `x` : input array of shape [I,k] (k: number of inputs; I: dimensionality of
            each input)
    - `M` : number of MC points used to approximate the spectral measure \Gamma for
            each layer up to L-1
    """
    k = x.shape[1]
    weight_b, point_b = gamma_bias(alpha, sigma_b, k)
    for i in range(L):
        if i > 0:
            x = phi(x)
        weights_w, points_w = gamma_weight(alpha, sigma_w, x)
        weights = np.concatenate([weight_b, weights_w])
        points = np.concatenate([point_b[None, ...], points_w])
        weights = 0.5 * np.concatenate([weights, weights])
        points = np.concatenate([points, -points])
        m = M if (i < L - 1) else M_last
        x = sample_mv_stable_discrete(alpha, weights, points, m)
    return x


# Example usage: 1 sample from a 1 hidden layer stable nn and from its infinitely wide
# limit for 2 one-dimensional inputs:
y = sample_net(
    alpha=2, sigma_w=1, sigma_b=1, phi=np.tanh, x=np.array([[-1], [1]]), D=200, L=2
)[:, 0]
s = sample_stable_net(
    alpha=2,
    sigma_w=1,
    sigma_b=1,
    phi=np.tanh,
    x=np.array([[-1, 1]]),
    M=100,
    M_last=1,
    L=2,
)[0]
