import numpy as np
from numba import jit


@jit(nopython=True, cache=True, fastmath=True)
def hadamard(n, dtype=int):
    """
    Construct an Hadamard matrix.

    Constructs an n-by-n Hadamard matrix.
    `n` must be a power of 2.

    # This function is compatible with scipy.linalg.hadamard
    # but about 40% faster
    # by doing array allocation up front
    """

    if n < 1:
        lg2 = 0
    else:
        lg2 = int(np.log2(n))
    if 2**lg2 != n:
        raise ValueError("n must be an positive integer, and n must be a power of 2")

    H = np.empty((n, n), dtype=dtype)
    H[0, 0] = 1
    for i in range(0, lg2):
        sz = 2**i
        H[sz : 2 * sz, 0:sz] = H[0:sz, 0:sz]
        H[0:sz, sz : 2 * sz] = H[0:sz, 0:sz]
        H[sz : 2 * sz, sz : 2 * sz] = -H[0:sz, 0:sz]

    # Sylvester's construction
    # for i in range(0, lg2):
    #    H = np.vstack((np.hstack((H, H)), np.hstack((H, -H))))

    return H


@jit(nopython=True, cache=True)
def matDotVec(a, b):
    """Matrix multiplication done right for MbyN * Nby1, i.e. matrix times vector.

    numpy overestimates the need memory size,
    it converts both arrays to the type which encompass
    resulting product before it does multiplication.
    For example, on beaglebone black computer which has only 512 MB memory
    with
    A=np.ones((16384, 16384), dtype=int8) -> 256 MB
    B=np.ones((16284,1), dtype=float) - 16284*8 = 132 kB
    A*B -> 132 kB, i.e. same as X
    But numpy complains:
    A @ B -> array is too big

    So here we will populate result element by element.
    """

    if a.ndim != 2:
        raise ValueError("1st argument should be 2 dimensional matrix")
    if b.ndim > 2:
        raise ValueError("2nd argument should look like a vector shape (N,) or (N,1)")
    if (b.ndim == 2) and (b.shape[0] != b.size):
        raise ValueError("2nd argument should look like a vector shape (N,) or (N,1)")
    if a.shape[1] != b.shape[0]:
        raise ValueError("Matrix dimension mismatch for matrix multiplication!")

    # t = a[0, 0] * b[0]  # test the product to find encompassing data type
    # ideally dtype should be t.dtype but numba crashes, so we force default type
    res = np.empty(b.shape)
    for i in range(a.shape[0]):
        res[i] = np.sum(a[i, :] * b.T)
    return res


if __name__ == "__main__":
    print("matrix times vector")
    a = hadamard(4**2, dtype=np.int8)
    b = np.ones((4**2), dtype=float)
    c = matDotVec(a, b)
    print(f"c.shape = {c.shape}")
    print("success")

    print("matrix times 1d matrix")
    a = hadamard(4**2, dtype=np.int8)
    b = np.ones((4**2, 1), dtype=float)
    c = matDotVec(a, b)
    print(f"c.shape = {c.shape}")
    print("success")

    print("large matrix times 1d matrix")
    a = hadamard(64**2, dtype=np.int8)
    b = np.ones((64**2, 1), dtype=float)
    c = matDotVec(a, b)
    a = hadamard(128**2, dtype=np.int8)
    b = np.ones((128**2, 1), dtype=float)
    c = matDotVec(a, b)
    print(f"c.shape = {c.shape}")
    print("success")
