#!/usr/bin/env python

import math
import numpy as np
import h5py
import scipy.ndimage
import sys

from matplotlib import pyplot as plt

F_COLOR_SUB_CARIER = 4.43361875e6
F_MASTER_CLK = F_COLOR_SUB_CARIER * 4.0
F_CPU = F_MASTER_CLK / 18.0
F_VIC_II = F_CPU * 8.0
T_HLINE_US = 1e6 / F_VIC_II * 63 * 8

SYNC_THRESHOLD = 0.1

class SyncDetector:

    def __init__(self, sample_rate, threshold):

        self._sample_rate      = sample_rate
        self._threshold        = threshold
        self._sync_count       = 0
        self._short_sync_count = 0
        self._long_sync        = False
        self._short_sync       = False
        self._normal_sync      = False
        self._vertical_sync    = False

    def input(self, luma):

        sync = (luma <= self._threshold)

        if sync and self._normal_sync and self._vertical_sync:
            self._vertical_sync = False

        if sync:
            self._sync_count += 1
        else:
            if self._sync_count > 0:
                sync_duration_us = self._sync_count * 1e6 / self._sample_rate

                if sync_duration_us > 10.0:
                    # Long sync pulse.
                    self._long_sync = True
                    self._normal_sync = False
                    self._short_sync = False
                    self._vertical_sync = True

                elif self._long_sync:
                    self._long_sync = False
                    self._short_sync = True
                    self._short_sync_count = self._sync_count

                elif self._short_sync:
                    if self._sync_count > 1.5 * self._short_sync_count:
                        self._short_sync = False
                        self._normal_sync = True

                elif not self._vertical_sync:
                    if self._sync_count < 1.5 * self._short_sync_count:
                        self._normal_sync = False
                        self._vertical_sync = True

            self._sync_count = 0

        return sync, self._vertical_sync

class LumaDecoder:

    def __init__(self, sample_rate, threshold):

        self._sync_detector = SyncDetector(sample_rate, threshold)
        self._sync = False
        self._vertical_sync = False

        self._num_x = math.ceil(63 * 8 * sample_rate / F_VIC_II)
        self._num_y = 312 - 9 - 1 - 1
        self._frame = None
        self._x = 0
        self._y = 0

    def input(self, luma):

        frame = None

        sync, vertical_sync = self._sync_detector.input(luma)

        new_frame = not vertical_sync and self._vertical_sync
        self._vertical_sync = vertical_sync

        new_h_line = sync and not self._sync
        self._sync = sync

        if vertical_sync:
            return None

        if new_frame:
            frame = self._frame
            self._frame = np.zeros((self._num_y, self._num_x))
            self._y = 0

        if self._frame is None:
            return None

        if new_h_line:
            self._x = 0

        if new_h_line and not new_frame:
            self._y += 1

        if self._x < self._num_x and self._y < self._num_y:
            self._frame[self._y, self._x] = luma
            self._x += 1

        return frame

def main(filename):

    with h5py.File(filename, "r") as hdf5_file:

        sample_rate = hdf5_file.attrs["sample_rate"]
        luma = hdf5_file["luma"][:]

    decoder = LumaDecoder(sample_rate, SYNC_THRESHOLD)

    for sample in luma:
        frame = decoder.input(sample)
        if frame is None:
            continue

        frame = scipy.ndimage.zoom(frame, (1.0, (F_VIC_II / sample_rate)))
        plt.imsave("decode-luma.png", frame, cmap="gray", vmin=0.0, vmax=1.0)

        plt.figure()
        plt.imshow(frame, cmap="gray", vmin=0.0, vmax=1.0)

    plt.show()

if __name__ == "__main__":

    main(*sys.argv[1:])