import sys, os, time
import numpy as num
import pylab as py

# simulation parameters

N           = 1024   # number of subcarriers
L           = 32    # length of cyclic prefix
N_OFDM      = 10000   # number of OFDM symbols
N_sub_ch    = 100   # plot constellation of this subchannel

ph4_rot = num.array([1, 1j, -1, -1j])

qam4_1q = num.array([1+1j])

qam16_1q = num.array([1+1j, 3+1j, 1+3j, 3+3j])

qam32_1q = num.array([1+1j, 3+1j, 5+1j, 1+3j, 3+3j, 5+3j, 1+5j, 3+5j])

qam64_1q = num.array([1+1j, 3+1j, 5+1j, 7+1j,\
                      1+3j, 3+3j, 5+3j, 7+3j,\
                      1+5j, 3+5j, 5+5j, 7+5j,\
                      1+7j, 3+7j, 5+7j, 7+7j])

qam128_1q = num.array([1+1j,  3+1j,  5+1j,  7+1j, 9+1j, 11+1j,\
                       1+3j,  3+3j,  5+3j,  7+3j, 9+3j, 11+3j,\
                       1+5j,  3+5j,  5+5j,  7+5j, 9+5j, 11+5j,
                       1+7j,  3+7j,  5+7j,  7+7j, 9+7j, 11+7j,\
                       1+9j,  3+9j,  5+9j,  7+9j,\
                       1+11j, 3+11j, 5+11j, 7+11j ])

qam256_1q = num.array([1+1j,  3+1j,  5+1j,  7+1j,  9+1j,  11+1j,  13+1j,  15+1j,\
                       1+3j,  3+3j,  5+3j,  7+3j,  9+3j,  11+3j,  13+3j,  15+3j,\
                       1+5j,  3+5j,  5+5j,  7+5j,  9+5j,  11+5j,  13+5j,  15+5j,\
                       1+7j,  3+7j,  5+7j,  7+7j,  9+7j,  11+7j,  13+7j,  15+7j,\
                       1+9j,  3+9j,  5+9j,  7+9j,  9+9j,  11+9j,  13+9j,  15+9j,\
                       1+11j, 3+11j, 5+11j, 7+11j, 9+11j, 11+11j, 13+11j, 15+11j,\
                       1+13j, 3+13j, 5+13j, 7+13j, 9+13j, 11+13j, 13+13j, 15+13j,\
                       1+15j, 3+15j, 5+15j, 7+15j, 9+15j, 11+15j, 13+15j, 15+15j])


def Source4QAM(block_len):
    """
    block_len  : nr of 4-QAM symbols in block
                
    returns    
    
    qam4_array : complex numpy array with block_len 4 QAM symbols
    """
    # random index vector for quadrant selection
    iquad = num.random.randint(0, 4, block_len)
    
    # random index vector for selection of symbol(s) from 1st quadrant
    isymbol = num.zeros(block_len, 'L') # with 4-QAM only 1 symbol per quadrant
    
    qam4_array = ph4_rot[iquad] * qam4_1q[isymbol]
    return qam4_array

def Source16QAM(block_len):
    """
    block_len  : nr of 16-QAM symbols in block
                
    returns    
    
    qam16_array : complex numpy array with block_len 16 QAM symbols
    """
    # random index vector for quadrant selection
    iquad = num.random.randint(0, 4, block_len)
    
    # random index vector for selection of symbol(s) from 1st quadrant
    N_symbols = len(qam16_1q) # nr of symbols in quadrant
    isymbol   = num.random.randint(0, N_symbols, block_len)
    
    qam16_array = ph4_rot[iquad] * qam16_1q[isymbol]
    return qam16_array

def Source32QAM(block_len):
    """
    block_len  : nr of 32-QAM symbols in block
                
    returns    
    
    qam32_array : complex numpy array with block_len 32 QAM symbols
    """
    # random index vector for quadrant selection
    iquad = num.random.randint(0, 4, block_len)
    
    # random index vector for selection of symbol(s) from 1st quadrant
    N_symbols = len(qam32_1q) # nr of symbols in quadrant
    isymbol   = num.random.randint(0, N_symbols, block_len)
    
    qam32_array = ph4_rot[iquad] * qam32_1q[isymbol]
    return qam32_array

def Source64QAM(block_len):
    """
    block_len  : nr of 64-QAM symbols in block
                
    returns    
    
    qam64_array : complex numpy array with block_len 64 QAM symbols
    """
    # random index vector for quadrant selection
    iquad = num.random.randint(0, 4, block_len)
    
    # random index vector for selection of symbol(s) from 1st quadrant
    N_symbols = len(qam64_1q)  # nr of symbols in quadrant
    isymbol   = num.random.randint(0, N_symbols, block_len)
    
    qam64_array = ph4_rot[iquad] * qam64_1q[isymbol]
    return qam64_array

def Source128QAM(block_len):
    """
    block_len  : nr of 128-QAM symbols in block
                
    returns    
    
    qam128_array : complex numpy array with block_len 128 QAM symbols
    """
    # random index vector for quadrant selection
    iquad = num.random.randint(0, 4, block_len)
    
    # random index vector for selection of symbol(s) from 1st quadrant
    N_symbols = len(qam128_1q) # nr of symbols in quadrant
    isymbol   = num.random.randint(0, N_symbols, block_len)
    
    qam128_array = ph4_rot[iquad] * qam128_1q[isymbol]
    return qam128_array

def Source256QAM(block_len):
    """
    block_len  : nr of 256-QAM symbols in block
                
    returns    
    
    qam256_array : complex numpy array with block_len 256 QAM symbols
    """
    # random index vector for quadrant selection
    iquad = num.random.randint(0, 4, block_len)
    
    # random index vector for selection of symbol(s) from 1st quadrant
    N_symbols = len(qam256_1q) # nr of symbols in quadrant
    isymbol   = num.random.randint(0, N_symbols, block_len)
    
    qam256_array = ph4_rot[iquad] * qam256_1q[isymbol]
    return qam256_array

#--------------------------------------------------------------------------------------------------------


# normalised frequency offset (df/fs)
f_offset_n = 1.0e-4


#g_channel_impulse = num.array([1, 0.1, -.1, 0.1, .3, -0.04])
g_channel_impulse = num.array([1, 0.0, 0.0, 0.0, .0, 0.0])

G_channel = num.fft.fft(g_channel_impulse, N)

M = len(g_channel_impulse)

overlap_old = num.zeros(M-1)

#error_total = []

sub_channel = num.zeros(N_OFDM, 'D')

t_start = time.clock()
for n_symbol in range(N_OFDM):
    
    # block of 4-QAM symbols
    qam_vec1 = Source4QAM(N/4)
    qam_vec2 = Source4QAM(N/4)
    qam_vec3 = Source4QAM(N/4)
    qam_vec4 = Source4QAM(N/4)
    
    qam_vec = num.concatenate( (qam_vec1, qam_vec2, qam_vec3, qam_vec4) )
    
    # useful part of OFDM symbol (N subcarriers)
    y_vec = N*num.fft.ifft(qam_vec)  # factor N compensates for the internal scaling of function ifft()
    
    # extract cyclic prefix
    cyclic_prefix = y_vec[N-L:]
    
    # prepend cyclic prefix to form the full OFDM symbol
    ofdm_vec = num.concatenate( (cyclic_prefix, y_vec) )
    
    # the channel ...
    r_vec = num.convolve(ofdm_vec, g_channel_impulse)    
    
    # split into vectors r1_vec, r2_vec 
    r1_vec        = r_vec[0:N+L]
    r2_vec        = r_vec[N+L:]
    
    # correct for channel output of last OFDM-symbol
    r1_vec[0:M-1] = r1_vec[0:M-1] + overlap_old
    
    # save overlap part for next ..
    overlap_old = r2_vec
    
    # frequency offset of receiver
    fo_vec = num.exp(2.j *num.pi * f_offset_n * num.arange(N+L)) 
    
    # apply frequency offset
    r1_vec_fo = r1_vec * fo_vec
    
    # discard cyclic prefix part
    ofdm_rec = r1_vec_fo[L:]   # should have N samples otherwise wrong
      
    # transform back
    # correct for scaling
    # correct for frequency dependent gain due to channel distortion
    
    data_rec = (1./ N) * num.fft.fft(ofdm_rec) / G_channel
    
    # deviation of received / sent data symbols
    error = data_rec - qam_vec
    
    # retrieve data of sub-channel Nr. N_sub_ch
    sub_channel[n_symbol] = data_rec[N_sub_ch]
    

# record elapsed time    
t_stop = time.clock()
t_elapsed = t_stop - t_start

print "time elapsed = %10.2f s" % t_elapsed

# some figures ...

# scatter plot of sub-channel Nr. N_sub_ch
py.figure()
py.scatter(sub_channel.real, sub_channel.imag, s=1, c='b')
py.axis('equal')
py.xlabel('I')
py.ylabel('Q')
py.title("constellation of sub-carrier Nr = %d  # N = %d" % (N_sub_ch, N) )
py.grid(True)
py.show()