from parse_cart import Pico8Cart
import struct
import zlib

def main():
    ff = load_seeds("input/fortunes_foundation.txt")

    ffdata = delta_4b(ff)
    offset_ffdata = 0

    augment_map("../main.p8", "../seed_constants.lua", ffdata, {
        "ffdata": 0
    })


def augment_map(target, target2, 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, delta_8b, delta_3b, delta_4b, delta_5b, zlib_delta_4b, zlib_delta_8b
    ]:
        print("{} encoding: {} bytes".format(encoding.__name__, len(encoding(seeds))))

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

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_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


if __name__ == "__main__":
    main()