# importing
import numpy as np
from scipy import stats

import matplotlib.pyplot as plt
import matplotlib

# showing figures inline
%matplotlib inline

# plotting options 
font = {'size'   : 24}
plt.rc('font', **font)
plt.rc('text', usetex=matplotlib.checkdep_usetex(True))

matplotlib.rc('figure', figsize=(14, 6) )

"""
Explanation: Content and Objectives

Show effects of inter-symbol-interference 
Illustrated as distortion in the continuous time signal, sampled symbols, and in the BER
BPSK symbols are being pulse-shaped, transmitted on a given multi-path channel (you may play around with this), and processed in the receiver

Import
End of explanation
"""


########################
# find impulse response of an RRC filter
########################
def get_rrc_ir(K, n_up, t_symb, r):
        
    ''' 
    Determines coefficients of an RRC filter 
    
    Formula out of: J. Huber, Trelliscodierung, Springer, 1992, S. 15
    At poles, values of wikipedia.de were used (without cross-checking)
    
    NOTE: Length of the IR has to be an odd number
    
    IN: length of IR, upsampling factor, symbol time, roll-off factor
    OUT: filter ceofficients
    '''

    assert K % 2 != 0, "Filter length needs to be odd"
    
    if r == 0:
        r = 1e-32

    # init
    rrc = np.zeros(K)
    t_sample = t_symb/n_up
    
        
    i_steps = np.arange( 0, K)
    k_steps = np.arange( -(K-1)/2.0, (K-1)/2.0 + 1 )    
    t_steps = k_steps*t_sample

    for i in i_steps:

        if t_steps[i] == 0:
            rrc[i] = 1.0/np.sqrt(t_symb) * (1.0 - r + 4.0 * r / np.pi )

        elif np.abs( t_steps[i] ) == t_symb/4.0/r:
            rrc[i] = r/np.sqrt(2.0*t_symb)*((1+2/np.pi)*np.sin(np.pi/4.0/r)+ \
                            ( 1.0 - 2.0/np.pi ) * np.cos(np.pi/4.0/r) )

        else:
            rrc[i] = 1.0/np.sqrt(t_symb)*( np.sin( np.pi*t_steps[i]/t_symb*(1-r) ) + \
                            4.0*r*t_steps[i]/t_symb * np.cos( np.pi*t_steps[i]/t_symb*(1+r) ) ) \
                            / (np.pi*t_steps[i]/t_symb*(1.0-(4.0*r*t_steps[i]/t_symb)**2.0))
 
    return rrc

"""
Explanation: Function for determining the impulse response of an RRC filter
End of explanation
"""


# number of symbols per sequence/packet
n_symb = 32

# modulation scheme and constellation points
M = 2
constellation = [ -1, 1 ]

# EbN0 range for simulation
EbN0_dB = np.arange( 0, 30, 2)
EbN0 = 10**(EbN0_dB/10)


# maximum number of errors and symbols to be simulated
max_errors = 1000
max_syms = 1e7

# parameters of the filter
r = 0.33

n_up = 8            # samples per symbol
syms_per_filt = 4   # symbols per filter (plus minus in both directions)

K_filt = 2 * syms_per_filt * n_up + 1         # length of the fir filter

# set symbol time 
t_symb = 1.0 

"""
Explanation: Parameters
End of explanation
"""


# define channel by characterizing delays and according attenuation
channel_delays_syms = range( 7 )
channel_factors = [ 1, .4, -.3, .1, .1, .05, .05 ]

assert( len(channel_delays_syms) == len(channel_factors) ), 'Length of delays and factors has to be the same!'

# get channel 
h_channel = np.zeros( (np.max(channel_delays_syms)+1) * n_up )

for k in np.arange(len(channel_delays_syms)):
    h_channel[ n_up*channel_delays_syms[k] ] = channel_factors[k]

"""
Explanation: Define channel characteristics and get channel impulse response
End of explanation
"""


# find rrc response and normalize to energy 1
rrc = get_rrc_ir( K_filt, n_up, t_symb, r)
rrc = rrc / np.linalg.norm(rrc)

# initialize BER  
ber = np.zeros_like( EbN0, dtype=float ) 

# theoretical values
ber_bpsk = 1 - stats.norm.cdf( np.sqrt( 2 * EbN0 ) )


# loop for snrs
for ind_ebn0, val_ebn0 in enumerate( EbN0 ):

    # get noise variance for simulation
    sigma2 = 1 / (np.log2(M) * val_ebn0)
    
    # initialize counter
    num_errors = 0
    num_syms = 0

    # loop for errors
    while (num_errors<max_errors and num_syms<max_syms):    
    
        # generate random binary vector and modulate 
        data = np.random.randint( 2, size=n_symb)
        s = [ constellation[d] for d in data ]
          
        # prepare sequence to be filtered by upsampling 
        s_up = np.zeros( n_symb * n_up, dtype=complex)        
        s_up[ : : n_up ] = s
    
        # apply RRC filtering for Tx pulse shaping     
        s_Tx = np.convolve( s_up, rrc)    
    
        # apply channel and add noise
        s_Rx = np.convolve( s_Tx, h_channel)

        n = np.sqrt(sigma2/2) * ( np.random.randn(len(s_Rx)) + 1j* np.random.randn(len(s_Rx)) )
    
        r_Rx = s_Rx + n
    
        # apply MF at the Rx 
        y_mf_rrc = np.convolve(r_Rx, rrc)
        
        # down-sampling from "high rate" (n_up samples per symbol) to symbol rate
        y_down = y_mf_rrc[ K_filt-1 : K_filt-1 + n_symb*n_up : n_up ]
        
        # demodulate
        data_est = [ int( np.real( rec ) > 0 ) for rec in y_down ]

        # count errors and symbols
        num_errors += sum( [ int( data_est[ _n ] != data[ _n ] ) for _n in range(len(data)) ] )
        num_syms += n_symb

        # get estimate of BER
        ber[ ind_ebn0 ] = num_errors/(num_syms*1.0*np.log2(M))
        
    print('EbN0 simulated:', 10*np.log10( val_ebn0) )
    
print('\n\nError floor at: \t\t\t{:0.4f}'.format(ber[-1]))
print('For comparison: 2^length of h = \t{:0.4f}'.format(2**(-len(channel_delays_syms))) )

"""
Explanation: Simulation loop
End of explanation
"""


# activate switches to plot whatever you like to see
show_signal = 1
show_ber = 1
show_symbols = 1


# plot signals
if show_signal:
    
    plt.figure()

    plt.plot( np.real(s_Tx), label='$s(t)$')
    plt.plot( np.real(s_Rx), label='$s(t)* h(t)$')
    plt.plot( np.real(r_Rx), label='$r(t)= s(t)*h(t)+n(t)$')
    plt.plot( np.real(y_mf_rrc), label='$y(t)=r(t)*g^\star(-t)$')
    
    plt.grid(True)
    plt.xlabel('$t$ (a.u.)')
    plt.autoscale(enable=True, tight=True)
    plt.legend(loc='upper right')
    plt.savefig('waveform_isi.pdf',bbox_inches='tight')
    
# show data symbols after processing
if show_symbols:
    plt.figure()

    markerline, stemlines, baseline = plt.stem( np.arange(len(s)), np.real(s), label='syms Tx')
    plt.setp(markerline, 'markersize', 8, 'markerfacecolor', 'b')

    markerline, stemlines, baseline = plt.stem( np.arange(len(y_down)), np.real(y_down), '+', label='syms Rx')
    plt.setp(markerline, 'markersize', 12,  'markerfacecolor', 'r',)

    plt.legend(loc='upper right')
    plt.grid(True)
    plt.xlabel('$n$')
    

# show BER
if show_ber:
    plt.figure()
    
    plt.plot(EbN0_dB, ber_bpsk,  label="BPSK, nur AWGN")
    plt.plot(EbN0_dB, ber, 'o', ms=14, label="Multipath, sim.")

    plt.yscale('log')
    plt.grid(True)
    plt.legend(loc='lower left')
    plt.xlabel('$E_b/N_0$ (dB)')
    plt.ylabel('BER')
    plt.autoscale(enable=True, tight=True)
    plt.ylim( (1e-6,1))
    plt.savefig('ber_isi.pdf',bbox_inches='tight')

"""
Explanation: Plotting
End of explanation
"""


