import numpy as np
import struct
pack32be = struct.Struct('>I').pack

def get_header():     ##1920*1080 이미지용 바이너리 헤더 생성
    header = bytearray(0)
    header += bytearray([-0x53, 0x70, 0x6c, 0x64])   ##Signature
    header += bytearray([1920 % 256, 1920//256])     ##Width
    header += bytearray([1080 % 256, 1080//256])     ##Height
    header += bytearray(4)                           ##NumBytes 4바이트
    header += bytearray([0xff]*8)                    
    header += bytearray(4)
    header.append(0)                                 ##Compression, 0=Uncompressed, 1=RLE, 2=Enhanced RLE
    header.append(2)
    header.append(1)
    header += bytearray(21)
    return header

header_template = get_header()

def merge(images):      ##최대 24개의 binary 이미지를 하나의 24-bit 컬러 이미지(RGB)로 합치는 함수

    if len(images) > 24:
        raise ValueError(f"merge 함수는 최대 24개의 이미지만 병합할 수 있습니다. (입력된 이미지 수: {len(images)})")

    image32 = np.zeros((1080, 1920), dtype=np.uint32)   ##1080*1920 해상도, 각 픽셀 32bit 정수, 포멧=0x00BBGGRR
    n_img = len(images)                                 ##한 컬러 채널당 8bit=8개 이미지 저장 가능
    batches = [8]*(n_img//8)
    if n_img % 8:
        batches.append(n_img %8)
    for i, batch_size in enumerate(batches):
        image8 = np.zeros((1080,1920), dtype=np.uint8)  ##image8 = 하나의 컬러 채널에 들어갈 8bit 값
        for j in range(batch_size):                     ## 하나의 8bit 이미지를 만드는 for문
            image8 += images[i*8+j]*(1<<j)
        image32 += image8.astype(np.uint32)*(1 << (i*8))
    return image32

def bgr(pixel):         ##unit32 형식의 한 픽셀 값(0x00BBGGRR)을 [B,G,R] 3바이트로 변환하는 함수
    return pack32be(pixel)[1:4]

def enc128(num):        ##0~32767 사이의 정수를 1바이트 또는 2바이트로 인코딩하는 함수 (저장공간 절약)
    return bytearray([(num & 0x7f) | 0x80, num >> 7]) if num >= 128 else bytearray([num])

def run_len(row, idx):  ##row에서 idx부터 시작하는 연속된 1의 길이를 계산하고 128픽셀 단위로 건너뛰는 최적화 함수
    stride = 128
    length = len(row)
    j = idx
    while j < length and row[j]:
        if j % stride == 0 and np.all(row[j:j+stride]):
            j += min(stride, length-j)
        else:
            j += 1
    return j-idx

def encode_row(row, same_prev):             
    ##ERLE 방식으로 압축하는 코드. 4가지 패턴 (이전 줄과 동일한 픽셀, 같은 픽셀이 반복되는 구간, 단일 픽셀, 여러개의 raw한 픽셀)
    same = np.logical_not(np.diff(row))
    same_either = np.logical_or(same_prev[:1919], same)

    j=0
    compressed = bytearray(0)
    while j < 1920:
        if same_prev[j]:                    ##이전 라인의 n 픽셀과 같을 경우
            r = run_len(same_prev, j+1) + 1
            j += r
            compressed += b'\x00\x01' + enc128(r)

        elif j < 1919 and same[j]:          ##같은 픽셀이 반복될 경우
            r = run_len(same, j+1) + 2
            j += r
            compressed += enc128(r) + bgr(row[j])
            j += 1

        elif j > 1917 or same_either[j+1]:   ##압축되지 않는 단일 픽셀
            compressed += b'\x01' + bgr(row[j])
            j += 1
        
        else:                                ##압축되지 않는 나머지 여러 픽셀들
            j_start = j
            pixels = bgr(row[j]) + bgr(row[j+1])
            j += 2
            while j == 1919 or not same_either[j]:
                pixels += bgr(row[j])
                j += 1
            compressed += b'\x00' + enc128(j-j_start) + pixels

    return compressed + b'\x00\x00'

def encode_chunk(images_chunk):         ##입력된 binary 패턴 이미지들을 ERLE압축을 하여 최종 이미지 데이터를 생성하는 함수
    encoded = bytearray(header_template)
    image = merge(images_chunk)

    for i in range(1080):
        same_prev = np.zeros(1920, dtype=bool) if i == 0 else image[i] == image[i-1]
        encoded += encode_row(image[i], same_prev)

    encoded += b'\x00\x01\x00'
    encoded += bytearray((-len(encoded)) % 4)

    struct.pack_into('<I', encoded, 8, len(encoded))

    return encoded, len(encoded)

def encode_all_patterns(images):
    Max_patterns_per_frame = 24
    encoded_frames = []

    for i in range(0, len(images), Max_patterns_per_frame):
        chunk = images[i : i + Max_patterns_per_frame]
        encoded_data, length = encode_chunk(chunk)
        encoded_frames.append({
            "frame_index":i//Max_patterns_per_frame,
            "num_patterns": len(chunk),
            "data": encoded_data,
            "length": length
        })
    
    return encoded_frames