import sys
from typing import List, Tuple

def flatten(l):
    for i in l:
        if isinstance(i, list):
            yield from flatten(i)
        else:
            yield i

from common import *


def encode_RLE(data: List[bool]) -> bytes:
    out = b""
    " values is a list of tuples where the first element is the count and the second is the value "
    " the count can only go upto 7 bit "
    values: List[Tuple[int,bool]] = []
    for i in range(len(data)):
        if i == 0:
            values.append((1, data[i]))
        elif data[i] != data[i-1]:
            values.append((1, data[i]))
        else:
            last = values[-1][0]
            if last < 127:
                values[-1] = (last+1, data[i])
            else:
                values.append((1, data[i]))
            
    
    for v in values:
        # high bit is set if value is true
        v = v[0] << 1 | int(v[1])
        out += v.to_bytes(1, "little")
    return out

stats = {
    "avg_frame_size": 0,
    "min_frame_size": float("inf"),
    "max_frame_size": 0,
    "total_size": 0,
}

def diff_encode(frames: List[List[List[bool]]]) -> bytes:
    out = b""
    for i, frame in enumerate(frames):
        if i > 0:
            last_frame = frames[i-1]
            frame = diff(last_frame, frame)

        o = encode_RLE(list(flatten(frame)))
        out += o
        sz_b = len(o)
        #print(f"Frame {i} size: {sz_b}")
        stats["min_frame_size"] = min(stats["min_frame_size"], sz_b)
        stats["max_frame_size"] = max(stats["max_frame_size"], sz_b)
    return out

def rle_encode(frames: List[List[List[bool]]]) -> bytes:
    out = b""
    for frame in frames:
        o = encode_RLE(list(flatten(frame)))
        out += o
        stats["min_frame_size"] = min(stats["min_frame_size"], len(o))
        stats["max_frame_size"] = max(stats["max_frame_size"], len(o))
    return out

def get_frame(f):
    pixels = []
    for y in range(res[1]):
        b = f.read(4)
        if len(b) < 4:
            return None
        row = [ int(a) == 1 for a in bin(int.from_bytes(b, "little"))[2:].zfill(32)]
        pixels.append(row)
    return pixels


def encode():
    f = open("badapple_input.bin", "rb")
    frames = []
    i = 0
    while True:
        i += 1
        pixels = get_frame(f)
        if not pixels:
            break
        if i % frameskip == 0:
            frames.append(pixels)
    stats["frame_count"] = len(frames)

    with open("out_diff.bin", "wb") as fo:
        o = diff_encode(frames)
        stats["avg_frame_size"] = len(o) / len(frames)
        stats["total_size"] = len(o)
        fo.write(o)
    print("diff", stats)

    with open("out_rle.bin", "wb") as fo:
        o = rle_encode(frames)
        stats["avg_frame_size"] = len(o) / len(frames)
        stats["total_size"] = len(o)
        fo.write(o)
    print("rle", stats)


if __name__ == "__main__":
    frameskip = int(sys.argv[1])
    encode()