import numpy as np 
import pandas as pd

class ZeroForcingBD:
    def __init__(self, H, Mrs_arr):
        Mr, Mt = np.shape(H)
        self.Mr = Mr
        self.Mt = Mt
        self.H = H
        self.Mrs_arr = Mrs_arr
    
    def __routines(self, H, mr, shift):
        
        # used in self.process() - See example above for illustration 
        # inputs: 
        #       H - the whole channel matrix
        #       mr - number of receive antennas of the i-th user
        #       shift - how much receive antennas were considered before
        # outputs:
        #       Uidx, Sigmaidx, Vhidx - SVD decomposition of the H_iP_i 
        #       d - rank of the hat H_i
        #       Hidx - H_i (channel matrix for the i-th user)
        #       r - rank of the H_i
        
        Hidx = H[0+shift:mr+shift,:] # H_i (channel matrix for the i-th user)
        r = np.linalg.matrix_rank(Hidx) # rank of the H_i
        del_idx = [i for i in range(0+shift, mr+shift, 1)] # row indeces of H_i in H
        H_hat_idx = np.delete(H, del_idx, 0) # hat H_i
        d = np.linalg.matrix_rank(H_hat_idx) # rank of the hat H_i
        U, Sigma, Vh = np.linalg.svd(H_hat_idx) # SVD
        Vhn = Vh[d:, :] # null-subspace of V^H
        Vn = np.matrix(Vhn).H # null-subspace of V
        Pidx = np.dot(Vn, np.matrix(Vn).H) # projection matrix
        Uidx, Sigmaidx, Vhidx = np.linalg.svd(np.dot(Hidx, Pidx)) # SVD of H_iP_i 
        return Uidx, Sigmaidx, Vhidx, d, Hidx, r
    
    def process(self):
        
        # used in self.obtain_matrices()
        # outputs:
        #       F - whole filtering (pre-coding) matrix (array of arrays)
        #       D - whole demodulator (post-processing) matrix (array of arrays)
        #       H - the whole channel matrix (array of arrays)
        
        shift = 0
        H = self.H
        F = []
        D = []
        Hs = []
        for mr in self.Mrs_arr:
            Uidx, Sigmaidx, Vhidx, d, Hidx, r = self.__routines(H, mr, shift)
            Vhidx1 = Vhidx[:r,:] # signal subspace
            Fidx = np.matrix(Vhidx1).H
            F.append(Fidx)
            D.append(Uidx)
            Hs.append(Hidx)
            shift = shift + mr
        return F, D, Hs
    
    def obtain_matrices(self):
        
        # used to obtain pre-coding and post-processing matrices
        # outputs:
        #       FF - whole filtering (pre-coding) matrix 
        #       DD - whole demodulator (post-processing) matrix (array of arrays)
        
        F, D, Hs = self.process()
        FF = np.hstack(F)
        # Home Task: calculation of the demodulator matrices :)
        return FF

Mrs_arr = [3,2,3] 
# 1st user have 3 receive antennas, 2nd user - 2 receive antennas, 3d user - 3 receive antennas 
Mr = sum(Mrs_arr) # total number of the receive antennas 
Mt = 8 # total number of the transmitt antennas
H = (np.random.randn(Mr,Mt) + 1j*np.random.randn(Mr, Mt))/np.sqrt(2); #Rayleigh flat faded channel matrix (MrxMt)

BD = ZeroForcingBD(H, Mrs_arr)
F, D, Hs = BD.process()
FF = BD.obtain_matrices()

df = pd.DataFrame(np.dot(H, FF))
df[abs(df).lt(1e-14)] = 0
print(pd.DataFrame(np.round(np.real(df),100)))