#!/usr/bin/env python3

import matplotlib.pyplot as plt
import struct

from PIL import Image

IMG_WIDTH = 640

SYNC_LENGTH = 105
SYNC_THRESH = 100
PORCH_LENGTH = 10

LINE_LENGTH = SYNC_LENGTH + PORCH_LENGTH + 4 * IMG_WIDTH

MAX_DEV = 600
CENTER = 1750
MIN_FREQ = 1200
MAX_FREQ = 2300

COLOR_LOW = 1500
COLOR_HIGH = MAX_FREQ

def to_freq(sample):
    freq = CENTER + sample * MAX_DEV
    freq = max(MIN_FREQ, freq)
    freq = min(MAX_FREQ, freq)
    return freq

def to_color(sample):
    sample = (sample - 1500) / (COLOR_HIGH - COLOR_LOW)
    sample = int(sample * 255)
    sample = max(sample, 0)
    sample = min(sample, 255)
    return sample

def ycbcr_to_rgb(y, cb, cr):
    #https://stackoverflow.com/questions/4041840/function-to-convert-ycbcr-to-rgb/15333019#15333019
    cr = cr - 128 ;
    cb = cb - 128 ;
    r = y + 45 * cr / 32 ;
    g = y - (11 * cb + 23 * cr) / 32 ;
    b = y + 113 * cb / 64 ;

    return (int(r), int(g), int(b))

last_sync = False
def is_sync(backlog):
    global last_sync
    count = 0
    for sample in backlog:
        if sample < COLOR_LOW:
            count += 1

    res = False
    if count > SYNC_THRESH and not last_sync:
        res = True

    last_sync = count > SYNC_THRESH
    return res


def render_lines(line):
    pixels = [[(0,0,0)] * IMG_WIDTH,  [(0,0,0)] * IMG_WIDTH]

    # Strip porch
    porch = len(line) - SYNC_LENGTH - 4 * IMG_WIDTH
    if porch < 0:
        return pixels
    line = line[PORCH_LENGTH:]

    for i in range(0, IMG_WIDTH):
        y0 = to_color(line[i])
        cr = to_color(line[IMG_WIDTH + i])
        cb = to_color(line[2 * IMG_WIDTH + i])
        y1 = to_color(line[3 * IMG_WIDTH + i])

        pixels[0][i] = ycbcr_to_rgb(y0, cb, cr)
        pixels[1][i] = ycbcr_to_rgb(y1, cb, cr)

    return pixels

def main():
    data_file = open("demod.dat", 'rb')

    samples = []
    syncs = []

    pixels = []

    backlog = [0.0] * SYNC_LENGTH
    line = []

    bytes = data_file.read(4)
    while len(bytes) == 4:
        sample = struct.unpack('f', bytes)[0]
        sample = to_freq(sample)
        samples += [sample]

        backlog = backlog[1:] + [sample]

        line += [sample]

        if is_sync(backlog) or len(line) > LINE_LENGTH:
            if len(line) > SYNC_LENGTH + 4 * IMG_WIDTH:
                syncs += [1000]
                print("%d \t %d" % (len(line), len(line) - LINE_LENGTH))
                pixels += render_lines(line)
            else:
                print("Short line dropped")
            line = []
        else:
            syncs += [0]

        bytes = data_file.read(4)


    img = Image.new('RGB', (IMG_WIDTH, len(pixels)))
    img_pix = img.load()
    for x in range(0, IMG_WIDTH):
        for y in range(0, len(pixels)):
            img_pix[x,y] = pixels[y][x]

    img.save('output.png')


if __name__ == '__main__':
    main()