#!/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):
    r = int(y + 1.402 * (cr - 128))
    g = int(y - 3.44136 * (cb - 128) - 0.714136 * (cr - 128))
    b = int(y + 1.772 * (cb - 128))
    return (r,g,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] = y0, cb, cr
        pixels[1][i] = 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)


    plt.plot(samples)
    plt.plot(syncs)
    plt.show()

    img = Image.new('YCbCr', (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 = img.convert('RGB')

    img.save('output.png')


if __name__ == '__main__':
    main()