#!/usr/bin/python3

# Quick and dirty Nagravision Syster VBI extractor
# -Philip Heron <phil@sanslogic.co.uk>

import struct
import numpy as np
import argparse

parser = argparse.ArgumentParser(description = 'Nagravision Syster VBI extractor')
parser.add_argument('-i', '--input', help = 'Input file name', required = True)
parser.add_argument('-w', '--width', help = 'Samples per line in VBI data', required = True)
parser.add_argument('-s', '--samplerate', help = 'Sample rate of VBI data', required = True)
args = parser.parse_args()

width = int(int(args.samplerate) / 625 / 25) # 1728

# The width to scale the image too (number of syster bits per line * 5)
dwidth = 284 * 5

# The syster sync pattern
sync = np.unpackbits(np.array((0xAA, 0x0B, 0x18, 0x36), dtype = np.uint8))
pattern = np.zeros(len(sync) * 5, dtype = np.float)

for x in range(0, len(sync)):
	v = 1 if sync[x] else -1
	pattern[1 + x * 5] = v * 0.7
	pattern[2 + x * 5] = v
	pattern[3 + x * 5] = v * 0.7

# Prepare the resampler
def _sinc(x):
	return np.sin(np.pi * x) / (np.pi * x)

def _raised_cosine(x, b, t):
	if x == 0:
		return 1.0
	
	return _sinc(x / t) * (np.cos(np.pi * b * x / t) / (1.0 - (4.0 * b * b * x * x / (t * t))))

def _init_filter(swidth, dwidth, level, beta):
	
	try:
		return np.load('filter_table_%d_%d.npy' % (swidth, dwidth))
	except:
		pass
	
	table = np.zeros((dwidth, swidth), dtype = np.float)
	
	for b in range(0, swidth):
		
		tt = (1.0 / swidth) * (0.5 + b)
		
		for x in range(0, dwidth):
			
			tv = (1.0 / dwidth) * (0.5 + x)
			tr = (tv - tt) * swidth
			table[x, b] += _raised_cosine(tr, beta, 1) * level
	
	np.save('filter_table_%d_%d.npy' % (swidth, dwidth), table)
	
	return table

def _filter(table, samples):
	
	t = np.zeros(len(table), dtype = np.float)
	
	for dx in range(0, len(t)):
		t[dx] = np.dot(samples, table[dx])
	
	return t

table = _init_filter(width, dwidth, 1, 0.7)

# Syster 16-bt CRC
def _crc(data):
	
	crc = 0x0000
	
	for b in bytearray(data):
		
		crc = crc ^ (b << 8)
		
		for _ in range(0, 8):
			
			if crc & 0x8000:
				crc = (crc << 1) ^ 0xC003
			
			else:
				crc <<= 1
	
	return crc & 0xFFFF

# Reverse order of an 8-bit byte
def _rev8(b):
	
	b = (b & 0xF0) >> 4 | (b & 0x0F) << 4
	b = (b & 0xCC) >> 2 | (b & 0x33) << 2
	b = (b & 0xAA) >> 1 | (b & 0x55) << 1
	
	return b

# Calculate hamming distance
def _hamming(b1, b2):
	
	d = 0
	b = b1 ^ b2
	
	while b > 0:
		d += b & 1
		b >>= 1
	
	return d

def _hamming84(c):
	
	codes = (
		0x15, 0x02, 0x49, 0x5E, 0x64, 0x73, 0x38, 0x2F,
		0xD0, 0xC7, 0x8C, 0x9B, 0xA1, 0xB6, 0xFD, 0xEA,
	)
	
	return codes[c & 0x0F]

def _unhamming84(c):
	
	codes = (
		  1,  0,  1,  1,  0,  0,  1,  0,  1,  2,  1,  1, 10,  0,  1,  7,
		  0,  0,  1,  0,  0,  0,  0,  0,  6,  0,  1, 11,  0,  0,  3,  0,
		  1, 12,  1,  1,  4,  0,  1,  7,  6,  2,  1,  7,  4,  7,  7,  7,
		  6,  0,  1,  5,  0,  0, 13,  0,  6,  6,  6,  5,  6,  0,  3,  7,
		  1,  2,  1,  1,  4,  0,  1,  9,  2,  2,  1,  2,  2,  2,  3,  2,
		  8,  0,  1,  5,  0,  0,  3,  0,  2,  2,  3,  2,  3,  0,  3,  3,
		  4,  2,  1,  5,  4,  4,  4,  4,  2,  2, 15,  2,  4,  2,  3,  7,
		  4,  5,  5,  5,  4,  0,  3,  5,  6,  2,  3,  5,  3, 14,  3,  3,
		  1, 12,  1,  1, 10,  0,  1,  9, 10,  2,  1, 11, 10, 10, 10,  7,
		  8,  0,  1, 11,  0,  0, 13,  0,  6, 11, 11, 11, 10,  0,  3, 11,
		 12, 12,  1, 12,  4, 12, 13,  7,  6, 12, 15,  7, 10,  7,  7,  7,
		  6, 12, 13,  5, 13,  0, 13, 13,  6,  6,  6, 11,  6, 14, 13,  7,
		  8,  2,  1,  9,  4,  9,  9,  9,  2,  2, 15,  2, 10,  2,  3,  9,
		  8,  8,  8,  5,  8,  0,  3,  9,  8,  2,  3, 11,  3, 14,  3,  3,
		  4, 12, 15,  5,  4,  4,  4,  9, 15,  2, 15, 15,  4, 14, 15,  7,
		  8,  5,  5,  5,  4, 14, 13,  5,  6, 14, 15,  5, 14, 14,  3, 14,
	)
	
	return codes[c & 0xFF]

def _sequence(id):
	
	seq = [ 0xA8, 0xBF, 0xCE, 0xD9, 0x7A, 0x6D, 0x92, 0x85, 0x40, 0x57 ]
	
	# Shortcut for exact matches
	if id in seq:
		return seq.index(id)
	
	# Calculate the hamming distance from each sequence value
	d = [_hamming(id, x) for x in seq]
	
	# Return the closest
	return d.index(min(d))

# Open the source file
f = open(args.input, 'rb')

i = 0
lwidth = int(args.width)
sequence = False

while True:
	
	# Load a line, crop it to the requested width
	line = np.fromfile(f, dtype = np.uint8, count = lwidth)
	
	if len(line) != lwidth:
		print("EOF")
		break
	
	line = line[:width].astype(np.float)
	if len(line) < width:
		line = np.pad(line, (0, width - len(line)), 'edge')
	
	# Resample the line
	line = _filter(table, line)
	
	# Find the best correlation point
	corr = [np.dot(pattern, line[x:x + len(pattern)]) for x in range(0, len(line) - len(pattern))]
	x = np.argmax(corr)
	
	# Calculate the bit positions
	offsets = np.array([int(round(x + 2 + (b * 5 * 1.000))) for b in range(0, 28 * 8)])
	
	# Extract the symbols and bytes
	if max(offsets) < len(line):
		
		symbols = np.array([line[offsets[x]] for x in range(0, 28 * 8)])
		symbols -= np.min(symbols)
		symbols /= np.max(symbols)
		
		data = bytearray(28)
		
		for b in range(0, len(data) * 8):
			data[b // 8] |= (0 if symbols[b] < 0.5 else 1) << (7 - (b % 8))
		
		has_sync = False
		has_crc = False
		has_fixes = 0
		
		s = _sequence(data[4])
		
		# Optional: Fix the sequence byte before the CRC calculation
		#data[4] = [ 0xA8, 0xBF, 0xCE, 0xD9, 0x7A, 0x6D, 0x92, 0x85, 0x40, 0x57 ][s]
		data[4] = _rev8(_hamming84(_unhamming84(_rev8(data[4]))))
		
		# Calculate the hamming distance of the sync pattern
		if _hamming(0xAA0B1836, struct.unpack('>I', data[:4])[0]) < 4:
			has_sync = True
		
		#if data[:4] == b'\xAA\x0B\x18\x36':
		#	has_sync = True
		
		if _crc(data[4:26]) == (data[26] << 8 | data[27]):
			has_crc = True
		
		if has_sync and not has_crc:
			
			w = [ abs(x - 0.5) for x in symbols[40:] ]
			w = sorted(range(len(w)), key = lambda k: w[k])
			
			for x in w:
				
				tdata = data.copy()
				tdata[w[x] // 8] ^= 1 << (7 - (w[x] % 8))
				
				if _crc(tdata[4:26]) == (tdata[26] << 8 | tdata[27]):
					has_crc = True
					has_fixes = 1
					data = tdata
					break
		
		if has_sync and not has_crc:
			
			w = [ abs(x - 0.5) for x in symbols ]
			w = sorted(range(len(w)), key = lambda k: w[k])
			
			for x in range(0, 256):
				
				tdata = data.copy()
				
				for xb in range(0, 8):
					if (x >> xb) & 1:
						tdata[w[xb] // 8] ^= 1 << (7 - (w[xb] % 8))
				
				if _crc(tdata[4:26]) == (tdata[26] << 8 | tdata[27]):
					has_crc = True
					has_fixes = bin(x).count('1')
					data = tdata
					break

		
		if False and has_sync and not has_crc:
			
			for l in range(2, 4):
				
				xbs = [0] * l
				
				for x in range(0, len(symbols) ** len(xbs)):
					
					if len(xbs) == len(set(xbs)):
						
						tdata = data.copy()
						
						for xb in xbs:
							if xb < 40:
								continue
							tdata[xb // 8] ^= 1 << (7 - (xb % 8))
						
						if _crc(tdata[4:26]) == (tdata[26] << 8 | tdata[27]):
							has_crc = True
							has_fixes = len(xbs)
							data = tdata
							break
					
					xb = 0
					xbs[xb] += 1
					
					while xb < len(xbs) - 1 and xbs[xb] == len(w):
						xbs[xb] = 0
						xb += 1
						xbs[xb] += 1
				
				if has_crc:
					break
		
		if has_sync and has_crc:
			
			# Test for a break in the sequence
			if sequence != False:
				while sequence != s:
					sequence = (sequence + 1) % 10
			
			sequence = (s + 1) % 10
			
			# Correct the header
			data = b'\xAA\x0B\x18\x36' + data[4:]
			
			s = ','.join(['0x%02X' % _rev8(x) for x in data])
			s += ', /* Line %d' % i
			
			#if not has_sync:
			#	s += ' Bad Sync'
			
			#if not has_crc:
			#	s += ' Bad CRC'
			
			if has_fixes > 0:
				s += ' Fixed %d' % has_fixes
			
			s += ' */'
			
			print(s)
	
	i += 1

f.close()