import numpy as np
from scipy import constants
import matplotlib.pyplot as plt

def to_decibel(x):
    return 10*np.log10(x)

def length_los(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, **kwargs):
    distance_sq = (x_tx-x_rx)**2 + (y_tx-y_rx)**2
    return np.sqrt(distance_sq + (h_tx-h_rx)**2)

def length_ground(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, **kwargs):
    distance_sq = (x_tx-x_rx)**2 + (y_tx-y_rx)**2
    return np.sqrt(distance_sq + (h_tx+h_rx)**2)

def length_left(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, **kwargs):
    return np.sqrt((x_tx+x_rx)**2 + (y_tx-y_rx)**2 + (h_tx-h_rx)**2)

def length_right(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, width, **kwargs):
    return np.sqrt((2*width-x_tx-x_rx)**2 + (y_tx-y_rx)**2 + (h_tx-h_rx)**2)

def length_top(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, height, **kwargs):
    return np.sqrt((x_tx-x_rx)**2 + (y_tx-y_rx)**2 + (2*height-h_tx-h_rx)**2)

def _get_combination_matrices(array):
    num_paths = np.shape(array)[-1]
    array = np.expand_dims(array, -1)
    array = np.tile(array, (num_paths,))
    array_T = np.swapaxes(array, -1, -2)
    return array, array_T

def receive_power(path_len, refl_coeff, freq, pow_tx=1, c=constants.c):
    omega = 2*np.pi*freq
    _factor = pow_tx * (c/(2*omega))**2
    if not np.all(np.shape(path_len) == np.shape(refl_coeff)):
        raise TypeError
    path_len, path_len_T = _get_combination_matrices(path_len)
    refl_coeff, refl_coeff_T = _get_combination_matrices(refl_coeff)
    _sum_components = np.cos(omega/c * (path_len-path_len_T))*(refl_coeff*refl_coeff_T)
    #path_sum = np.sum(np.triu(_sum_components), axis=(-1, -2))
    path_sum = np.sum(_sum_components, axis=(-1, -2))
    rec_power = _factor * path_sum
    return rec_power

def lower_bound_rec_power(refl_coeff, freq, pow_tx=1, c=constants.c):
    omega = 2*np.pi*freq
    _factor = pow_tx * (c/(2*omega))**2
    refl_coeff, refl_coeff_T = _get_combination_matrices(refl_coeff)
    _sum_components = refl_coeff*refl_coeff_T
    _sum_off_diag = np.sum(np.abs(_sum_components), axis=(-1, -2))
    _sum_diag = np.sum(np.diagonal(_sum_components, axis1=-1, axis2=-2), axis=1)
    rec_power = _factor * (2*_sum_diag - _sum_off_diag)
    return rec_power

def main(freq=10e9):
    y_rx = np.logspace(1, 2, 2000)
    y_tx = np.zeros_like(y_rx)
    h_tx = 3
    h_rx = 1.5
    x_tx = 5
    x_rx = 5
    width = 10
    height = 10
    func_length = [length_los, length_ground, length_left, length_right, length_top]
    #func_length = [length_los, length_ground]
    #func_length = [length_los, length_ground, length_left]
    path_len = np.vstack([_length(x_tx, x_rx, y_tx, y_rx, h_tx, h_rx, width=width, height=height) for _length in func_length]).T
    refl_coeff = -1./path_len
    refl_coeff[:, 0] = -refl_coeff[:, 0]
    rec_pow = receive_power(path_len, refl_coeff, freq)
    print(rec_pow)
    print(to_decibel(rec_pow))
    plt.semilogx(y_rx, to_decibel(rec_pow))
    #rec_lower = lower_bound_rec_power(refl_coeff, freq)
    #plt.semilogx(y_rx, to_decibel(rec_lower), '--')
    #rec_pow2 = receive_power(path_len, refl_coeff, freq+50e6)
    #plt.semilogx(y_rx, to_decibel(0.5*(rec_pow+rec_pow2)), 'g--', alpha=.5)
    return

if __name__ == "__main__":
    main()
    plt.show()
