from parse_cart import Pico8Cart
import struct
import zlib

def main():
    # No need to store any seeds for level 2 or 3: all 5000 of the first 5000 seeds are fine
    # We could use negencode_delta_4b, which is pretty uneventful
    # level_1 = load_seeds("input/level_1.txt")[:5000]
    # analyze(level_1)
    # level_2 = load_seeds("input/level_2.txt")[:5000]
    # analyze(level_2)
    # level_3 = load_seeds("input/level_3.txt")[:5000]
    # analyze(level_3)
    level_4 = load_seeds("input/level_4.txt")[:5000]
    # analyze(level_4)
    level_5 = load_seeds("input/level_5.txt")[:5000]
    # analyze(level_5)
    ff = load_seeds("input/fortunes_foundation.txt")
    # analyze(ff)

    # level_7 = load_seeds("input/level_7.txt")[:5000]
    # analyze(level_7)

    level_4_data = negencode_16b(level_4)
    level_5_data = negencode_delta_4b(level_5)
    ff_data = delta_4b(ff)

    all_data = b""

    offsets = {}
    for block, data in [
        ("l4", level_4_data),
        ("l5", level_5_data),
        ("ff", ff_data)
    ]:
        offsets[f"{block}_start"] = len(all_data)
        all_data += data
        offsets[f"{block}_end"] = len(all_data)


    augment_map("../main.p8", "../seed_constants.lua", all_data, offsets)


def augment_map(target, constants_file, binary, offsets):
    assert isinstance(binary, bytes) and len(binary) < 8192  # length of mapdata
    print(f"Length of basic extra map data: {len(binary)}")
    mapdata = (binary + bytes([0] * 8192))[:8192]

    cart = Pico8Cart.load(target)
    def touch_map(memory):
        memory[0x0:0x1000] = mapdata[0x1000:0x2000]
    def touch_gfx(memory):
        memory[0x1000:0x2000] = mapdata[0x0000:0x1000]

    cart.touch("__map__", touch_map)
    cart.touch("__gfx__", touch_gfx)
    cart.save(target)

    with open(constants_file, "wt") as f:
        f.write("seed_constants={\n")
        for i, (k, v) in enumerate(offsets.items()):
            sep = "," if i < len(offsets) - 1 else ""
            f.write(f" {k}={v+0x1000}{sep}\n")
        f.write("}\n")


def load_seeds(fname):
    seeds = set()
    with open(fname, "rt") as f:
        for line in f:
            seeds.add(int(line))

        return list(sorted(seeds))


def analyze(seeds):
    def peek_at_seeds(seeds):
        print("Seeds modulo various")
        for i in range(0, 30):
            matches = [s for s in seeds if s%30 == i]
            print("- {}: {} (max {})".format(i, len(matches), max(matches)))
        print()

    peek_at_seeds(seeds)
    # seeds=seeds[:8192]  # stick to the range with a realistic distribution

    print("{} seeds".format(len(seeds)))
    for encoding in [
        naive, bitfield, delta_8b, delta_2b, delta_3b, delta_4b, delta_5b, zlib_delta_4b, zlib_delta_8b,
        negencode_delta_4b, negencode_16b,
    ]:
        print("{} encoding: {} bytes".format(encoding.__name__, len(encoding(seeds))))

def naive(seeds):
    return b"".join(struct.pack("<I", s) for s in seeds)

def bitfield(seeds):
    mx = max(seeds)
    out = bytearray((mx+7)//8)
    for s in seeds:
        out[s // 8] |= 1 << (s % 8)
    return bytes(out)

def delta_8b(seeds):
    out = b""
    acc = 0
    i = 0
    while i < len(seeds):
        diff = seeds[i] - acc
        if diff > 255:
            out += struct.pack("<B", 0)
            acc += 255
        else:
            out += struct.pack("<B", diff)
            acc = seeds[i]
            i += 1

    return out

def zlib_delta_8b(seeds):
    return zlib.compress(delta_8b(seeds))


def delta_4b(seeds):
    out_nibbles = []
    acc = 0
    i = 0
    while i < len(seeds):
        diff = seeds[i] - acc
        if diff > 15:
            out_nibbles.append(0)
            acc += 15
        else:
            out_nibbles.append(diff)
            acc = seeds[i]
            i += 1

    while len(out_nibbles) % 2 != 0:
        out_nibbles.append(0)

    out = b""
    for i in range(0, len(out_nibbles), 2):
        out += bytes([(out_nibbles[i] << 4) + out_nibbles[i+1]])
    return out

def zlib_delta_4b(seeds):
    return zlib.compress(delta_4b(seeds))

def delta_2b(seeds):
    return delta_nb(seeds, 2)

def delta_3b(seeds):
    return delta_nb(seeds, 3)

def delta_5b(seeds):
    return delta_nb(seeds, 5)

def delta_nb(seeds, n):
    out_fibbles = []
    acc = 0
    i = 0
    while i < len(seeds):
        diff = seeds[i] - acc
        if diff > (1<<n)-1:
            out_fibbles.append(0)
            acc += (1<<n)-1
        else:
            out_fibbles.append(diff)
            acc = seeds[i]
            i += 1

    while len(out_fibbles) % 8 != 0:
        out_fibbles.append(0)

    out = b""
    for i in range(0, len(out_fibbles), 8):
        chunk = out_fibbles[i:i+8]
        chunk_bits = sum([
            [i&(1<<x)==0 for x in reversed(range(n))]
            for i in chunk
        ], [])
        assert len(chunk_bits)==n*8
        chunk_bytes = [
            sum([v*1<<(7-i) for i, v in enumerate(byte)])
            for byte in [
                chunk_bits[o:o+8]
                for o in range(0, n*8, 8)
            ]
        ]
        out += bytes(chunk_bytes)
    return out

def negencode_delta_4b(seeds):
    mx, new = negencode(seeds)
    return struct.pack("<H", mx) + delta_4b(new)

def negencode_16b(seeds):
    mx, new = negencode(seeds)
    out = b""
    out += struct.pack("<H", mx)
    for i in new:
        out += struct.pack("<H", i) 
    return out

def negencode(seeds):
    mn = 0
    mx = max(seeds) + 1
    everyone = set(range(mn, mx))
    for i in seeds:
        everyone.remove(i)
    return mx, sorted(everyone)


if __name__ == "__main__":
    main()