import argparse
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from itertools import islice

parser = argparse.ArgumentParser()
parser.add_argument('--revert', action='store_true')
args = parser.parse_args()


NUM_POINTS = 1000
LIMIT = 100


def part1_by1(x):
    x &= 0x0000ffff
    x = (x ^ (x << 8)) & 0x00ff00ff
    x = (x ^ (x << 4)) & 0x0f0f0f0f
    x = (x ^ (x << 2)) & 0x33333333
    x = (x ^ (x << 1)) & 0x55555555
    return x


def part_even(x):
    x = x & 0x5555555555555555
    x = (x | (x >> 1)) & 0x3333333333333333
    x = (x | (x >> 2)) & 0x0F0F0F0F0F0F0F0F
    x = (x | (x >> 4)) & 0x00FF00FF00FF00FF
    x = (x | (x >> 8)) & 0x0000FFFF0000FFFF
    return x


def encode_morton_2d(x, y):
    code = (part1_by1(y) << 1) + part1_by1(x)
    return code


def decode_morton_2d(code):
    x = part_even(code)
    y = part_even(code >> 1)
    return x, y


def revert_morton(code):
    code = ((code >> 1) & 0x55555555) | ((code & 0x55555555) << 1)
    code = ((code >> 2) & 0x33333333) | ((code & 0x33333333) << 2)
    code = ((code >> 4) & 0x0f0f0f0f) | ((code & 0x0f0f0f0f) << 4)
    code = ((code >> 8) & 0x00ff00ff) | ((code & 0x00ff00ff) << 8)
    code = ((code >> 16) & 0xffff) | ((code & 0xffff) << 16)
    return code


def identity(code):
    return code


a = np.random.randint(1, 100, size=(NUM_POINTS, 2))

func = revert_morton if args.revert else identity
codes = ((p[0], p[1], func(encode_morton_2d(p[0], p[1]))) for p in a)
codes_sorted = sorted(codes, key=lambda c: c[2])

b = np.array(list([c[0], c[1]] for c in islice(codes_sorted, LIMIT)))


def update_line(num, data, line):
    line.set_data(data[:num, 0], data[:num, 1])
    return line,


fig = plt.figure()
plt.plot(a[:, 0], a[:, 1], 'b.')
plt.title('Revert Morton' if args.revert else 'Morton')
l, = plt.plot([], [], 'ro')
anim = animation.FuncAnimation(fig, update_line, LIMIT, fargs=(b, l),
                               interval=50, blit=True, repeat=False)
anim.save('{}morton.mp4'.format('revert' if args.revert else ''))
plt.show()