text: #!/usr/bin/env python3
  
  import array
  import cmath
  import collections
  import io
  import itertools
  import math
  import os
  
  import numpy as np
  import wltrace
  
  import commpy.channelcoding.convcode as cc
  
  LONG_PREAMBLE_TXT = \
      """
  0 -0.078 0.000 40 0.098 -0.026 80 0.062 0.062 120 -0.035 -0.151
  1 0.012 -0.098 41 0.053 0.004 81 0.119 0.004 121 -0.122 -0.017
  2 0.092 -0.106 42 0.001 -0.115 82 -0.022 -0.161 122 -0.127 -0.021
  3 -0.092 -0.115 43 -0.137 -0.047 83 0.059 0.015 123 0.075 -0.074
  4 -0.003 -0.054 44 0.024 -0.059 84 0.024 0.059 124 -0.003 0.054
  5 0.075 0.074 45 0.059 -0.015 85 -0.137 0.047 125 -0.092 0.115
  6 -0.127 0.021 46 -0.022 0.161 86 0.001 0.115 126 0.092 0.106
  7 -0.122 0.017 47 0.119 -0.004 87 0.053 -0.004 127 0.012 0.098
  8 -0.035 0.151 48 0.062 -0.062 88 0.098 0.026 128 -0.156 0.000
  9 -0.056 0.022 49 0.037 0.098 89 -0.038 0.106 129 0.012 -0.098
  10 -0.060 -0.081 50 -0.057 0.039 90 -0.115 0.055 130 0.092 -0.106
  11 0.070 -0.014 51 -0.131 0.065 91 0.060 0.088 131 -0.092 -0.115
  12 0.082 -0.092 52 0.082 0.092 92 0.021 -0.028 132 -0.003 -0.054
  13 -0.131 -0.065 53 0.070 0.014 93 0.097 -0.083 133 0.075 0.074
  14 -0.057 -0.039 54 -0.060 0.081 94 0.040 0.111 134 -0.127 0.021
  15 0.037 -0.098 55 -0.056 -0.022 95 -0.005 0.120 135 -0.122 0.017
  16 0.062 0.062 56 -0.035 -0.151 96 0.156 0.000 136 -0.035 0.151
  17 0.119 0.004 57 -0.122 -0.017 97 -0.005 -0.120 137 -0.056 0.022
  18 -0.022 -0.161 58 -0.127 -0.021 98 0.040 -0.111 138 -0.060 -0.081
  19 0.059 0.015 59 0.075 -0.074 99 0.097 0.083 139 0.070 -0.014
  20 0.024 0.059 60 -0.003 0.054 100 0.021 0.028 140 0.082 -0.092
  21 -0.137 0.047 61 -0.092 0.115 101 0.060 -0.088 141 -0.131 -0.065
  22 0.001 0.115 62 0.092 0.106 102 -0.115 -0.055 142 -0.057 -0.039
  23 0.053 -0.004 63 0.012 0.098 103 -0.038 -0.106 143 0.037 -0.098
  24 0.098 0.026 64 -0.156 0.000 104 0.098 -0.026 144 0.062 0.062
  25 -0.038 0.106 65 0.012 -0.098 105 0.053 0.004 145 0.119 0.004
  26 -0.115 0.055 66 0.092 -0.106 106 0.001 -0.115 146 -0.022 -0.161
  27 0.060 0.088 67 -0.092 -0.115 107 -0.137 -0.047 147 0.059 0.015
  28 0.021 -0.028 68 -0.003 -0.054 108 0.024 -0.059 148 0.024 0.059
  29 0.097 -0.083 69 0.075 0.074 109 0.059 -0.015 149 -0.137 0.047
  30 0.040 0.111 70 -0.127 0.021 110 -0.022 0.161 150 0.001 0.115
  31 -0.005 0.120 71 -0.122 0.017 111 0.119 -0.004 151 0.053 -0.004
  32 0.156 0.000 72 -0.035 0.151 112 0.062 -0.062 152 0.098 0.026
  33 -0.005 -0.120 73 -0.056 0.022 113 0.037 0.098 153 -0.038 0.106
  34 0.040 -0.111 74 -0.060 -0.081 114 -0.057 0.039 154 -0.115 0.055
  35 0.097 0.083 75 0.070 -0.014 115 -0.131 0.065 155 0.060 0.088
  36 0.021 0.028 76 0.082 -0.092 116 0.082 0.092 156 0.021 -0.028
  37 0.060 -0.088 77 -0.131 -0.065 117 0.070 0.014 157 0.097 -0.083
  38 -0.115 -0.055 78 -0.057 -0.039 118 -0.060 0.081 158 0.040 0.111
  39 -0.038 -0.106 79 0.037 -0.098 119 -0.056 -0.022 159 -0.005 0.120
  """
  
  SHORT_PREAMBLE_TXT = \
      """
  16 0.046 0.046 56 0.046 0.046 96 0.046 0.046 136 0.046 0.046
  17 -0.132 0.002 57 0.002 -0.132 97 -0.132 0.002 137 0.002 -0.132
  18 -0.013 -0.079 58 -0.079 -0.013 98 -0.013 -0.079 138 -0.079 -0.013
  19 0.143 -0.013 59 -0.013 0.143 99 0.143 -0.013 139 -0.013 0.143
  20 0.092 0.000 60 0.000 0.092 100 0.092 0.000 140 0.000 0.092
  21 0.143 -0.013 61 -0.013 0.143 101 0.143 -0.013 141 -0.013 0.143
  22 -0.013 -0.079 62 -0.079 -0.013 102 -0.013 -0.079 142 -0.079 -0.013
  23 -0.132 0.002 63 0.002 -0.132 103 -0.132 0.002 143 0.002 -0.132
  24 0.046 0.046 64 0.046 0.046 104 0.046 0.046 144 0.046 0.046
  25 0.002 -0.132 65 -0.132 0.002 105 0.002 -0.132 145 -0.132 0.002
  26 -0.079 -0.013 66 -0.013 -0.079 106 -0.079 -0.013 146 -0.013 -0.079
  27 -0.013 0.143 67 0.143 -0.013 107 -0.013 0.143 147 0.143 -0.013
  28 0.000 0.092 68 0.092 0.000 108 0.000 0.092 148 0.092 0.000
  29 -0.013 0.143 69 0.143 -0.013 109 -0.013 0.143 149 0.143 -0.013
  30 -0.079 -0.013 70 -0.013 -0.079 110 -0.079 -0.013 150 -0.013 -0.079
  31 0.002 -0.132
  """
  
  LONG_PREAMBLE = dict()
  for idx, i, q in zip(*(iter(LONG_PREAMBLE_TXT.split()),) * 3):
      LONG_PREAMBLE[int(idx)] = complex(float(i), float(q))
  LONG_PREAMBLE = [LONG_PREAMBLE[i] for i in sorted(LONG_PREAMBLE.keys())]
  
  SHORT_PREAMBLE = dict()
  for idx, i, q in zip(*(iter(SHORT_PREAMBLE_TXT.split()),) * 3):
      idx = int(idx)
      if idx >= 16 and idx <= 31:
          SHORT_PREAMBLE[idx - 16] = complex(float(i), float(q))
  SHORT_PREAMBLE = [SHORT_PREAMBLE[i] for i in sorted(SHORT_PREAMBLE.keys())]
  
  # Table 78 - Rate-dependent parameters
  RATE_PARAMETERS = {
      # rate : (n_bpsc, n_cbps, n_dbps),
      6: (1, 48, 24),
      9: (1, 48, 36),
      12: (2, 96, 48),
      18: (2, 96, 72),
      24: (4, 192, 96),
      36: (4, 192, 144),
      48: (6, 288, 192),
      54: (6, 288, 216),
  }
  
  HT_MCS_PARAMETERS = {
      0: (1, 52, 26),
      1: (2, 104, 52),
      2: (2, 104, 78),
      3: (4, 208, 104),
      4: (4, 208, 156),
      5: (6, 312, 208),
      6: (6, 312, 234),
      7: (6, 312, 260),
  }
  
  PILOT_SUBCARRIES = [-21, -7, 7, 21]
  
  PHASE_SCALE = 256
  
  # 802.11a/g
  SUBCARRIERS = list(range(-26, 0)) + list(range(1, 27))
  FFT_MAPPING = dict((c, c if c > 0 else 64 + c) for c in SUBCARRIERS)
  LTS_REF = dict(
      list(zip(SUBCARRIERS,
               [1, 1, -1, -1, 1, 1, -1, 1, -1, 1, 1, 1, 1, 1, 1, -1, -1, 1, 1,
                -1, 1, -1, 1, 1, 1, 1, 1, -1, -1, 1, 1, -1, 1, -1, 1, -1, -1, -1,
                -1, -1, 1, 1, -1, -1, 1, -1, 1, -1, 1, 1, 1, 1])))
  polarity = [1, 1, 1, 1, -1, -1, -1, 1, -1, -1, -1, -1, 1, 1, -1, 1, -1, -1, 1,
              1, -1, 1, 1, -1, 1, 1, 1, 1, 1, 1, -1, 1, 1, 1, -1, 1, 1, -1, -1, 1,
              1, 1, -1, 1, -1, -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, 1, 1, 1, 1, 1,
              -1, -1, 1, 1, -1, -1, 1, -1, 1, -1, 1, 1, -1, -1, -1, 1, 1, -1, -1,
              -1, -1, 1, -1, -1, 1, -1, 1, 1, 1, 1, -1, 1, -1, 1, -1, 1, -1, -1,
              -1, -1, -1, 1, -1, 1, 1, -1, 1, -1, 1, 1, 1, -1, -1, 1, -1, -1, -1,
              1, 1, 1, -1, -1, -1, -1, -1, -1, -1]
  
  RATE_BITS = {
      '1101': 6,
      '1111': 9,
      '0101': 12,
      '0111': 18,
      '1001': 24,
      '1011': 36,
      '0001': 48,
      '0011': 54,
  }
  
  # 802.11n
  HT_SUBCARRIERS = list(range(-28, 0)) + list(range(1, 29))
  HT_FFT_MAPPING = dict((c, c if c > 0 else 64 + c) for c in HT_SUBCARRIERS)
  
  # append 1 for negative sub carriers, -1 for positive sub carriers
  HT_LTS_REF = dict(
      list(zip(HT_SUBCARRIERS,
               [1, 1, 1, 1, -1, -1, 1, 1, -1, 1, -1, 1, 1, 1, 1, 1, 1, -1, -1,
                1, 1, -1, 1, -1, 1, 1, 1, 1, 1, -1, -1, 1, 1, -1, 1, -1, 1, -1,
                -1, -1, -1, -1, 1, 1, -1, -1, 1, -1, 1, -1, 1, 1, 1, 1, -1,
                -1])))
  
  
  def to_hex(n, n_bits):
      return hex((n + (1 << n_bits)) % (1 << n_bits))
  
  
  class ChannelEstimator(object):
      """
      Frequency offset correction and channel equalization.
      """
  
      def __init__(self, samples):
          # samples[0] is the first sample of STS/L-STS
          self.samples = samples
          self.ht = False
          self.subcarriers = SUBCARRIERS
          self.fft_mapping = FFT_MAPPING
          self.short_gi = False
  
          self.fix_freq_offset()
  
          assert len(self.lts_samples) == 128
          lts1 = self.do_fft(self.lts_samples[:64])
          lts2 = self.do_fft(self.lts_samples[64:128])
  
          # print '[RAW LTS] (%s)' % (self.samples[160+32:160+32+4])
          for i in range(0):
              print('[%d] (%d, %d) -> (%d, %d) phase: %f (diff: %d)' % \
                    (i,
                     int(self.samples[192 + i].real),
                     int(self.samples[192 + i].imag),
                     int(self.lts_samples[i].real),
                     int(self.lts_samples[i].imag),
                     cmath.phase(self.lts_samples[i]),
                     int((cmath.phase(self.lts_samples[i]) - cmath.phase(samples[192 + i])) * PHASE_SCALE),
                     ))
          self.gain = dict((c, (lts1[c] + lts2[c]) / 2 * LTS_REF[c])
                           for c in self.subcarriers)
          self.idx = 0
          self.polarity = itertools.cycle(polarity)
  
      def next_symbol(self):
          if self.short_gi:
              symbol = self.do_fft(self.data_samples[self.idx + 8:self.idx + 72])
              self.idx += 72
          else:
              symbol = self.do_fft(self.data_samples[self.idx + 16:self.idx + 80])
              self.idx += 80
  
          # remove randomness of pilot carriers
          p = next(self.polarity)
          if not self.ht:
              polarity = [p, p, p, -p]
          else:
              polarity = [p * p1 for p1 in self.ht_polarity]
              self.ht_polarity.rotate(-1)
          # print '[POLARITY] %d: %s' % (p, polarity)
  
          for c, p in zip(PILOT_SUBCARRIES, polarity):
              symbol[c] *= p
  
          # Correct Residual Carrier Frequency Offset
          prod_sum = 0
          for c in PILOT_SUBCARRIES:
              prod = symbol[c].conjugate() * self.gain[c]
              prod_sum += prod
          beta = cmath.phase(prod_sum)
          print("[PILOT OFFSET] %f (%d)" % (beta, int(beta * PHASE_SCALE)))
          carriers = []
          for c in self.subcarriers:
              if c in PILOT_SUBCARRIES:
                  continue
              symbol[c] *= cmath.exp(complex(0, beta))
              symbol[c] /= self.gain[c]
              carriers.append(symbol[c])
          return carriers
  
      def switch_ht(self):
          """
          Switch to HT channel estimator.
          """
          self.ht = True
          self.subcarriers = HT_SUBCARRIERS
          self.fft_mapping = HT_FFT_MAPPING
  
          ht_sts = self.data_samples[self.idx:self.idx + 80]
          self.idx += 80
  
          ht_offset = cmath.phase(sum([ht_sts[i].conjugate() * ht_sts[i + 16]
                                       for i in range(len(ht_sts) - 16)])) / 16
          print('[HT OFFSET] %f (%d)' % (ht_offset, int(ht_offset * PHASE_SCALE)))
          ht_offset = 0
          self.data_samples = [c * cmath.exp(complex(0, -n * ht_offset))
                               for n, c in enumerate(self.data_samples[self.idx:])]
          self.idx = 0
  
          ht_lts = self.do_fft(self.data_samples[self.idx + 16:self.idx + 80])
          self.idx += 80
  
          self.gain = dict((c, ht_lts[c] * HT_LTS_REF[c]) for c in self.subcarriers)
          self.ht_polarity = collections.deque([1, 1, 1, -1])
  
      def do_fft(self, samples):
          assert len(samples) == 64, len(samples)
          freq = np.fft.fft(samples)
          return dict((c, freq[self.fft_mapping[c]]) for c in self.subcarriers)
  
      def fix_freq_offset(self):
          sts = self.samples[80:160]
          lts = self.samples[160 + 32:160 + 160]
  
          coarse_offset = cmath.phase(sum([sts[i] * sts[i + 16].conjugate()
                                           for i in range(len(sts) - 16)])) / 16
          print('[COARSE OFFSET] %f (%d)' % (coarse_offset, int(coarse_offset * PHASE_SCALE)))
          # coarse_offset = 0
  
          # coarse correction
          lts = [c * cmath.exp(complex(0, n * coarse_offset))
                 for n, c in enumerate(lts)]
  
          fine_offset = cmath.phase(sum([lts[i] * lts[i + 64].conjugate()
                                         for i in range(len(lts) - 64)])) / 64
          print('[FINE OFFSET] %f (%d)' % (fine_offset, int(fine_offset * PHASE_SCALE)))
          fine_offset = 0
  
          self.lts_samples = [c * cmath.exp(complex(0, n * fine_offset))
                              for n, c in enumerate(lts)]
  
          self.freq_offset = coarse_offset + fine_offset
          print('[FREQ OFFSET] %f (%d)' % (self.freq_offset, int(self.freq_offset * PHASE_SCALE)))
  
          self.data_samples = [c * cmath.exp(complex(0, n * self.freq_offset))
                               for n, c in enumerate(self.samples[320:], start=128)]
  
  
  class Demodulator(object):
      QAM16_MAPPING = {
          0b00: -3,
          0b01: -1,
          0b11: 1,
          0b10: 3,
      }
  
      QAM64_MAPPING = {
          0b000: -7,
          0b001: -5,
          0b011: -3,
          0b010: -1,
          0b110: 1,
          0b111: 3,
          0b101: 5,
          0b100: 7,
      }
  
      def __init__(self, rate=6, mcs=0, ht=False):
          if (not ht and rate in [6, 9]) or (ht and mcs == 0):
              # BPSK
              self.scale = 1
              self.bits_per_sym = 1
              self.cons_points = np.array([complex(-1, 0), complex(1, 0)])
          elif (not ht and rate in [12, 18]) or (ht and mcs in [1, 2]):
              # QPSK
              self.scale = 1
              self.bits_per_sym = 2
              self.cons_points = np.array(
                  [complex(-1, -1), complex(-1, 1),
                   complex(1, -1), complex(1, 1)]
              )
          elif (not ht and rate in [24, 36]) or (ht and mcs in [3, 4]):
              # 16-QAM
              self.scale = 3
              self.bits_per_sym = 4
              self.cons_points = np.array(
                  [complex(self.QAM16_MAPPING[i >> 2],
                           self.QAM16_MAPPING[i & 0b11]) for i in range(16)])
          elif (not ht and rate in [48, 54]) or (ht and mcs in [5, 6, 7]):
              # 64-QAM
              self.scale = 7
              self.bits_per_sym = 6
              self.cons_points = np.array(
                  [complex(self.QAM64_MAPPING[i >> 3],
                           self.QAM64_MAPPING[i & 0b111]) for i in range(64)])
  
      def demodulate(self, carriers):
          bits = []
          for sym in carriers:
              idx = np.argmin(abs(sym * self.scale - self.cons_points))
              bits.extend([int(b) for b in ('{0:0%db}' % (self.bits_per_sym)).format(idx)])
          return bits
  
  
  class Signal(object):
  
      def __init__(self, bits):
          assert len(bits) == 24
          str_bits = ''.join([str(b) for b in bits])
          self.rate_bits = str_bits[:4]
          self.rsvd = str_bits[4]
          self.len_bits = str_bits[5:17]
          self.parity_bits = str_bits[17]
          self.tail_bits = str_bits[18:]
  
          self.mcs = 0
          self.rate = RATE_BITS.get(self.rate_bits, 0)
          self.length = int(self.len_bits[::-1], 2)
          self.parity_ok = sum(bits[:18]) % 2 == 0
          self.ht = False
          self.mcs = 0
  
  
  class HTSignal(object):
  
      def __init__(self, bits):
          assert len(bits) == 48
          str_bits = ''.join([str(b) for b in bits])
          self.mcs_bits = str_bits[:6]
          self.cbw = str_bits[7]
          self.len_bits = str_bits[8:24]
          self.smoothing = str_bits[24]
          self.not_sounding = str_bits[25]
          self.rsvd = str_bits[26]
          self.aggregation = str_bits[27]
          self.stbc = str_bits[28:30]
          self.fec = str_bits[30]
          self.short_gi = str_bits[31]
          self.num_ext_stream = str_bits[32:34]
          self.crc = str_bits[34:42]
          self.tail_bits = str_bits[42:48]
  
          self.mcs = int(self.mcs_bits[::-1], 2)
          try:
              self.rate = wltrace.dot11.mcs_to_rate(self.mcs)
          except:
              self.rate = 0
          self.length = int(self.len_bits[::-1], 2)
          self.expected_crc = ''.join(['%d' % c for c in self.calc_crc(bits[:34])])
          self.crc_ok = self.expected_crc == self.crc
          self.ht = True
  
      def calc_crc(self, bits):
          c = [1] * 8
  
          for b in bits:
              next_c = [0] * 8
              next_c[0] = b ^ c[7]
              next_c[1] = b ^ c[7] ^ c[0]
              next_c[2] = b ^ c[7] ^ c[1]
              next_c[3] = c[2]
              next_c[4] = c[3]
              next_c[5] = c[4]
              next_c[6] = c[5]
              next_c[7] = c[6]
              c = next_c
  
          return [1 - b for b in c[::-1]]
  
  
  class Decoder(object):
  
      def __init__(self, path, power_thres=200, skip=0, window=80):
          if path is not None:
              self.fh = open(path, 'rb')
              size = os.path.getsize(path)
              if skip * 4 < size:
                  self.fh.seek(skip * 4)
              else:
                  print("[WARN] try to seek beyond end of file %d/%d" % (skip * 4, size))
              self.power_thres = 200
              self.window = window
              self.count = skip
  
      def decode_next(self, *args, **kwargs):
          trigger = False
          samples = []
          glbl_index = 0
  
          while True:
              chunk = array.array('h', self.fh.read(self.window * 4))
              chunk = [complex(i, q) for i, q in zip(chunk[::2], chunk[1::2])]
              if not trigger and any([abs(c) > self.power_thres for c in chunk]):
                  trigger = True
                  samples = []
                  print("Power trigger at %d" % (self.count))
                  glbl_index = self.count
  
              self.count += self.window
  
              if trigger:
                  samples.extend(chunk)
  
              if trigger and all([abs(c) < self.power_thres for c in chunk]):
                  start = self.find_pkt(samples)
                  if start is None:
                      trigger = False
                  else:
                      print("Decoding packet starting from sample %d" % \
                            (glbl_index + start))
                      return (glbl_index,) + \
                             self.decode(samples[start:], *args, **kwargs)
  
      def find_pkt(self, samples):
          """
          Returns the index of the first sample of STS
          None if no packet was detected.
          """
          lts = LONG_PREAMBLE[-64:]
          peaks = np.array([abs(c) for c in np.correlate(
              samples, lts, mode='valid')]).argsort()[-2:]
          if (max(peaks) - min(peaks) == 64) and (min(peaks) - 32 - 160) > 0:
              return min(peaks) - 32 - 160
          else:
              return None
  
      def demodulate(self, carriers, signal=None):
          if signal is None:
              demod = Demodulator(rate=6)
          else:
              demod = Demodulator(signal.rate, signal.mcs, signal.ht)
          return demod.demodulate(carriers)
  
      def deinterleave(self, in_bits, rate=6, mcs=0, ht=False, verbose=False):
          if not ht:
              n_bpsc, n_cbps, n_dbps = RATE_PARAMETERS[rate]
              n_col = 16
              n_row = 3 * n_bpsc
          else:
              n_bpsc, n_cbps, n_dbps = HT_MCS_PARAMETERS[mcs]
              n_col = 13
              n_row = 4 * n_bpsc
          s = max(n_bpsc / 2, 1)
  
          first_perm = dict()
          for j in range(0, n_cbps):
              first_perm[j] = (s * (j // s)) + ((j + n_col * j // n_cbps) % s)
  
          second_perm = dict()
          for i in range(0, n_cbps):
              second_perm[i] = n_col * i - (n_cbps - 1) * (i // n_row)
  
          if verbose:
              print("Bit rate: %f Mb/s" % (rate))
              print("Bits per sub carrier: %d" % (n_bpsc))
              print("Coded bits per symbol: %d" % (n_cbps))
              print("Data bits per symbol %d" % (n_dbps))
              print("S = %d" % (s))
              print("====== Overall permutation =========")
              for j in range(0, n_cbps):
                  print('%d -> %d -> %d' % (j, first_perm[j], second_perm[first_perm[j]]))
  
          if in_bits is None:
              idx_map = []
              for k in range(n_cbps):
                  idx_map.append((second_perm[first_perm[k]], k))
              return sorted(idx_map)
  
          if verbose:
              print('%d bits, %f samples' % (len(in_bits), float(len(in_bits)) / n_cbps))
  
          out_bits = [0] * len(in_bits)
          for n in range(len(in_bits) // n_cbps):
              for j in range(n_cbps):
                  base = n * n_cbps
                  out_bits[base + second_perm[first_perm[j]]] = in_bits[base + j]
          return out_bits
  
      def viterbi_decode(self, bits, signal=None):
          if signal is None:
              ht = False
              rate = 6
          else:
              ht, rate, mcs = signal.ht, signal.rate, signal.mcs
  
          if (not ht and rate in [9, 18, 36, 54]) or (ht and mcs in [2, 4, 6]):
              print('[PUNCTURE] 3/4')
              # 3/4
              new_bits = []
              for i in range(0, len(bits), 12):
                  new_bits.extend(bits[i:i + 3])
                  new_bits.extend([2, 2])
                  new_bits.extend(bits[i + 3:i + 7])
                  new_bits.extend([2, 2])
                  new_bits.extend(bits[i + 7:i + 11])
                  new_bits.extend([2, 2])
                  new_bits.extend(bits[i + 11:i + 12])
              bits = new_bits
          elif (not ht and rate in [48]) or (ht and mcs in [5]):
              print('[PUNCTURE] 2/3')
              # 2/3
              new_bits = []
              for i in range(0, len(bits), 9):
                  new_bits.extend(bits[i:i + 3])
                  new_bits.append(2)
                  new_bits.extend(bits[i + 3:i + 6])
                  new_bits.append(2)
                  new_bits.extend(bits[i + 6:i + 9])
                  new_bits.append(2)
              bits = new_bits
          elif ht and mcs in [7]:
              print('[PUNCTURE] 5/6')
              # 5/6
              new_bits = []
              for i in range(0, len(bits), 6):
                  new_bits.extend(bits[i:i + 3])
                  new_bits.extend([2, 2])
                  new_bits.extend(bits[i + 3:i + 5])
                  new_bits.extend([2, 2])
                  new_bits.append(bits[i + 5])
              bits = new_bits
          else:
              print('[NO PUNCTURE]')
  
          extended_bits = np.array([0] * 2 + bits + [0] * 12)
          trellis = cc.Trellis(np.array([7]), np.array([[0o133, 0o171]]))
          return list(cc.viterbi_decode(extended_bits, trellis, tb_depth=35))[:-7]
  
      def descramble(self, bits):
          X = [0] * 7
          X[0] = bits[2] ^ bits[6]
          X[1] = bits[1] ^ bits[5]
          X[2] = bits[0] ^ bits[4]
          X[3] = X[0] ^ bits[3]
          X[4] = X[1] ^ bits[2]
          X[5] = X[2] ^ bits[1]
          X[6] = X[3] ^ bits[0]
  
          out_bits = []
          for i, b in enumerate(bits):
              feedback = X[6] ^ X[3]
              out_bits.append(feedback ^ b)
              X = [feedback] + X[:-1]
  
          return out_bits
  
      def check_signal(self, signal):
          if signal.rate_bits not in RATE_BITS:
              print('[SIGNAL] invalid rate: %s' % (signal.rate_bits))
              return False
  
          if signal.rsvd != '0':
              print('[SIGNAL] wrong rsvd')
              return False
  
          if not signal.parity_ok:
              print('[SIGNAL] wrong parity')
              return False
  
          if not all([b == '0' for b in signal.tail_bits]):
              print('[SIGNAL] wrong tail: %s' % (signal.tail_bits))
              return False
  
          return True
  
      def check_ht_signal(self, signal):
          if signal.mcs > 7:
              print('[HT SIGNAL] mcs not supported: %d' % (signal.mcs))
              return False
  
          if signal.cbw != '0':
              print('[HT SIGNAL] CBW not supported')
              return False
  
          if signal.rsvd != '1':
              print('[HT SIGNAL] wrong rsvd')
              return False
  
          if signal.stbc != '00':
              print('[HT SIGNAL] stbc not supported: %s' % (signal.stbc))
              return False
  
          if signal.fec != '0':
              print('[HT SIGNAL] FEC not supported')
              return False
  
          if signal.num_ext_stream != '00':
              print('[HT SIGNAL] EXT spatial stream not supported: %s' % (signal.num_ext_stream))
              return False
  
          if not all([b == '0' for b in signal.tail_bits]):
              print('[HT SIGNAL] wrong tail: %s' % (signal.tail_bits))
              return False
  
          return True
  
      def decode(self, samples, *args, **kwargs):
          eq = ChannelEstimator(samples)
  
          carriers = eq.next_symbol()
          assert len(carriers) == 48
          carrier_bits = self.demodulate(carriers)
          raw_bits = self.deinterleave(carrier_bits)
          bits = self.viterbi_decode(raw_bits)
          signal = Signal(bits)
          print("[SIGNAL] %s" % (signal.__dict__))
  
          if not self.check_signal(signal):
              return
  
          cons = []
  
          if signal.rate == 6:
              # decode next two OFDM SYMs to detect potential 802.11n packets
              demod_out = []
              for _ in range(2):
                  carriers = eq.next_symbol()
                  assert len(carriers) == 48
                  cons.extend(carriers)
  
                  # HT-SIG was rotated by 90 degrees, rotate it back
                  carriers = [c * complex(0, -1) for c in carriers]
  
                  carrier_bits = self.demodulate(carriers)
                  demod_out.extend(carrier_bits)
  
              raw_bits = self.deinterleave(demod_out)
              signal_bits = self.viterbi_decode(raw_bits)
  
              ht_signal = HTSignal(signal_bits)
              if ht_signal.crc_ok:
                  print("[HT SIGNAL] %s" % (ht_signal.__dict__))
                  if not self.check_ht_signal(ht_signal):
                      return
                  signal = ht_signal
                  eq.switch_ht()
                  cons = []
                  eq.short_gi = ht_signal.short_gi == '1'
  
          if signal.ht:
              num_symbol = int(math.ceil(float((16 + signal.length * 8 + 6)) /
                                         HT_MCS_PARAMETERS[signal.mcs][2]))
          else:
              num_symbol = int(math.ceil(float((16 + signal.length * 8.0 + 6)) /
                                         RATE_PARAMETERS[signal.rate][2]))
  
          print("%d DATA OFDM symbols to decode" % (num_symbol))
  
          if signal.rate == 6 and not signal.ht:
              num_symbol -= 2
  
          for _ in range(num_symbol):
              carriers = eq.next_symbol()
              if signal.ht:
                  assert len(carriers) == 52, len(carriers)
              else:
                  assert len(carriers) == 48, len(carriers)
              cons.extend(carriers)
  
          demod_out = self.demodulate(cons, signal)
          deinter_out = self.deinterleave(demod_out, signal.rate, signal.mcs, signal.ht)
          conv_out = self.viterbi_decode(deinter_out, signal)
          descramble_out = self.descramble(conv_out)
  
          if not all([b == 0 for b in descramble_out[:16]]):
              print('[SERVICE] not all bits are 0')
  
          # skip the SERVICE bits
          data_bits = descramble_out[16:]
          num_bytes = min(len(data_bits) / 8, signal.length)
  
          data_bytes = [self.array_to_int(data_bits[i * 8:(i + 1) * 8])
                        for i in range(num_bytes)]
          assert len(data_bytes) == num_bytes
  
          for i in range(0, num_bytes, 16):
              print('[%3d] %s' % \
                    (i, ' '.join([format(data_bytes[j], '02x')
                                  for j in range(i, min(i + 16, num_bytes))])))
  
          fh = io.StringIO(''.join([chr(b) for b in data_bytes]))
          # TODO: pkt = dot11.Dot11Packet(fh)
          pkt = None
  
          return signal, cons, demod_out, deinter_out, conv_out, descramble_out, data_bytes, pkt
  
      def array_to_int(self, arr):
          assert len(arr) == 8
          return int(''.join(['%d' % (b) for b in arr[::-1]]), 2)
  