import numpy as np

#pythran export discrepancy_scipy(float64[:,:], bool, str)
def discrepancy_scipy(sample, iterative=False, method="CD"):
    sample = np.asarray(sample)

    n, d = sample.shape

    if iterative:
        n += 1

    if method == "CD":
        # reference [1], page 71 Eq (3.7)
        abs_ = abs(sample - 0.5)
        disc1 = np.sum(np.prod(1 + 0.5 * abs_ - 0.5 * abs_ ** 2, axis=1))

        prod_arr = np.ones((sample.shape[0], sample.shape[1]))
        for i in range(d):
            s0 = sample[:, i]
            prod_arr *= (
                1
                + 0.5 * abs(s0[:, None] - 0.5)
                + 0.5 * abs(s0 - 0.5)
                - 0.5 * abs(s0[:, None] - s0)
            )
        disc2 = prod_arr.sum()

        return (13.0 / 12.0) ** d - 2.0 / n * disc1 + 1.0 / (n ** 2) * disc2
    elif method == "WD":
        # reference [1], page 73 Eq (3.8)
        prod_arr = np.ones((sample.shape[0], sample.shape[1]))
        for i in range(d):
            s0 = sample[:, i]
            x_kikj = abs(s0[:, None] - s0)
            prod_arr *= 3.0 / 2.0 - x_kikj + x_kikj ** 2

        # typo in the book sign missing: - (4.0 / 3.0) ** d
        return -((4.0 / 3.0) ** d) + 1.0 / (n ** 2) * prod_arr.sum()
    elif method == "MD":
        # reference [2], page 290 Eq (18)
        abs_ = abs(sample - 0.5)
        disc1 = np.sum(
            np.prod(5.0 / 3.0 - 0.25 * abs_ - 0.25 * abs_ ** 2, axis=1)
        )

        prod_arr = np.ones((sample.shape[0], sample.shape[1]))
        for i in range(d):
            s0 = sample[:, i]
            prod_arr *= (
                15.0 / 8.0
                - 0.25 * abs(s0[:, None] - 0.5)
                - 0.25 * abs(s0 - 0.5)
                - 3.0 / 4.0 * abs(s0[:, None] - s0)
                + 0.5 * abs(s0[:, None] - s0) ** 2
            )
        disc2 = prod_arr.sum()

        disc = (19.0 / 12.0) ** d
        disc1 = 2.0 / n * disc1
        disc2 = 1.0 / (n ** 2) * disc2

        return disc - disc1 + disc2
    elif method == "L2-star":
        # reference [1], page 69 Eq (3.5)
        disc1 = np.sum(np.prod(1 - sample ** 2, axis=1))

        xik = sample[None, :, :]
        xjk = sample[:, None, :]
        disc2 = np.sum(
            np.sum(np.prod(1 - np.maximum(xik, xjk), axis=2), axis=1)
        )

        return np.sqrt(
            3 ** (-d) - 1 / n * 2 ** (1 - d) * disc1 + 1 / (n ** 2) * disc2
        )

