/*
 * Copyright (C) 2013 Bastian Bloessl <bloessl@ccs-labs.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <ieee802-11/ofdm_decode_mac.h>

#include "utils.h"

#include <boost/crc.hpp>
#include <gnuradio/io_signature.h>
#include <itpp/itcomm.h>
#include <iomanip>

using namespace gr::ieee802_11;
using namespace itpp;

class ofdm_decode_mac_impl : public ofdm_decode_mac {

public:
ofdm_decode_mac_impl(bool log, bool debug) : block("ofdm_decode_mac",
			gr::io_signature::make(1, 1, 48 * sizeof(gr_complex)),
			gr::io_signature::make(0, 0, 0)),
			d_log(log),
			d_debug(debug),
			ofdm(BPSK_1_2),
			tx(ofdm, 0) {

	message_port_register_out(pmt::mp("out"));

	// bpsk
	std::complex<double> bpsk_syms[] = {std::complex<double>(-1, 0), std::complex<double>(1, 0)};
	int bpsk_bits[] = {0, 1};
	bpsk.set(cvec(bpsk_syms, 2), ivec(bpsk_bits, 2));

	// qpsk
	std::complex<double> qpsk_syms[] = {
			std::complex<double>(-1, -1), std::complex<double>(-1,  1),
			std::complex<double>( 1, -1), std::complex<double>( 1,  1)};
	int qpsk_bits[] = {0, 1, 2, 3};
	qpsk.set(cvec(qpsk_syms, 4), ivec(qpsk_bits, 4));

	// qam16
	std::complex<double> qam16_syms[] = {
		std::complex<double>(-0.9487, -0.9487), std::complex<double>(-0.9487, -0.3162),
		std::complex<double>(-0.9487, 0.9487), std::complex<double>(-0.9487, 0.3162),
		std::complex<double>(-0.3162, -0.9487), std::complex<double>(-0.3162, -0.3162),
		std::complex<double>(-0.3162, 0.9487), std::complex<double>(-0.3162, 0.3162),
		std::complex<double>(0.9487, -0.9487), std::complex<double>(0.9487, -0.3162),
		std::complex<double>(0.9487, 0.9487), std::complex<double>(0.9487, 0.3162),
		std::complex<double>(0.3162, -0.9487), std::complex<double>(0.3162, -0.3162),
		std::complex<double>(0.3162, 0.9487), std::complex<double>(0.3162, 0.3162)};
	int qam16_bits[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
	qam16.set(cvec(qam16_syms, 16), ivec(qam16_bits, 16));

	// qam64
	std::complex<double> qam64_syms[] = {
		std::complex<double>(-1.0801, -1.0801), std::complex<double>(-1.0801, -0.7715),
		std::complex<double>(-1.0801, -0.1543), std::complex<double>(-1.0801, -0.4629),
		std::complex<double>(-1.0801, 1.0801), std::complex<double>(-1.0801, 0.7715),
		std::complex<double>(-1.0801, 0.1543), std::complex<double>(-1.0801, 0.4629),
		std::complex<double>(-0.7715, -1.0801), std::complex<double>(-0.7715, -0.7715),
		std::complex<double>(-0.7715, -0.1543), std::complex<double>(-0.7715, -0.4629),
		std::complex<double>(-0.7715, 1.0801), std::complex<double>(-0.7715, 0.7715),
		std::complex<double>(-0.7715, 0.1543), std::complex<double>(-0.7715, 0.4629),
		std::complex<double>(-0.1543, -1.0801), std::complex<double>(-0.1543, -0.7715),
		std::complex<double>(-0.1543, -0.1543), std::complex<double>(-0.1543, -0.4629),
		std::complex<double>(-0.1543, 1.0801), std::complex<double>(-0.1543, 0.7715),
		std::complex<double>(-0.1543, 0.1543), std::complex<double>(-0.1543, 0.4629),
		std::complex<double>(-0.4629, -1.0801), std::complex<double>(-0.4629, -0.7715),
		std::complex<double>(-0.4629, -0.1543), std::complex<double>(-0.4629, -0.4629),
		std::complex<double>(-0.4629, 1.0801), std::complex<double>(-0.4629, 0.7715),
		std::complex<double>(-0.4629, 0.1543), std::complex<double>(-0.4629, 0.4629),
		std::complex<double>(1.0801, -1.0801), std::complex<double>(1.0801, -0.7715),
		std::complex<double>(1.0801, -0.1543), std::complex<double>(1.0801, -0.4629),
		std::complex<double>(1.0801, 1.0801), std::complex<double>(1.0801, 0.7715),
		std::complex<double>(1.0801, 0.1543), std::complex<double>(1.0801, 0.4629),
		std::complex<double>(0.7715, -1.0801), std::complex<double>(0.7715, -0.7715),
		std::complex<double>(0.7715, -0.1543), std::complex<double>(0.7715, -0.4629),
		std::complex<double>(0.7715, 1.0801), std::complex<double>(0.7715, 0.7715),
		std::complex<double>(0.7715, 0.1543), std::complex<double>(0.7715, 0.4629),
		std::complex<double>(0.1543, -1.0801), std::complex<double>(0.1543, -0.7715),
		std::complex<double>(0.1543, -0.1543), std::complex<double>(0.1543, -0.4629),
		std::complex<double>(0.1543, 1.0801), std::complex<double>(0.1543, 0.7715),
		std::complex<double>(0.1543, 0.1543), std::complex<double>(0.1543, 0.4629),
		std::complex<double>(0.4629, -1.0801), std::complex<double>(0.4629, -0.7715),
		std::complex<double>(0.4629, -0.1543), std::complex<double>(0.4629, -0.4629),
		std::complex<double>(0.4629, 1.0801), std::complex<double>(0.4629, 0.7715),
		std::complex<double>(0.4629, 0.1543), std::complex<double>(0.4629, 0.4629)};

	int qam64_bits[] = {
		 0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
		10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
		20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
		30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
		40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
		50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
		60, 61, 62, 63};
	qam64.set(cvec(qam64_syms, 64), ivec(qam64_bits, 64));
}

~ofdm_decode_mac_impl(){
}

int general_work (int noutput_items, gr_vector_int& ninput_items,
		gr_vector_const_void_star& input_items,
		gr_vector_void_star& output_items) {

	const gr_complex *in = (const gr_complex*)input_items[0];

	int i = 0;

	std::vector<gr::tag_t> tags;
	const uint64_t nread = this->nitems_read(0);

	dout << "Decode MAC: input " << ninput_items[0] << std::endl;

	while(i < ninput_items[0]) {

		get_tags_in_range(tags, 0, nread + i * 48, nread + (i + 1) * 48 - 1,
			pmt::string_to_symbol("ofdm_start"));

		if(tags.size()) {
			pmt::pmt_t tuple = tags[0].value;
			int len_data = pmt::to_uint64(pmt::car(tuple));
			int encoding = pmt::to_uint64(pmt::cdr(tuple));

			ofdm = ofdm_param((Encoding)encoding);
			tx = tx_param(ofdm, len_data);

			copied = 0;

			// this frame is garbage
			if(tx.n_sym > 100) {
				// FIXME: make this nicer, this way garbage frames are skipped
				copied = tx.n_sym + 1;
			}
			dout << "Decode MAC: frame start -- len " << len_data
				<< "  symbols " << tx.n_sym << "  encoding "
				<< encoding << std::endl;
		}

		if(copied < tx.n_sym) {
			std::memcpy(sym + (copied * 48), in, 48 * sizeof(gr_complex));
			copied++;

			if(copied == tx.n_sym) {
				dout << "received complete frame - decoding" << std::endl;
				decode();
				in += 48;
				i++;
				break;
			}
		}

		in += 48;
		i++;
	}

	consume(0, i);
	return 0;
}

void decode() {
	demodulate();
	deinterleave();
	decode_conv();
	descramble();
	print_output();

	// skip service field
	boost::crc_32_type result;
	result.process_bytes(out_bytes + 2, tx.psdu_size);
	if(result.checksum() != 558161692) {
		dout << "checksum wrong -- dropping" << std::endl;
		return;
	}

	mylog(boost::format("encoding: %1% - length: %2% - symbols: %3%")
			% ofdm.encoding % tx.psdu_size % tx.n_sym);

	// create PDU
	pmt::pmt_t blob = pmt::make_blob(out_bytes + 2, tx.psdu_size - 4);
	pmt::pmt_t enc = pmt::from_uint64(ofdm.encoding);
	pmt::pmt_t dict = pmt::make_dict();
	dict = pmt::dict_add(dict, pmt::mp("encoding"), enc);
	message_port_pub(pmt::mp("out"), pmt::cons(dict, blob));
}

void demodulate() {

	cvec symbols;
	symbols.set_length(tx.n_sym * 48);
	for(int i = 0; i < tx.n_sym * 48; i++) {
		symbols[i] = std::complex<double>(sym[i]);
	}

	switch(ofdm.encoding) {
	case BPSK_1_2:
	case BPSK_3_4:

		bits = to_vec(bpsk.demodulate_bits(symbols));
		break;

	case QPSK_1_2:
	case QPSK_3_4:

		bits = to_vec(qpsk.demodulate_bits(symbols));
		break;

	case QAM16_1_2:
	case QAM16_3_4:
		bits = to_vec(qam16.demodulate_bits(symbols));
		break;
	case QAM64_2_3:
	case QAM64_3_4:
		bits = to_vec(qam64.demodulate_bits(symbols));
		break;
	}

	// I hate the guy who wrote itpp
	bits = bits * (-2) + 1;
}

void deinterleave() {

	int n_cbps = ofdm.n_cbps;
	int first[n_cbps];
	int second[n_cbps];
	int s = std::max(ofdm.n_bpsc / 2, 1);

	for(int j = 0; j < n_cbps; j++) {
		first[j] = s * (j / s) + ((j + int(floor(16.0 * j / n_cbps))) % s);
	}

	for(int i = 0; i < n_cbps; i++) {
		second[i] = 16 * i - (n_cbps - 1) * int(floor(16.0 * i / n_cbps));
	}

	for(int i = 0; i < tx.n_sym; i++) {
		for(int k = 0; k < n_cbps; k++) {
			deinter[i * n_cbps + second[first[k]]] = bits[i * n_cbps + k];
		}
	}
}

void decode_conv() {
	Punctured_Convolutional_Code code;
	ivec generator(2);
	generator(0)=0133;
	generator(1)=0171;
	code.set_generator_polynomials(generator, 7);

	bmat puncture_matrix;
	switch(ofdm.encoding) {
	case BPSK_1_2:
	case QPSK_1_2:
	case QAM16_1_2:
		puncture_matrix = "1 1; 1 1";
		break;
	case BPSK_3_4:
	case QPSK_3_4:
	case QAM16_3_4:
	case QAM64_3_4:
		puncture_matrix = "1 1 0; 1 0 1;";
		break;
	case QAM64_2_3:
		puncture_matrix = "1 1 1 1 1 1; 1 0 1 0 1 0;";
		break;
	}
	code.set_puncture_matrix(puncture_matrix);
	code.set_truncation_length(30);

	dout << "coding rate " << code.get_rate() << std::endl;
	dout << tx.n_encoded_bits << std::endl;

	vec rx_signal(deinter, tx.n_encoded_bits);

	code.reset();
	decoded_bits.set_length(tx.n_encoded_bits);
	code.decode_tail(rx_signal, decoded_bits);

	//dout << "length decoded " << decoded_bits.size() << std::endl;
	//std::cout << decoded_bits << std::endl;
}

void descramble () {
	int index = 0;
	for(int i = 0; i < 7; i++) {
		if(decoded_bits(i)) {
			index |= 1 << (6 - i);
		}
	}
	int state = scrambler_init[index];

	int feedback;

	for(int i = 0; i < decoded_bits.size(); i++) {
		feedback = ((!!(state & 64))) ^ (!!(state & 8));
		out_bits[i] = feedback ^ decoded_bits(i);
		state = ((state << 1) & 0x7e) | feedback;
	}

	for(int i = 0; i < decoded_bits.size(); i++) {
		int bit = i % 8;
		int byte = i / 8;
		if(bit == 0) {
			out_bytes[byte] = 0;
		}

		if(out_bits[i]) {
			out_bytes[byte] |= (1 << bit);
		}
	}
}

void print_output() {

	dout << std::endl;
	for(int i = 0; i < decoded_bits.size() / 8; i++) {
		dout << std::setfill('0') << std::setw(2) << std::hex << ((unsigned int)out_bytes[i] & 0xFF) << std::dec << " ";
		if(i % 16 == 15) {
			dout << std::endl;
		}
	}
	dout << std::endl;
	for(int i = 0; i < decoded_bits.size() / 8; i++) {
		if((out_bytes[i] > 31) && (out_bytes[i] < 127)) {
			dout << ((char) out_bytes[i]);
		} else {
			dout << ".";
		}
	}
	dout << std::endl;
}

private:
	gr_complex sym[1000 * 48 * 100];
	vec bits;
	double deinter[1000 * 48];
	char out_bits[40000];
	char out_bytes[40000];
        bvec decoded_bits;
	bool   d_debug;
	bool   d_log;
	tx_param tx;
	ofdm_param ofdm;
	int copied;
	static int scrambler_init[128];

	Modulator<std::complex<double> > bpsk;
	Modulator<std::complex<double> > qpsk;
	Modulator<std::complex<double> > qam16;
	Modulator<std::complex<double> > qam64;
};

ofdm_decode_mac::sptr
ofdm_decode_mac::make(bool log, bool debug) {
	return gnuradio::get_initial_sptr(new ofdm_decode_mac_impl(log, debug));
}

int ofdm_decode_mac_impl::scrambler_init[128] = {
0, 73, 18, 91, 36, 109, 54, 127, 72, 1, 90, 19, 108, 37, 126, 55, 89, 16, 75, 2, 125, 52, 111, 38, 17, 88, 3, 74, 53, 124, 39, 110, 50, 123, 32, 105, 22, 95, 4, 77, 122, 51, 104, 33, 94, 23, 76, 5, 107, 34, 121, 48, 79, 6, 93, 20, 35, 106, 49, 120, 7, 78, 21, 92, 100, 45, 118, 63, 64, 9, 82, 27, 44, 101, 62, 119, 8, 65, 26, 83, 61, 116, 47, 102, 25, 80, 11, 66, 117, 60, 103, 46, 81, 24, 67, 10, 86, 31, 68, 13, 114, 59, 96, 41, 30, 87, 12, 69, 58, 115, 40, 97, 15, 70, 29, 84, 43, 98, 57, 112, 71, 14, 85, 28, 99, 42, 113, 56};