import numpy as np

L = 64
Dim = 16
BLK = 4

# Q[L, Dim]
# K[L, Dim]
# V[L, Dim]

MIN_M = -10000


def np_softmax(x):
    x -= np.max(x, axis=1, keepdims=True)
    f_x = np.exp(x) / np.sum(np.exp(x), axis=1, keepdims=True)
    return f_x


def normal_softmax(x):
    out = np.array(x)
    for r in range(0, L):
        maxi = 0
        for i in range(0, L):
            maxi = max(maxi, x[r, i])
        e_sum = 0
        for i in range(0, L):
            e_sum += np.exp(x[r, i] - maxi)
        for i in range(0, L):
            out[r, i] = np.exp(x[r, i] - maxi) / e_sum
    return out


def online_softmax_update(m0, d0, m1, d1):
    #                             x    1
    # Init value:       MIN_M 0
    m = max(m0, m1)
    d = d0 * np.exp(m0 - m) + d1 * np.exp(m1 - m)
    return m, d


def online_softmax(x):
    out = np.zeros(x.shape, x.dtype)
    for r in range(0, L):
        m = MIN_M
        d = 0
        for i in range(0, L):
            m, d = online_softmax_update(m, d, x[r, i], 1)
        for i in range(0, L):
            out[r, i] = np.exp(x[r, i] - m) / d
    return out


def online_block_softmax(x):
    assert L % BLK == 0
    out = np.array(x)
    for r in range(0, L):

        m = MIN_M
        d = 0
        for b in range(0, L // BLK):
            # Calculate block
            mm = MIN_M
            dd = 0
            for i in range(0, BLK):
                mm, dd = online_softmax_update(mm, dd, x[r, b * BLK + i], 1)

            # Merge to total
            m, d = online_softmax_update(m, d, mm, dd)

        for i in range(0, L):
            out[r, i] = np.exp(x[r, i] - m) / d
    return out


a = np.random.uniform(-1, 1, [L, L]).astype("float32")

np_res = np_softmax(a)
# print(np_res)

normal_res = normal_softmax(a)
# print(normal_res)
np.testing.assert_allclose(np_res, normal_res, 1e-6, 1e-6)

online_res = online_softmax(a)
np.testing.assert_allclose(np_res, online_res, 1e-6, 1e-6)

online_block_res = online_block_softmax(a)
np.testing.assert_allclose(np_res, online_block_res, 1e-6, 1e-6)


def np_selfattn(q, k, v):
    x = np.matmul(q, np.transpose(k))
    softmax = np_softmax(x)
    o = np.matmul(softmax, v)
    return o


def normal_selfattn(q, k, v):
    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L]
    softmax = np_softmax(x)

    # [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        for i in range(0, Dim):
            for j in range(0, L):
                o[r, i] += softmax[r, j] * v[j, i]

    return o


def flashattn_0(q, k, v):
    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L] -> [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        m = MIN_M
        d = 0
        for i in range(0, L):
            m, d = online_softmax_update(m, d, x[r, i], 1)

        softmax = np.zeros([L], "float32")
        for i in range(0, L):
            softmax[i] = np.exp(x[r, i] - m) / d
        for c in range(0, Dim):
            for i in range(0, L):
                o[r, c] += softmax[i] * v[i, c]

    return o


def flashattn_update(m, d, m0, d0, s0, m1, d1, s1):
    #                      |   |   |   |   |   |
    #                      |   |   |   x   v   1
    # Init value:        MIN_M 0   0
    s = s0 * np.exp(m0 - m) * d0 / d + s1 * np.exp(m1 - m) * d1 / d
    return s


def flashattn_1(q, k, v):
    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L] -> [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        m = []
        d = []
        for i in range(0, L):
            mm, dd = online_softmax_update(
                m[-1] if i > 0 else MIN_M, d[-1] if i > 0 else 0, x[r, i], 1
            )
            m.append(mm)
            d.append(dd)

        for c in range(0, Dim):
            s = 0
            for i in range(0, L):
                s = flashattn_update(
                    m[i],
                    d[i],
                    m[i - 1] if i > 0 else MIN_M,
                    d[i - 1] if i > 0 else 0,
                    s,
                    x[r, i],
                    v[i, c],
                    1,
                )
            o[r, c] = s
    return o


def flashattn_1_block(q, k, v):
    assert L % BLK == 0
    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L] -> [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        m = np.zeros([L // BLK], "float32")
        d = np.zeros([L // BLK], "float32")
        mm = np.zeros([L], "float32")
        dd = np.zeros([L], "float32")
        for b in range(0, L // BLK):
            # Calculate block
            for i in range(0, BLK):
                mm[b * BLK + i], dd[b * BLK + i] = online_softmax_update(
                    mm[b * BLK + i - 1] if i > 0 else MIN_M,
                    dd[b * BLK + i - 1] if i > 0 else 0,
                    x[r, b * BLK + i],
                    1,
                )

            # Merge to total
            m[b], d[b] = online_softmax_update(
                m[b - 1] if b > 0 else MIN_M,
                d[b - 1] if i > 0 else 0,
                mm[(b + 1) * BLK - 1],
                dd[(b + 1) * BLK - 1],
            )

        for c in range(0, Dim):
            s = 0
            for b in range(0, L // BLK):
                # Calculate block
                ss = 0
                for i in range(0, BLK):
                    ss = flashattn_update(
                        mm[b * BLK + i],
                        dd[b * BLK + i],
                        mm[b * BLK + i - 1] if i > 0 else MIN_M,
                        dd[b * BLK + i - 1] if i > 0 else 0,
                        ss,
                        x[r, b * BLK + i],
                        v[b * BLK + i, c],
                        1,
                    )

                # Merge to total
                s = flashattn_update(
                    m[b],
                    d[b],
                    m[b - 1] if b > 0 else MIN_M,
                    d[b - 1] if b > 0 else 0,
                    s,
                    mm[(b + 1) * BLK - 1],
                    dd[(b + 1) * BLK - 1],
                    ss,
                )
            o[r, c] = s
    return o


def flashattn_1_reorder(q, k, v):
    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L] -> [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        s = np.zeros([Dim], "float32")

        m = MIN_M
        d = 0
        for i in range(0, L):
            mm, dd = m, d
            m, d = online_softmax_update(mm, dd, x[r, i], 1)
            for c in range(0, Dim):
                s[c] = flashattn_update(
                    m, d, mm, dd, s[c] if i > 0 else 0, x[r, i], v[i, c], 1
                )

        for c in range(0, Dim):
            o[r, c] = s[c]
    return o


def flashattn_1_reorder_block(q, k, v):
    assert L % BLK == 0

    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L] -> [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        s = np.zeros([Dim], "float32")

        m = MIN_M
        d = 0
        for b in range(0, L // BLK):
            # Calculate block
            mm = MIN_M
            dd = 0
            ss = np.zeros([Dim], "float32")
            for i in range(0, BLK):
                mmm, ddd = mm, dd
                mm, dd = online_softmax_update(mmm, ddd, x[r, b * BLK + i], 1)
                for c in range(0, Dim):
                    ss[c] = flashattn_update(
                        mm,
                        dd,
                        mmm,
                        ddd,
                        ss[c] if i > 0 else 0,
                        x[r, b * BLK + i],
                        v[b * BLK + i, c],
                        1,
                    )

            # Merge to total
            last_m, last_d = m, d
            m, d = online_softmax_update(m, d, mm, dd)
            for c in range(0, Dim):
                s[c] = flashattn_update(
                    m, d, last_m, last_d, s[c] if b > 0 else 0, mm, dd, ss[c]
                )

        for c in range(0, Dim):
            o[r, c] = s[c]
    return o


q = np.random.uniform(-1, 1, [L, Dim]).astype("float32")
k = np.random.uniform(-1, 1, [L, Dim]).astype("float32")
v = np.random.uniform(-1, 1, [L, Dim]).astype("float32")

np_attn_res = np_selfattn(q, k, v)
# print(np_attn_res)

normal_selfattn_res = normal_selfattn(q, k, v)
np.testing.assert_allclose(np_attn_res, normal_selfattn_res, 1e-6, 1e-6)

flashattn_0_res = flashattn_0(q, k, v)
np.testing.assert_allclose(np_attn_res, flashattn_0_res, 1e-6, 1e-6)

flashattn_1_res = flashattn_1(q, k, v)
np.testing.assert_allclose(np_attn_res, flashattn_1_res, 1e-6, 1e-6)

flashattn_1_block_res = flashattn_1_block(q, k, v)
np.testing.assert_allclose(np_attn_res, flashattn_1_block_res, 1e-6, 1e-6)

flashattn_1_reorder_res = flashattn_1_reorder(q, k, v)
np.testing.assert_allclose(np_attn_res, flashattn_1_reorder_res, 1e-6, 1e-6)

flashattn_1_reorder_block_res = flashattn_1_reorder_block(q, k, v)
np.testing.assert_allclose(np_attn_res, flashattn_1_reorder_block_res, 1e-6, 1e-6)

#                      m, d, m0, d0, s0, m1, d1, s1):
def flashattn_2_update(m,    m0,     s0, m1, d1, s1):
    #                        |       |   |   |   |
    #                        |       |   x   v   1
    # Init value:           MIN_M    0
    s = s0 * np.exp(m0 - m) + s1 * np.exp(m1 - m) * d1
    return s


def flashattn_2_block(q, k, v):
    assert L % BLK == 0
    # [L, Dim] * [Dim, L] -> [L, L]
    x = np.zeros([L, L], "float32")
    for r in range(0, L):
        for i in range(0, L):
            for j in range(0, Dim):
                x[r, i] += q[r, j] * k[i, j]

    # [L, L] -> [L, L] * [L, Dim] -> [L, Dim]
    o = np.zeros([L, Dim], "float32")
    for r in range(0, L):
        m = np.zeros([L // BLK], "float32")
        d = np.zeros([L // BLK], "float32")
        mm = np.zeros([L], "float32")
        dd = np.zeros([L], "float32")
        for b in range(0, L // BLK):
            # Calculate block
            for i in range(0, BLK):
                mm[b * BLK + i], dd[b * BLK + i] = online_softmax_update(
                    mm[b * BLK + i - 1] if i > 0 else MIN_M,
                    dd[b * BLK + i - 1] if i > 0 else 0,
                    x[r, b * BLK + i],
                    1,
                )

            # Merge to total
            m[b], d[b] = online_softmax_update(
                m[b - 1] if b > 0 else MIN_M,
                d[b - 1] if i > 0 else 0,
                mm[(b + 1) * BLK - 1],
                dd[(b + 1) * BLK - 1],
            )

        for c in range(0, Dim):
            s = 0
            for b in range(0, L // BLK):
                # Calculate block
                ss = 0
                for i in range(0, BLK):
                    ss = flashattn_2_update(
                        mm[b * BLK + i],
                        # dd[b * BLK + i],
                        mm[b * BLK + i - 1] if i > 0 else MIN_M,
                        # dd[b * BLK + i - 1] if i > 0 else 0,
                        ss,
                        x[r, b * BLK + i],
                        v[b * BLK + i, c],
                        1,
                    )

                # Merge to total
                s = flashattn_2_update(
                    m[b],
                    # d[b],
                    m[b - 1] if b > 0 else MIN_M,
                    # d[b - 1] if b > 0 else 0,
                    s,
                    mm[(b + 1) * BLK - 1],
                    dd[(b + 1) * BLK - 1],
                    ss / dd[(b + 1) * BLK - 1],
                )
            o[r, c] = s / d[L // BLK - 1]
    return o

flashattn_2_block_res = flashattn_2_block(q, k, v)
np.testing.assert_allclose(np_attn_res, flashattn_2_block_res, 1e-6, 1e-6)